基于tensorflow2.x的多GPU并行训练

由于最近训练transformer，在单卡上显存不够，另外一块卡上也无法加载，故尝试使用双卡并行的策略。将基本的流程、遇见的难题汇总在这里。

分布策略解释

使用官方给出的tf.distribute.MirroredStrategy作为分布策略。这个策略通过如下的方式运行：
1）所有变量和模型计算图都会在副本之间复制。
2）输入都均匀分布在副本中。
3）每个副本在收到输入后计算输入的损失和梯度。
4）通过求和，每一个副本上的梯度都能同步。
5）同步后，每个副本上的复制的变量都可以同样更新。

正文

初始化分布策略

可以使用如下的命令，查看当前设备有几块GPU可以供使用。

strategy = tf.distribute.MirroredStrategy()
print(strategy.num_replicas_in_sync)

一、数据加载

使用分布式训练，会将总的batch分散到多块GPU上。我这里有两块GPU，使用的batch是32，那么在每个上面就是16。这里，在数据加载的时候就需要做处理，具体处理过程如下：

1）创建一个总的batchsize

GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync

2) 加载数据集

train_ds = DataLoader().make_batch(PARA['train_'], GLOBAL_BATCH_SIZE, PARA['max_len_sequence'])
valid_ds = DataLoader().make_batch(PARA['vaild_'], GLOBAL_BATCH_SIZE, PARA['max_len_sequence'])
test_ds = DataLoader().make_batch(PARA['testt_'],  GLOBAL_BATCH_SIZE, PARA['max_len_sequence'])

3）对数据做分发

train_ds = strategy.experimental_distribute_dataset(train_ds)
valid_ds = strategy.experimental_distribute_dataset(valid_ds)
test_ds = strategy.experimental_distribute_dataset(test_ds)

经过上面这些操作，数据已经处理好了，接下来处理训练策略。

二、 定义损失函数

注：这里有几个地方需要特别注意，tf.losses/tf.keras.losses 中的损失函数通常会返回输入最后一个维度的平均值。损失类封装这些函数。在创建损失类的实例时传递 reduction=Reduction.NONE，表示“无额外缩减”。对于样本输入形状为 [batch, W, H, n_classes] 的类别损失，会缩减 n_classes 维度。对于类似 losses.mean_squared_error 或 losses.binary_crossentropy 的逐点损失，应包含一个虚拟轴，使 [batch, W, H, 1] 缩减为 [batch, W, H]。如果没有虚拟轴，则 [batch, W, H] 将被错误地缩减为 [batch, W]。
增加虚拟轴的方式也很简单，labels = labels[:, tf.newaxis]。如果没有这个，回归模型是跑不起来的！！！

1）使用 tf.distribute.Strategy 时应如何计算损失？

例如，假设有 2 个 GPU，批次大小为 64。一个批次的输入会分布在各个副本（2 个 GPU）上，每个副本获得一个大小为 32 的输入。

每个副本上的模型都会使用其各自的输入进行前向传递，并计算损失。现在，不将损失除以其相应输入中的样本数 (BATCH_SIZE_PER_REPLICA = 32)，而应将损失除以 GLOBAL_BATCH_SIZE (64)。

之所以需要这样做，是因为在每个副本上计算完梯度后，会通过对梯度求和在副本之间同步梯度。

2）计算方法

如果使用自定义训练循环，则应将每个样本的损失相加，然后将总和除以 GLOBAL_BATCH_SIZE: scale_loss = tf.reduce_sum(loss) * (1. / GLOBAL_BATCH_SIZE)，或者使用 tf.nn.compute_average_loss，它会将每个样本的损失、可选样本权重和 GLOBAL_BATCH_SIZE 作为参数，并返回经过缩放的损失。比较而言，选择tf.nn.compute_average_loss这个会好一些。

由于我这里使用的是 tf.keras.losses 类，则需要将损失归约显式指定为 NONE 或 SUM。与 tf.distribute.Strategy 一起使用时，不允许使用 AUTO 和 SUM_OVER_BATCH_SIZE。不允许使用 AUTO，因为用户应明确考虑他们想要的归约量，以确保在分布式情况下归约量正确。不允许使用 SUM_OVER_BATCH_SIZE，因为当前它只能按副本批次大小进行划分，而将按副本数量划分留给用户，这可能很容易遗漏。因此，您需要自己显式执行归约操作。

我做的是回归任务，具体的代码如下，可以看到，loss损失里面使用了reduction=tf.keras.losses.Reduction.NONE，返回损失值的时候使用了tf.nn.compute_average_loss。

GLOBAL_BATCH_SIZE = PARA['batch_size']*strategy.num_replicas_in_sync
with strategy.scope():
    # Set reduction to `NONE` so you can do the reduction afterwards and divide by
    # global batch size.
    loss_object = tf.keras.losses.Huber(reduction=tf.keras.losses.Reduction.NONE)
    def compute_loss(labels, predictions):
        # 这里有个坑，见最开始的注
        # 使用Reduction.NONE之后，回归损失会减少一个维度，故要在后面添加一列
        # https://tensorflow.google.cn/tutorials/distribute/custom_training?hl=zh-cn
        labels = labels[:,tf.newaxis]
        predictions = predictions[:, tf.newaxis]
        per_example_loss = loss_object(labels, predictions)
        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)

三、定义评价指标

评价指标根据自己的实际情况来，我这里使用了loss跟rmse

with strategy.scope():
    test_loss = tf.keras.metrics.Mean(name='test_loss')
    train_rmse = tf.keras.metrics.RootMeanSquaredError(name='train_rmse')
    test_rmse = tf.keras.metrics.RootMeanSquaredError(name='test_rmse')

四、初始化模型

模型、优化器和checkpoint务必要放在strategy.scope

with strategy.scope():
    model = Transformer(PARA['num_layers'], PARA['input_vocab_size'], PARA['target_vocab_size'],PARA['target_class'],PARA['max_len_sequence'],PARA['d_model'],PARA['num_heads'],PARA['dff'],rate=PARA['dropout_rate'])
    # 加载优化器：
    learning_rate = CustomizedSchedule(PARA['d_model'])
    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
    # 记录模型
    check = tf.train.Checkpoint(model=model, optimizer=optimizer)
    check_manager = tf.train.CheckpointManager(check, PARA['model_save'], max_to_keep=5)
    if check_manager.latest_checkpoint:
        check.restore(check_manager.latest_checkpoint)

五、构建训练策略

1) 先构建并行的策略，再构建train_step

with strategy.scope():
    # `run` replicates the provided computation and runs it
    # with the distributed input.
    @tf.function
    def distributed_train_step(dataset_inputs):
        per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
        return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

    @tf.function
    def distributed_test_step(dataset_inputs):
        return strategy.run(test_step, args=(dataset_inputs,))

2) 构建train_step

def train_step(inputs):
    train_rmse.reset_states()
    sequence, tm, label = inputs
    with tf.GradientTape() as tape:
        predictions = model(sequence, training=True)
        loss = compute_loss(tm, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_rmse.update_state(tm, predictions)
    return loss

def test_step(inputs):
    sequence, tm, label = inputs
    predictions = model(sequence, training=False)
    t_loss = loss_object(tm, predictions)
    test_loss.update_state(t_loss)
    test_rmse.update_state(tm, predictions)

六、自定义训练过程

def fit(train_ds, valid_ds, test_ds):
        steps = 0
        start = time.time()
        for epoch in range(PARA['EPOCH']):
            # TRAIN LOOP
            total_loss, num_batches, batch = 0.0, 0, 0
            for (batch, x) in enumerate(train_ds):
            	# 这里返回每一个批次的损失值
                per_loss= distributed_train_step(x)
                total_loss += per_loss
                steps += 1
                # 这是自定义的记录函数，可以直接print当前值
                save_smurry('train','-', epoch, batch, steps, [per_loss, train_rmse.result()])
                
                if batch % (PARA['REPORT_STEP']*2) == 0 and batch:
                # 每次处理完之后，需要对test_loss及test_rmse做重置
                    for (batch, x) in enumerate(valid_ds):
                        distributed_test_step(x)
                    # 这里需要得到的是在整个验证集上的结果
                    save_smurry('vaild','-', epoch, batch, steps, [test_loss.result(), test_rmse.result()])
                    test_loss.reset_states()
                    test_rmse.reset_states()

                # 每50次做一次benchmark验证
                if batch % (PARA['REPORT_STEP']*5) == 0 and batch:
                    for x in test_ds:
                        distributed_test_step(x)
                    save_smurry('test','-', epoch, batch, steps, [test_loss.result(), test_rmse.result()])
                    test_loss.reset_states()
                    test_rmse.reset_states()
        time_used = 'Time take for 1 epoch:{} secs\n'.format(time.time()-start)
        fout(time_used)

至此，分布程序构建完成。欢迎一起讨论