如何在Tensorflow 2.0 + Keras中进行并行GPU推理？

2024-04-29 • 问答

让我们从一个前提开始，即我刚开始接触TensorFlow和一般的深度学习。

我有使用tf.Model.train()训练的TF 2.0 Keras风格模型，两个可用的GPU，并且我希望缩短推理时间。

我使用非常方便的tf.distribute.MirroredStrategy().scope()上下文管理器训练了在GPU上分布的模型。

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model.compile(...)
  model.train(...)

两个GPU都得到有效使用（即使我对结果的准确性不太满意）。

我似乎无法找到一种类似的策略来通过tf.Model.predict()方法在GPU之间分配推理：当我运行model.predict()时，（显然）只有两个GPU之一可以使用。 / p>

是否有可能在两个GPU上建立相同的模型并并行地向它们提供不同的数据块？

有些帖子建议如何在TF 1.x中执行此操作，但我似乎无法在TF2.0中复制结果

我在这个问题上的精神挣扎主要是

TF 1.x是基于tf.Session()的，而会话在TF2.0中是隐式的，如果我正确理解的话，我阅读的解决方案将为每个GPU使用单独的会话，而我真的不知道如何复制它在TF2.0中
我不知道如何在特定会话中使用model.predict()方法。

我知道这个问题的格式可能不正确，但我将其总结为：

有人知道如何在TF2.0中的多个GPU上运行Keras风格的model.predict()（以并行方式推断每个GPU上不同批次的数据）吗？

在此先感谢您的帮助。