让我们从我刚开始接触 TensorFlow 和深度学习的前提开始。
我有使用tf.Model.train()两个可用 GPU训练的 TF 2.0 Keras 样式模型,我希望缩短推理时间。
我使用极其方便的tf.distribute.MirroredStrategy().scope()上下文管理器训练了分布在 GPU 上的模型
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
model.compile(...)
model.train(...)
Run Code Online (Sandbox Code Playgroud)
两个 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
我对这个问题的心理斗争主要是
tf.Session()基于 TF2.0 中隐含的会话,如果我理解正确,我阅读的解决方案对每个 GPU 使用单独的会话,我真的不知道如何在 TF2.0 中复制它model.predict()特定会话中使用该方法。我知道这个问题可能没有很好地表述,但我总结为:
有没有人知道如何model.predict()在 TF2.0 中在多个 GPU 上运行 Keras 风格(以并行方式对每个 GPU 上的不同批次数据进行推理)?
在此先感谢您的帮助。