基于 UAI-Train 平台的分布式訓練

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

TensorFlow 采用了資料流範式， 使用節點和邊的有向圖來表示計算。TensorFlow 需要使用者靜态聲明這種符号計算圖，并對該圖使用複寫和分區rewrite & partitioning，将其配置設定到機器上進行分布式執行。

TensorFlow 中的分布式機器學習訓練使用了如圖所示的參數伺服器方法 。

Cluster、Job、Task 關于 TensorFlow 的分布式訓練，主要概念包括 Cluster、Job、Task，其關聯關系如下：

• TensorFlow 分布式 Cluster 由多個 Task 組成，每個 Task 對應一個 train.Server 執行個體，作為 Cluster 的一個單獨節點；
• 多個相同作用的 Task 可以被劃分為一個 Job，在分布式深度學習架構中，我們一般把 Job 劃分為 Parameter Server （PS）和 Worker，Parameter Server Job 管理參數的存儲和更新工作，而 Worker Job 運作 OPs，作為計算節點隻執行計算密集型的 Graph 計算；
• Cluster 中的 Task 會相對進行通信，以便進行狀态同步、參數更新等操作，如果參數的數量過大，一台機器處理不了，這就要需要多個 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 機器還是在用上一次疊代的舊版參數值。

分布式訓練步驟

1. 指令行參數解析，擷取叢集的資訊 ps_hosts 和 worker_hosts，以及目前節點的角色資訊 job_name 和 task_index
2. 建立目前 Task 結點的 Server

   1. cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})

   2. server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)

3. 如果目前節點是Parameter Server，則調用join()無休止等待；如果是Worker，則執行下一步

   1. if FLAGS.job_name == "ps":

   2. server.join()

4. 建構要訓練的模型

   1. # build tensorflow graph model

5. 建立 train.Supervisor 來管理模型的訓練過程

   1. # Create a "supervisor", which oversees the training process.

   2. sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0), logdir="/tmp/train_logs")

   3. # The supervisor takes care of session initialization and restoring from a checkpoint.

   4. sess = sv.prepare_or_wait_for_session(server.target)

   5. # Loop until the supervisor shuts down

   6. while not sv.should_stop()

   7. # 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 做如下修改：

1. 添加相關參數：--data_dir、 --output_dir、 --work_dir、 --log_dir、 --num_gpu，UAI-Train 平台将會自動生成這些參數；
2. 在代碼中增加 UAI 參數：使用 data_dir 配置輸入檔案夾、使用 output_dir 配置輸出檔案夾。

具體案例代碼可以在

https://github.com/ucloud/uai-sdk/tree/master/examples/tensorflow/train/cifar

擷取。

在 UAI-Train 平台執行訓練

1. 根據 https://github.com/tensorflow/models/tree/master/tutorials/image/cifar10_estimator 的說明生成 CIFAR-10 的 tfrecords；
2. 使用 UAI-SDK 提供的 Python 生成 CIFAR-10 樣例的 Docker 鏡像；
3. 確定 Docker 鏡像已經上傳至 UHub，在 UAI-Train 平台上執行。

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

在 UAI 平台上的分布式訓練

CIFAR-10 樣例代碼使用 tf.estimator.Estimator API，隻需一個分布式環境和分布式環境配置，便可直接進行分布式訓練，該配置需要适用于 tf.estimator.Estimator API 的标準，即定義一個 TF_CONFIG 配置。

1. TF_CONFIG = {

2. "cluster":{

3. "master":["ip0:2222"],

4. "ps":["ip0:2223","ip1:2223"],

5. "worker":["ip1:2222"]},

6. "task":{"type":"worker","index":0},

7. "environment":"cloud"

8. }

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

1. /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 平台上進行分布式訓練。

原文釋出時間為：2018-05-24

本文來自雲栖社群合作夥伴“

Linux中國開源社群

”，了解相關資訊可以關注“

”。