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

最新推荐文章于 2026-03-01 03:04:21 发布

原创

最新推荐文章于 2026-03-01 03:04:21 发布 · 3.5k 阅读

标签

#tensorflow #分布式训练 #单机多GPU #GPU并行训练

本文介绍了如何在GPU资源不足的情况下，利用TensorFlow的MirroredStrategy进行分布式训练，特别是针对Transformer模型。主要步骤包括初始化分布策略、数据加载与分发、定义损失函数、评价指标、模型初始化、训练策略构建以及自定义训练过程。文章强调了在处理损失函数和数据分发时的注意事项，以及如何确保在分布式环境中正确计算损失和更新模型。

由于最近训练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。不允许