歡迎光臨
每天分享高質量文章

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

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


致謝

 作者 | Ucloud

在大型資料集上進行訓練的現代神經網路架構,可以跨廣泛的多種領域獲取可觀的結果,涵蓋從影象識別、自然語言處理到欺詐檢測和推薦系統等各個方面,但訓練這些神經網路模型需要大量浮點計算能力。雖然,近年來 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)和 WorkerParameter 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 機器還是在用上一次迭代的舊版引數值。

分散式訓練步驟

☉ 命令列引數解析,獲取叢集的資訊 ps_hosts 和 worker_hosts,以及當前節點的角色資訊 job_name 和 task_index
☉ 建立當前 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)

☉ 如果當前節點是Parameter Server,則呼叫join()無休止等待;如果是Worker,則執行下一步

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

  2.     server.join()

☉ 構建要訓練的模型

  1. # build tensorflow graph model

☉ 建立 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 做如下修改:

☉ 新增相關引數:–data_dir、 –output_dir、 –work_dir、 –log_dir、 –num_gpu,UAI-Train 平臺將會自動生成這些引數;
☉ 在程式碼中增加 UAI 引數:使用 data_dir 配置輸入檔案夾、使用 output_dir 配置輸出檔案夾。

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

在 UAI-Train 平臺執行訓練

☉ 根據 https://github.com/tensorflow/models/tree/master/tutorials/image/cifar10_estimator 的說明生成 CIFAR-10 的 tfrecords;
☉ 使用 UAI-SDK 提供的 Python 生成 CIFAR-10 樣例的 Docker 映象;
☉ 確保 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 平臺上進行分散式訓練。

贊(0)

分享創造快樂