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

Question

让我们从我刚开始接触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中的一个获得使用。

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

有些帖子建议如何在TF1.x中做到这一点，但我似乎无法在TF2.0中复制结果

https://medium.com/@sbp3624/tensorflow-多gpu参考-测试时间-58e952a2ed95

Tensorflow：GPU和CPU 上的同时预测

我对这个问题的心理挣扎主要是

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

我知道这个问题可能没有很好地表述，但我将其总结为：

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

提前感谢您的帮助。

Answer 1

尝试在tf.distribute.MirroredStrategy中加载模型并使用更大的batch_size

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
model = tf.keras.models.load_model(saved_model_path)
result = model.predict(batch_size=greater_batch_size)

Answer 2

分布式推理似乎还没有正式的例子。使用tf.distribute.MirroredStrategy有一个潜在的解决方案：https://github.com/tensorflow/tensorflow/issues/37686.然而，它似乎没有完全利用多gpus