跨多个gpu镜像多gpu模型

Question

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

我有一个4个GPU的系统，我想用来训练，我想镜像这个模型，这样GPU 1和2包含一个模型，GPU 3和4包含镜像的模型。

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

如果tf.distribute.MirroredStrategy不能处理这个问题，有没有人对如何定制MirroredStrategy来实现这个培训结构有什么建议呢？

Answer 1

这听起来很像你需要做一个定制的训练循环。镜像策略在每个GPU上复制模型，但由于您的模型在两个GPU上都已就绪，我认为它不会正常工作。

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