如何進行多 GPU 的分布式訓練？

在之前的學習之中，我們進行的訓練都是在一臺機器上進行的?；蛘吒唧w的說，我們之前的學習任務都是在同一臺機器的 CPU 或者一臺機器的同一個 GPU 上進行訓練的。那么我們如何才能將機器之中的多個 GPU 充分地利用起來呢？那么我們這節(jié)課就來學習一下如何進行多 GPU 的分布式訓練。

1. 什么是分布式訓練

分布式訓練，顧名思義，就是在多個設備之上進行訓練。它可以充分的使用硬件資源，使得訓練與學習任務可以在更短的時間內完成。

具體來說，分布式任務大體可以分為以下幾個模塊：

• 主程序將具體的任務進行分割，分割成多個小型的任務；
• 將分割后的小型任務分配到不同的設備之中去，并讓他們獨立執(zhí)行；
• 不同的設備在完成任務后會將產生的結果返回到主程序；
• 主程序會將結果進行合并，從而得到最終的結果。

既然明白了分布式任務的基本原理，那么其中的 “不同的設備” 是如何定義的呢？

其實這個設備的含義很廣泛，它可以包括以下的含義：

• 同一臺計算機上的不同 GPU ；
• 不同計算機的不同 GPU ；
• 不同的計算機本身作為一個單獨的設備。

在實際的應用過程之中，我們使用最多的情況是在同一臺設備上會有多張 GPU 卡，因此我們大多數的分布式訓練實在同一臺機器上的不同 GPU 中進行的。因此我們這節(jié)課著重介紹如何在多張 GPU 顯卡之中進行機器學習的相關任務。

在實際的 TensorFlow 的分布式訓練之中，包括很多的實現方式，結合我們之前采用的大多數訓練方式是使用 tf.keras 進行訓練，因此我們這節(jié)課會著重介紹如何使用 tf.keras 的模型進行單機器多 GPU 分布式訓練。

2. 如何使用 tf.keras 模型進行單機器多 GPU 分布式訓練

在 TensorFlow 之中進行石激起多 GPU 顯卡進行分布式訓練的總體實現方式大致為：

• 定義要使用的 GPU ；
• 將完全一樣的模型放在不同的 GPU 之上并分別進行編譯；
• 將數據送到不同的模型之中訓練并得到相應的結果；
• 主程序將結果進行整合。

這個過程看上去很復雜，但是在 TensorFlow 之中，我們只需要進行第二部的操作即可。

那么我們?yōu)槭裁床挥枚x GPU 呢？因為 TensorFlow 的 GPU 版本會自動檢測 GPU ，并且進行識別。

因此我們要使用的 API 為：

tf.distribute.MirroredStrategy()

在具體使用的時候，我們需要實例化一個對象，然后在以下語句的范圍內進行模型的編譯與定義即可：

tf.distribute.MirroredStrategy()
with strategy.scope():
    # 模型的定義與編譯
    ......

該 API 會將在 Scope 之中定義模型創(chuàng)建不同的鏡像，并將這些鏡像分不到不同的 GPU 之中去，從而實現分布式的訓練。

而之后的訓練步驟，我們只需要和平常進行訓練即可，剩下的工作 TensorFlow 會幫助我們完成。

3. 使用 tf.keras 模型進行單機器多 GPU 分布式訓練的程序實例

在這節(jié)課之中，我們會使用一個簡單的實例來進行演示使用多個 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 上的批次大?。?/li>
• 然后我們需要定義總體的批次大小 = 單個 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

可以看到，我們最終完成了訓練，說明我們的分布式訓練得到了正確的結果。

感興趣的同學可以嘗試一下不使用分布式訓練的結果和以上結果有什么區(qū)別。

4. 小結

在這節(jié)課之中，我們學習了什么是分布式訓練，同時了解了分布式訓練的種類。然后我們具體的了解了如何在同一臺機器上使用多 GPU 進行訓練，最后我們通過一個簡單的示例進行了演示。