如何進(jìn)行多 GPU 的分布式訓(xùn)練？

在之前的學(xué)習(xí)之中，我們進(jìn)行的訓(xùn)練都是在一臺(tái)機(jī)器上進(jìn)行的。或者更具體的說，我們之前的學(xué)習(xí)任務(wù)都是在同一臺(tái)機(jī)器的 CPU 或者一臺(tái)機(jī)器的同一個(gè) GPU 上進(jìn)行訓(xùn)練的。那么我們?nèi)绾尾拍軐C(jī)器之中的多個(gè) GPU 充分地利用起來呢？那么我們這節(jié)課就來學(xué)習(xí)一下如何進(jìn)行多 GPU 的分布式訓(xùn)練。

1. 什么是分布式訓(xùn)練

分布式訓(xùn)練，顧名思義，就是在多個(gè)設(shè)備之上進(jìn)行訓(xùn)練。它可以充分的使用硬件資源，使得訓(xùn)練與學(xué)習(xí)任務(wù)可以在更短的時(shí)間內(nèi)完成。

具體來說，分布式任務(wù)大體可以分為以下幾個(gè)模塊：

• 主程序?qū)⒕唧w的任務(wù)進(jìn)行分割，分割成多個(gè)小型的任務(wù)；
• 將分割后的小型任務(wù)分配到不同的設(shè)備之中去，并讓他們獨(dú)立執(zhí)行；
• 不同的設(shè)備在完成任務(wù)后會(huì)將產(chǎn)生的結(jié)果返回到主程序；
• 主程序會(huì)將結(jié)果進(jìn)行合并，從而得到最終的結(jié)果。

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

其實(shí)這個(gè)設(shè)備的含義很廣泛，它可以包括以下的含義：

• 同一臺(tái)計(jì)算機(jī)上的不同 GPU ；
• 不同計(jì)算機(jī)的不同 GPU ；
• 不同的計(jì)算機(jī)本身作為一個(gè)單獨(dú)的設(shè)備。

在實(shí)際的應(yīng)用過程之中，我們使用最多的情況是在同一臺(tái)設(shè)備上會(huì)有多張 GPU 卡，因此我們大多數(shù)的分布式訓(xùn)練實(shí)在同一臺(tái)機(jī)器上的不同 GPU 中進(jìn)行的。因此我們這節(jié)課著重介紹如何在多張 GPU 顯卡之中進(jìn)行機(jī)器學(xué)習(xí)的相關(guān)任務(wù)。

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

2. 如何使用 tf.keras 模型進(jìn)行單機(jī)器多 GPU 分布式訓(xùn)練

在 TensorFlow 之中進(jìn)行石激起多 GPU 顯卡進(jìn)行分布式訓(xùn)練的總體實(shí)現(xiàn)方式大致為：

• 定義要使用的 GPU ；
• 將完全一樣的模型放在不同的 GPU 之上并分別進(jìn)行編譯；
• 將數(shù)據(jù)送到不同的模型之中訓(xùn)練并得到相應(yīng)的結(jié)果；
• 主程序?qū)⒔Y(jié)果進(jìn)行整合。

這個(gè)過程看上去很復(fù)雜，但是在 TensorFlow 之中，我們只需要進(jìn)行第二部的操作即可。

那么我們?yōu)槭裁床挥枚x GPU 呢？因?yàn)?TensorFlow 的 GPU 版本會(huì)自動(dòng)檢測(cè) GPU ，并且進(jìn)行識(shí)別。

因此我們要使用的 API 為：

tf.distribute.MirroredStrategy()

在具體使用的時(shí)候，我們需要實(shí)例化一個(gè)對(duì)象，然后在以下語句的范圍內(nèi)進(jìn)行模型的編譯與定義即可：

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

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

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

在這節(jié)課之中，我們會(huì)使用一個(gè)簡(jiǎn)單的實(shí)例來進(jìn)行演示使用多個(gè) GPU 進(jìn)行訓(xùn)練。這里以我們之前使用過的 fashion_mnist 數(shù)據(jù)集的分類進(jìn)行演示。

具體的程序?yàn)椋?/p>

import tensorflow as tf

# 獲取fashion_mnist數(shù)據(jù)集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()

# 預(yù)處理圖片數(shù)據(jù)，使其歸一化
x_train, x_test = x_train / 255.0, x_test / 255.0

# 定義鏡像策略
strategy = tf.distribute.MirroredStrategy()

# 單個(gè)GPU之上的Batch數(shù)量
batch_per_gpu = 64
# 同時(shí)訓(xùn)練的Batch大小應(yīng)該為：?jiǎn)蝹€(gè)GPU的Batch * GPU個(gè)數(shù)
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)

這里網(wǎng)絡(luò)模型，我們采用了與之前相同的網(wǎng)絡(luò)模型。

值得注意的是，這邊的 Batch 大小：

• 我們首先定義單個(gè) GPU 上的批次大?。?/li>
• 然后我們需要定義總體的批次大小 = 單個(gè) GPU 上的批次大小 * GPU 數(shù)量；
• 在訓(xùn)練過程之中，TensorFlow 會(huì)將批次大小自動(dòng)分割分給不同的 GPU。

最終我們可以得到如下結(jié)果：

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

可以看到，我們最終完成了訓(xùn)練，說明我們的分布式訓(xùn)練得到了正確的結(jié)果。

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

4. 小結(jié)

在這節(jié)課之中，我們學(xué)習(xí)了什么是分布式訓(xùn)練，同時(shí)了解了分布式訓(xùn)練的種類。然后我們具體的了解了如何在同一臺(tái)機(jī)器上使用多 GPU 進(jìn)行訓(xùn)練，最后我們通過一個(gè)簡(jiǎn)單的示例進(jìn)行了演示。