如何将多个 GPU 用于协同工作的多个模型？

Posted 2023-02-16

【中文标题】如何将多个 GPU 用于协同工作的多个模型？

【英文标题】：How to use multiple GPUs for multiple models that work together?

【发布时间】：2021-12-27 03:37:04

【问题描述】：

我在 tensorflow 的不同设备范围下定义了三个模型，我正在使用 GradientTape 来训练这些网络。当我这样做时，内存会增加几百兆字节，以表明模型已加载到各自的 GPU 中。问题是，当我开始训练时，即使批量非常小，也只有 GPU @位置 0 内存增加。我也在使用 GradientTape 进行培训。有什么方法可以确保只有分配给模型的 GPU 用于该模型？

with tf.device('/device:GPU:0'):
    model1 = model1Class().model()

with tf.device('/device:GPU:1'):
    model2 = model2Class().model()

with tf.device('/device:GPU:2'):
    model3 = model3Class().model()


for epoch in range(10):
    dataGen = DataGenerator(...)
    X, y = next(dataGen)

    with tf.GradientTape() as tape1:
         X = model1(X)
         loss1 = lossFunc(X, y[1])
    grads1 = suppressionTape.gradient(tape1,model1.trainable_weights)
    optimizer1.apply_gradients(zip(model1.trainable_weights))

    with tf.GradientTape() as tape2:
         X = model2(X)          # Uses output from model2
         loss2 = lossFunc(X, y[2])
    grads2 = suppressionTape.gradient(tape2,model2.trainable_weights)
    optimizer2.apply_gradients(zip(model2.trainable_weights))

    with tf.GradientTape() as tape3:
         X = model3(X)          # Uses output from model3
         loss3 = lossFunc(X, y[3])
    grads3 = suppressionTape.gradient(tape3,model3.trainable_weights)
    optimizer3.apply_gradients(zip(model3.trainable_weights))

【参考方案1】：

我必须承认，我一直在寻找一些正确的解决方案来解决您的问题。看来您的问题的答案在这里（学分归 Laplace Ricky）：

@Laplace Ricky：它应该在单个 gpu 中运行（可能是 第一个 gpu, GPU:0) 用于任何代码之外的代码 mirrored_strategy.run()。另外，因为您想要渐变 从副本返回，mirrored_strategy.gather() 也是必需的。

除此之外，还必须使用创建分布式数据集 mirrored_strategy.experimental_distribute_dataset。分布式数据集 尝试将单批数据均匀地分布在副本之间。一个 下面包括关于这些要点的示例。

model.fit()、model.predict() 等...以分布式方式运行 自动只是因为他们已经处理了所有事情 上面为你提到的。

在此处查看此主题：Tensorflow - Multi-GPU doesn’t work for model(inputs) nor when computing the gradients。

您需要使用mirrored_strategy.experimental_distribute_dataset(dataset) 并根据您的需要调整代码。