tensorflow - 跨多个 GPU 镜像多 GPU 模型

Question

我有一个使用多个 GPU 进行计算的 tensorflow (tf2.0)/keras 模型。模型中有 2 个分支，每个分支都在单独的 GPU 上。

我有一个要用于训练的 4 GPU 系统，我想镜像这个模型，以便 GPU 1 和 2 包含一个模型，而 GPU 3 和 4 包含镜像模型。

tf.distribute.MirroredStrategy 会自动 处理这个镜像吗？还是假设我的模型将是单个 GPU 模型？

如果tf.distribute.MirroredStrategy不能处理这个问题，是否有人对如何定制MirroredStrategy以实现这种培训结构有任何建议？

标签： tensorflowtf.kerasmulti-gpu

Answer 1

这听起来很像您需要进行自定义训练循环。镜像策略在每个 GPU 上复制模型，但由于您的模型已经在两个 GPU 上，我认为它不能正常工作。

但是您可以尝试一下，并使用 nvidia-smi 检查 tensorflow 正在做什么。