训练操作顺序分布在多个GPU上的TensorFlow模型

2024-03-29 14:14:02 发布

您现在位置:Python中文网/ 问答频道 /正文

我想实现一个模型,它的操作在不同的GPU上执行(相应的权重放在不同的GPU上)。我知道这不是最优的,我可能应该做的是完全的模型并行,但这只是第一步

当我尝试手动放置设备时,我想我确实在多个GPU上执行了正向传递,但没有执行反向传播。 在进行模型并行时,如何确保反向传播也发生在TensorFlow中的多个GPU上(如果可能的话)?

对于2 GPU,我的代码的最小可复制示例如下所示:

import tensorflow as tf
from tensorflow.keras.layers import Conv3D
from tensorflow.keras.models import Model

tf.debugging.set_log_device_placement(True)

class MyBigModel(Model):
    def __init__(self, n_convs=2, n_filters=32, **kwargs):
        super().__init__(**kwargs)
        self.convs = []
        self.n_convs = n_convs
        for i in range(n_convs - 1):
            with tf.device(f'/gpu:{int(i<self.n_convs//2)}'):
                conv = Conv3D(
                    n_filters,
                    3,
                    padding='same',
                    activation='relu',
                )
            self.convs.append(conv)
        with tf.device('/gpu:1'):
            last_conv = Conv3D(
                1,
                3,
                padding='same',
            )
        self.convs.append(last_conv)
        
    def call(self, inputs):
        outputs = inputs
        for i, conv in enumerate(self.convs):
            with tf.device(f'/gpu:{int(i<self.n_convs//2)}'):
                outputs = conv(outputs)
        return outputs

model = MyBigModel(n_convs=10, n_filters=256)
size = 128
inputs = tf.random.normal([1, size, size, size, 1])
outputs = tf.random.normal([1, size, size, size, 1])
res = model(inputs)  # just a forward pass
model.compile(loss='mse', optimizer='sgd')
model.train_on_batch(inputs, outputs)  # forward and backward pass

我试图确保反向传播也发生在多个GPU上的方法是通过肉眼监视GPU的使用情况。 我注意到在执行train_on_batch方法的过程中有以下步骤:

  • GPU0为100%
  • GPU1为100%
  • GPU0在相当长的一段时间内为100%

因此,我的结论是,前向传递确实发生在2个GPU上,而不是只发生在第一个GPU上的反向传播。 也许我没有用正确的方法诊断,在这种情况下,请告诉我怎么做

我只是想澄清一下,我对数据并行性(在不同的GPU之间分割批处理)不感兴趣,而是对在不同的GPU上按顺序执行模型的不同组件感兴趣。 当然mesh-tensorflow是我正在考虑的真实模型并行性的替代方案,但我首先想按顺序进行尝试