基於 UAI-Train 平臺的分散式訓練

在大型資料集上進行訓練的現代神經網路架構，可以跨廣泛的多種領域獲取可觀的結果，涵蓋從影象識別、自然語言處理到欺詐檢測和推薦系統等各個方面，但訓練這些神經網路模型需要大量浮點計算能力。雖然，近年來 GPU 硬體算力和訓練方法上均取得了重大進步，但在單一機器上，網路訓練所需要的時間仍然長得不切實際，因此需要藉助分散式 GPU 環境來提升神經網路訓練系統的浮點計算能力。

TensorFlow 分散式訓練

TensorFlow 採用了資料流正規化， 使用節點和邊的有向圖來表示計算。TensorFlow 需要使用者靜態宣告這種符號計算圖，並對該圖使用複寫和分割槽rewrite & partitioning，將其分配到機器上進行分散式執行。 

TensorFlow 中的分散式機器學習訓練使用瞭如圖所示的引數伺服器方法 。

Cluster、Job、Task

關於 TensorFlow 的分散式訓練，主要概念包括 Cluster、Job、Task，其關聯關係如下：

TensorFlow 分散式計算樣式

In-graph 樣式

In-graph 樣式，將模型計算圖的不同部分放在不同的機器上執行。把計算從單機多 GPU 擴充套件到了多機多 GPU， 不過資料分發還是在一個節點。這樣配置簡單， 多機多 GPU 的計算節點只需進行 join 操作， 對外提供一個網路介面來接受任務。訓練資料的分發依然在一個節點上， 把訓練資料分發到不同的機器上， 將會影響併發訓練速度。在大資料訓練的情況下， 不推薦使用這種樣式。

Between-graph 樣式

Between-graph 樣式下，資料並行，每臺機器使用完全相同的計算圖。訓練的引數儲存在引數伺服器，資料不用分發，而是分佈在各個計算節點自行計算， 把要更新的引數通知引數伺服器進行更新。這種樣式不需要再練資料的分發， 資料量在 TB 級時可以節省大量時間，目前主流的分散式訓練樣式以 Between-graph 為主。

引數更新方式

同步更新

各個用於平行計算的節點，計算完各自的 batch 後，求取梯度值，把梯度值統一送到 PS 引數服務機器中，並等待 PS 更新模型引數。PS 引數伺服器在收集到一定數量計算節點的梯度後，求取梯度平均值，更新PS引數伺服器上的引數，同時將引數推送到各個 worker 節點。

非同步更新

PS 引數伺服器只要收到一臺機器的梯度值，就直接進行引數更新，無需等待其它機器。這種迭代方法比較不穩定，因為當 A 機器計算完更新了 PS 引數伺服器中的引數，可能 B 機器還是在用上一次迭代的舊版引數值。

分散式訓練步驟

cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)

if FLAGS.job_name == "ps":
    server.join()

# build tensorflow graph model

# Create a "supervisor", which oversees the training process.
sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0), logdir="/tmp/train_logs")
# The supervisor takes care of session initialization and restoring from a checkpoint.
sess = sv.prepare_or_wait_for_session(server.target)
# Loop until the supervisor shuts down
while not sv.should_stop()
     # train model

UAI-Train 分散式訓練部署

UCloud AI 訓練服務（UCloud AI Train）是面向 AI 訓練任務的大規模分散式計算平臺，基於高效能 GPU 計算節點提供一站式託管 AI 訓練任務服務。使用者在提交 AI 訓練任務後，無需擔心計算節點排程、訓練環境準備、資料上傳下載以及容災等問題。

目前，UAI-Train 平臺支援 TensorFlow 和 MXNet 框架的分散式訓練。需要將 PS 程式碼和 Worker 程式碼實現在同一個程式碼入口中，執行過程中，PS 和 Worker 將使用相同的 Docker 容器映象和相同的 Python 程式碼入口進行執行，系統將自動生成 PS 和 Worker 的 env 環境引數。TensorFlow 分散式訓練採用 PS-Worker 的分散式格式，並提供 Python 的介面執行分散式訓練。

UAI-Train 分散式訓練採用 Parameter Server 和 Worker 混合部署的方法，所有計算節點均由 GPU 物理雲主機組成。PS 僅使用 CPU 進行計算，Worker 則同時使用 GPU 和 CPU 進行計算，PS 和 Worker 的比例為 1:1。

資料儲存

分散式訓練所使用的輸入資料可以來自不同的資料源，目前 UAI-Train 僅支援 UFS 作為資料的儲存。

Input 資料儲存

指定一個 UFS 網盤作為 Input 資料源，UAI-Train 平臺在訓練執行過程中會將對應的 UFS 資料對映到訓練執行的 Worker 容器的 /data/data 目錄下，系統會自動將資料對映到執行的容器中，如 ip:/xxx/data/imagenet/tf → /data/data/。

Output 資料儲存

指定一個 UFS 網盤作為 output 資料源，UAI-Train 平臺在訓練執行過程中會將對應的 UFS 資料對映到訓練執行的每一個 PS 容器和 Worker 容器的 /data/output 目錄下，並以共享的方式訪問同一份資料。同時，在訓練過程，可以透過其它雲主機實時訪問訓練儲存的模型 checkpoint。

案例分析：透過 CIFAR-10 進行影象識別

CIFAR-10 是機器學習中常見的影象識別資料集，該資料集共有 60000 張彩色影象。這些影象分為 10 個類，每類 6000 張圖，有 50000 張用於訓練，另外 10000 用於測試。

http://groups.csail.mit.edu/vision/TinyImages/

調整訓練程式碼

為了在 UAI 平臺上進行訓練，首先下載原始碼，並對 cifar10_main.py 做如下修改：

具體案例程式碼可以在 https://github.com/ucloud/uai-sdk/tree/master/examples/tensorflow/train/cifar 獲取。

在 UAI-Train 平臺執行訓練

/data/cifar10_main.py --train-batch-size=16

在 UAI 平臺上的分散式訓練

CIFAR-10 樣例程式碼使用 tf.estimator.Estimator API，只需一個分散式環境和分散式環境配置，便可直接進行分散式訓練，該配置需要適用於 tf.estimator.Estimator API 的標準，即定義一個 TF_CONFIG 配置。

TF_CONFIG = {
    "cluster":{
        "master":["ip0:2222"],
        "ps":["ip0:2223","ip1:2223"],
        "worker":["ip1:2222"]},
    "task":{"type":"worker","index":0},
    "environment":"cloud"
}

UAI-Train 平臺的分散式訓練功能可以自動生成 TensorFlow 分散式訓練的 GPU 叢集環境，同時為每個訓練節點自動生成TF_CONFIG。因此，在 UAI-Train 平臺上執行 CIFAR-10 的分散式訓練和單機訓練一樣，僅需要指定 input/output 的 UFS 地址並執行如下指令即可：

/data/cifar10_main.py --train-batch-size=16

總結

UAI-Train TensorFlow 的分散式訓練環境實現基於 TensorFlow 的分散式訓練系統實現，採用預設的 gRPC 協議進行資料交換。PS 和 Worker 採用混合部署的方式部署，PS 使用純 CPU 計算，Worker 使用 GPU+CPU 計算。

在 UAI-Train 平臺中可以非常方便的開展分散式計算，提高效率、壓縮訓練時間。最後透過 CIFAR-10 案例解析在 UAI-Train 平臺上進行訓練所需作出的修改，併在 UAI-Train 平臺上進行分散式訓練。