技術乾貨 | 基於UAI-Train平臺的分佈式訓練

在大型數據集上進行訓練的現代神經網絡架構，能夠跨普遍的多種領域獲取可觀的結果，涵蓋從圖像識別、天然語言處理到欺詐檢測和推薦系統等各個方面，但訓練這些神經網絡模型須要大量浮點計算能力。雖然，近年來 GPU 硬件算力和訓練方法上均取得了重大進步，但在單一機器上，網絡訓練所須要的時間仍然長得不切實際，所以須要藉助分佈式GPU環境來提高神經網絡訓練系統的浮點計算能力。

TensorFlow分佈式訓練

（TensorFlow分佈式訓練概述圖）

TensorFlow 採用了數據流範式， 使用節點和邊的有向圖來表示計算。TensorFlow 須要用戶靜態聲明這種符號計算圖，並對該圖使用複寫和分區（rewrite & partitioning），將其分配到機器上進行分佈式執行。

（TensorFlow訓練數據流轉圖）

TensorFlow 中的分佈式機器學習訓練使用瞭如圖所示的參數服務器方法 。

Cluster、Job、Task

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

**1.**TensorFlow分佈式Cluster由多個Task組成，每一個Task對應一個tf.train.Server實例，做爲Cluster的一個單獨節點；**2.**多個相同做用的Task能夠被劃分爲一個Job，在分佈式深度學習框架中,咱們通常把Job劃分爲Parameter Server和Worker，Parameter Job是管理參數的存儲和更新工做，而Worker Job運行OPs，做爲計算節點只執行計算密集型的Graph計算；**3.**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

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)

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

if FLAGS.job_name == 「ps」:server.join()

4. 構建要訓練的模型

# build tensorflow graph model

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

# 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的接口運行分佈式訓練。

（圖：AI訓練平臺概述）

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

（圖：AI 訓練平臺分佈式訓練集羣部署範例）

數據存儲

分佈式訓練所使用的輸入數據能夠來自不一樣的數據源，目前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/

（圖：CIFAR-10數據集簡介）

調整訓練代碼

爲了在UAI平臺上進行訓練，首先下載源代碼，並對cifar10_main.py作以下修改：

1. 添加相關參數：–data_dir, –output_dir, –work_dir, –log_dir, –num_gpus，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提供的tf_tools.py 生成CIFAR-10樣例的Docker鏡像；

3. 確保Docker鏡像已經上傳至UHub，在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平臺上進行分佈式訓練。

想要獲取更多技術和活動資訊，可掃描如下二維碼，關注「UCloud技術公告牌」微信公衆號；或搜索微信ID：ucloud_tech進行關注。