TensorFlow分布式訓練¶

當我們擁有大量計算資源時，透過使用合適的分散式策略，我們可以充分利用這些計算資源，大幅壓縮模型訓練的時間。針對不同的使用場景，TensorFlow 在 tf.distribute.Strategy 中為我們提供了許多種分散式策略，使得我們能夠更高效的訓練模型。

單機多卡訓練： `MirroredStrategy`¶

tf.distribute.MirroredStrategy 是一種簡單且高性能的，資料並行的同步式分散式策略，主要支援多個 GPU 在同一台主機上訓練。使用這種策略時，我們只需實例化一個 MirroredStrategy 策略:

strategy = tf.distribute.MirroredStrategy()

並將模型建構的程式碼放入 strategy.scope() 的上下文環境中:

with strategy.scope():
    # 模型建構程式碼

小技巧

可以在參數中指定設備，如:

strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

即指定只使用第0、1號GPU參與分散式策略。

以下程式碼展示了使用 MirroredStrategy 策略，在 TensorFlow Datasets 中的部分圖像資料集上使用 Keras 訓練 MobileNetV2 的過程：

import tensorflow as tf
import tensorflow_datasets as tfds

num_epochs = 5
batch_size_per_replica = 64
learning_rate = 0.001

strategy = tf.distribute.MirroredStrategy()
print('Number of devices: %d' % strategy.num_replicas_in_sync)  # 輸出設備數量
batch_size = batch_size_per_replica * strategy.num_replicas_in_sync

# 載入資料集並預處理
def resize(image, label):
    image = tf.image.resize(image, [224, 224]) / 255.0
    return image, label

# 使用 TensorFlow Datasets 載入貓狗分類資料集，詳見“TensorFlow Datasets資料集載入”一章
dataset = tfds.load("cats_vs_dogs", split=tfds.Split.TRAIN, as_supervised=True)
dataset = dataset.map(resize).shuffle(1024).batch(batch_size)

with strategy.scope():
    model = tf.keras.applications.MobileNetV2()
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy]
    )

model.fit(dataset, epochs=num_epochs)

在以下的測試中，我們使用同一台主機上的 4 塊 NVIDIA GeForce GTX 1080 Ti 顯卡進行單機多卡的模型訓練。所有測試的 epoch 數均為 5。使用單機無分散式配置時，雖然機器依然具有 4 塊顯卡，但程式不使用分散式的設置，直接進行訓練，Batch Size 設置為 64。使用單機四卡時，測試總 Batch Size 為 64（分發到單台機器的 Batch Size 為 16）和總 Batch Size 為 256（分發到單台機器的 Batch Size 為 64）兩種情況。

資料集	單機無分散式（Batch Size爲64）	單機四卡（總Batch Size爲64）	單機四卡（總Batch Size爲256）
cats_vs_dogs	146s/epoch	39s/epoch	29s/epoch
tf_flowers	22s/epoch	7s/epoch	5s/epoch

可見，使用 MirroredStrategy 後，模型訓練的速度有了大幅度的提高。在所有顯卡性能接近的情況下，訓練時間與顯卡的數目接近於反比關係。

MirroredStrategy 過程簡介

MirroredStrategy的步驟如下：

訓練開始前，該策略在所有 N 個計算設備上均各複製一份完整的模型；
每次訓練傳入一個批量的資料時，將資料分成 N 份，分別傳入 N 個計算設備（即資料平行處理）；
N 個計算設備使用區域變數（鏡像變數）分別計算自己所獲得部分的資料梯度；
使用分散式計算的 All-reduce 操作，在計算設備間高效的，交換梯度資料並進行求和，使得最終每個設備都有了所有設備的梯度和；
使用梯度求和的結果更新區域變數（鏡像變數）；
當所有設備均更新區域變數後，進行下一輪訓練（該並行策略是同步的）。

默認情況下，TensorFlow 中的 MirroredStrategy 策略使用 NVIDIA NCCL 進行 All-reduce 操作。

多機訓練： `MultiWorkerMirroredStrategy`¶

多機訓練的方法和單機多卡類似，將 MirroredStrategy 更換為適合多機訓練的 MultiWorkerMirroredStrategy 即可。不過，由於涉及到多台電腦之間的通訊，還需要進行一些額外的設置。具體而言，需要設置環境變數 TF_CONFIG ，範例如下:

os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:20000", "localhost:20001"]
    },
    'task': {'type': 'worker', 'index': 0}
})

TF_CONFIG 由 cluster 和 task 兩部分組成：

cluster 說明了整個多機集群的結構和每台機器的網路位置（IP + 埠號）。對於每一台機器，cluster 的值都是相同的；
task 說明了目前機器的角色。例如， {'type': 'worker', 'index': 0} 說明目前機器是 cluster 中的第 0 個 worker（即 localhost:20000 ）。每一台機器的 task 值都需要針對目前主機進行分別的設置。

以上內容設置完成後，在所有的機器上逐一執行訓練程式碼即可。先執行的程式碼在尚未與其他主機連接時會進入監聽狀態，待整個集群的連接建立完畢後，所有的機器即會同時開始訓練。

提示

請在各台機器上均注意防火牆的設置，尤其是需要開放與其他主機通信的埠號。如上例的 0 號 worker 需要開放 20000 埠號，1 號 worker 需要開放 20001 埠號。

以下範例的訓練任務與之前章節相同，只不過遷移到了多機訓練環境。假設我們有兩台機器，即首先在兩台機器上均部署下面的程式，唯一的區別是 task 部分，第一台機器設置為 {'type': 'worker', 'index': 0} ，第二台機器設置為 {'type': 'worker', 'index': 1} 。接下來，在兩台機器上依序執行程式，待通訊成功後，即會自動開始訓練流程。

import tensorflow as tf
import tensorflow_datasets as tfds
import os
import json

num_epochs = 5
batch_size_per_replica = 64
learning_rate = 0.001

num_workers = 2
os.environ['TF_CONFIG'] = json.dumps({
    'cluster': {
        'worker': ["localhost:20000", "localhost:20001"]
    },
    'task': {'type': 'worker', 'index': 0}
})
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
batch_size = batch_size_per_replica * num_workers

def resize(image, label):
    image = tf.image.resize(image, [224, 224]) / 255.0
    return image, label

dataset = tfds.load("cats_vs_dogs", split=tfds.Split.TRAIN, as_supervised=True)
dataset = dataset.map(resize).shuffle(1024).batch(batch_size)

with strategy.scope():
    model = tf.keras.applications.MobileNetV2()
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss=tf.keras.losses.sparse_categorical_crossentropy,
        metrics=[tf.keras.metrics.sparse_categorical_accuracy]
    )

model.fit(dataset, epochs=num_epochs)

在以下測試中，我們在 Google Cloud Platform 分別建立兩台具有單張 NVIDIA Tesla K80 的虛擬機實例（具體建立方式參見後文介紹），並分別測試在使用一個 GPU 時的訓練時間和使用兩台虛擬機實例進行分散式訓練的訓練時間。所有測試的 epoch 數均為 5。使用單機單卡時，Batch Size 設置為 64。使用雙機單卡時，測試總 Batch Size 為 64（分發到單台機器的 Batch Size 為 32）和總 Batch Size 為 128（分發到單台機器的 Batch Size 為 64）兩種情況。

資料集	單機單卡（Batch Size爲64）	雙機單卡（總Batch Size爲64）	雙機單卡（總Batch Size爲128）
cats_vs_dogs	1622s	858s	755s
tf_flowers	301s	152s	144s

可見模型訓練的速度同樣有大幅度的提高。在所有機器性能接近的情況下，訓練時間與機器的數目接近於反比關係。

TensorFlow分布式訓練¶

單機多卡訓練： MirroredStrategy¶

多機訓練： MultiWorkerMirroredStrategy¶

單機多卡訓練： `MirroredStrategy`¶

多機訓練： `MultiWorkerMirroredStrategy`¶