如何进行多 GPU 的分布式训练？

在之前的学习之中，我们进行的训练都是在一台机器上进行的。或者更具体的说，我们之前的学习任务都是在同一台机器的 CPU 或者一台机器的同一个 GPU 上进行训练的。那么我们如何才能将机器之中的多个 GPU 充分地利用起来呢？那么我们这节课就来学习一下如何进行多 GPU 的分布式训练。

1. 什么是分布式训练

分布式训练，顾名思义，就是在多个设备之上进行训练。它可以充分的使用硬件资源，使得训练与学习任务可以在更短的时间内完成。

具体来说，分布式任务大体可以分为以下几个模块：

主程序将具体的任务进行分割，分割成多个小型的任务；
将分割后的小型任务分配到不同的设备之中去，并让他们独立执行；
不同的设备在完成任务后会将产生的结果返回到主程序；
主程序会将结果进行合并，从而得到最终的结果。

既然明白了分布式任务的基本原理，那么其中的 “不同的设备” 是如何定义的呢？

其实这个设备的含义很广泛，它可以包括以下的含义：

同一台计算机上的不同 GPU ；
不同计算机的不同 GPU ；
不同的计算机本身作为一个单独的设备。

在实际的应用过程之中，我们使用最多的情况是在同一台设备上会有多张 GPU 卡，因此我们大多数的分布式训练实在同一台机器上的不同 GPU 中进行的。因此我们这节课着重介绍如何在多张 GPU 显卡之中进行机器学习的相关任务。

在实际的 TensorFlow 的分布式训练之中，包括很多的实现方式，结合我们之前采用的大多数训练方式是使用 tf.keras 进行训练，因此我们这节课会着重介绍如何使用 tf.keras 的模型进行单机器多 GPU 分布式训练。

2. 如何使用 tf.keras 模型进行单机器多 GPU 分布式训练

在 TensorFlow 之中进行石激起多 GPU 显卡进行分布式训练的总体实现方式大致为：

定义要使用的 GPU ；
将完全一样的模型放在不同的 GPU 之上并分别进行编译；
将数据送到不同的模型之中训练并得到相应的结果；
主程序将结果进行整合。

这个过程看上去很复杂，但是在 TensorFlow 之中，我们只需要进行第二部的操作即可。

那么我们为什么不用定义 GPU 呢？因为 TensorFlow 的 GPU 版本会自动检测 GPU ，并且进行识别。

因此我们要使用的 API 为：

tf.distribute.MirroredStrategy()

在具体使用的时候，我们需要实例化一个对象，然后在以下语句的范围内进行模型的编译与定义即可：

tf.distribute.MirroredStrategy()
with strategy.scope():
    # 模型的定义与编译
    ......

该 API 会将在 Scope 之中定义模型创建不同的镜像，并将这些镜像分不到不同的 GPU 之中去，从而实现分布式的训练。

而之后的训练步骤，我们只需要和平常进行训练即可，剩下的工作 TensorFlow 会帮助我们完成。

3. 使用 tf.keras 模型进行单机器多 GPU 分布式训练的程序实例

在这节课之中，我们会使用一个简单的实例来进行演示使用多个 GPU 进行训练。这里以我们之前使用过的 fashion_mnist 数据集的分类进行演示。

具体的程序为：

import tensorflow as tf

# 获取fashion_mnist数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# 预处理图片数据，使其归一化
x_train, x_test = x_train / 255.0, x_test / 255.0

# 定义镜像策略
strategy = tf.distribute.MirroredStrategy()

# 单个GPU之上的Batch数量
batch_per_gpu = 64
# 同时训练的Batch大小应该为：单个GPU的Batch * GPU个数
batch = batch_per_gpu * strategy.num_replicas_in_sync


train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch).shuffle(buffer_size=10000)
eval_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch)


with strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
  ])

  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                optimizer=tf.keras.optimizers.Adam(),
                metrics=['accuracy'])
  

model.fit(train_dataset, epochs=10)

eval_loss, eval_acc = model.evaluate(eval_dataset)
print(eval_loss, eval_acc)

这里网络模型，我们采用了与之前相同的网络模型。

值得注意的是，这边的 Batch 大小：

我们首先定义单个 GPU 上的批次大小；
然后我们需要定义总体的批次大小 = 单个 GPU 上的批次大小 * GPU 数量；
在训练过程之中，TensorFlow 会将批次大小自动分割分给不同的 GPU。

最终我们可以得到如下结果：

INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)
Epoch 1/10
938/938 [==============================] - 3s 3ms/step - loss: 1.7785 - accuracy: 0.6877
Epoch 2/10
938/938 [==============================] - 4s 4ms/step - loss: 1.7312 - accuracy: 0.7306
......
Epoch 10/10
938/938 [==============================] - 3s 3ms/step - loss: 1.5702 - accuracy: 0.8915
157/157 [==============================] - 0s 2ms/step - loss: 1.5899 - accuracy: 0.8721
1.589921236038208, 0.8720999956130981

可以看到，我们最终完成了训练，说明我们的分布式训练得到了正确的结果。

感兴趣的同学可以尝试一下不使用分布式训练的结果和以上结果有什么区别。

4. 小结

在这节课之中，我们学习了什么是分布式训练，同时了解了分布式训练的种类。然后我们具体的了解了如何在同一台机器上使用多 GPU 进行训练，最后我们通过一个简单的示例进行了演示。