如何將多個 GPU 用於多個協同工作的模型？

Question

我在 tensorflow 的不同設備范圍下定義了三個模型，我正在使用 GradientTape 來訓練這些網絡。 當我這樣做時，內存會增加幾百兆字節，以表明模型已加載到相應的 GPU 中。 問題是當我開始訓練時，即使批量很小，也只有 GPU @ position 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))

Answer 1

我必須承認，我一直在尋找為您提供解決問題的正確方法。 似乎您的問題的答案就在這里（感謝拉普拉斯·瑞奇（Laplace Ricky））：

@Laplace Ricky：對於 mirrored_strategy.run() 之外的任何代碼，它應該在單個 gpu（可能是第一個 gpu，GPU:0）中運行。 此外，由於您希望從副本返回梯度，因此還需要 mirrored_strategy.gather()。

除此之外，必須使用 mirrored_strategy.experimental_distribute_dataset 創建分布式數據集。 分布式數據集嘗試在副本之間均勻分布單批數據。 以下是關於這些要點的示例。

model.fit()、model.predict() 等...以分布式方式自動運行，因為它們已經為您處理了上面提到的所有內容。

請在此處查看此線程： Tensorflow - Multi-GPU 不適用於模型（輸入）或計算梯度。

您需要使用mirrored_strategy.experimental_distribute_dataset(dataset)並根據您的需要調整代碼。