如何在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上实例化相同的模型，并以并行方式将不同的数据块馈送给它们？

有一些帖子建议如何在TF 1.x中实现，但我似乎无法在TF2.0中复制结果。

https://medium.com/@sbp3624/tensorflow-multi-gpu-for-inferencing-test-time-58e952a2ed95

Tensorflow：同时在GPU和CPU上进行预测

我与此问题的斗争主要是：

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

我知道这个问题可能没有明确定义，但我总结一下：

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

非常感谢任何帮助。

- Filippo Castelli

---

这可能会有所帮助。我遇到了同样的问题，似乎我的问题源于我将模型保存在hdf5格式中，而在将模型加载回分布式策略时不支持该格式。 - Will.Evo

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)

- DachuanZhao

---

这个在使用 model.predict_on_batch(greater_batch_size) 时不起作用，对吧？ - Derk

Answer 2

目前还没有官方的分布式推理示例。这里提供一种可能的解决方案，使用tf.distribute.MirroredStrategy：https://github.com/tensorflow/tensorflow/issues/37686。然而，它似乎没有完全利用多个GPU。

- Ryan abbu