如何借助分布式GPU環(huán)境來提升神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)的浮點(diǎn)計(jì)算能力

前言

在大型數(shù)據(jù)集上進(jìn)行訓(xùn)練的現(xiàn)代神經(jīng)網(wǎng)絡(luò)架構(gòu)，可以跨廣泛的多種領(lǐng)域獲取可觀的結(jié)果，涵蓋從圖像識(shí)別、自然語言處理，到欺詐檢測(cè)和推薦系統(tǒng)等各個(gè)方面。但訓(xùn)練這些神經(jīng)網(wǎng)絡(luò)模型需要大量浮點(diǎn)計(jì)算能力。雖然，近年來 GPU 硬件算力和訓(xùn)練方法上均取得了重大進(jìn)步，但在單一機(jī)器上，網(wǎng)絡(luò)訓(xùn)練所需要的時(shí)間仍然長(zhǎng)得不切實(shí)際，因此需要借助分布式GPU環(huán)境來提升神經(jīng)網(wǎng)絡(luò)訓(xùn)練系統(tǒng)的浮點(diǎn)計(jì)算能力。

TensorFlow分布式訓(xùn)練

TensorFlow 采用了數(shù)據(jù)流范式， 使用節(jié)點(diǎn)和邊的有向圖來表示計(jì)算。TensorFlow 需要用戶靜態(tài)聲明這種符號(hào)計(jì)算圖，并對(duì)該圖使用復(fù)寫和分區(qū)（rewrite & partitioning）將其分配到機(jī)器上進(jìn)行分布式執(zhí)行。

TensorFlow 中的分布式機(jī)器學(xué)習(xí)訓(xùn)練使用了如圖所示的參數(shù)服務(wù)器方法 。

Cluster、Job、Task

關(guān)于TensorFlow的分布式訓(xùn)練，主要概念包括Cluster、Job、Task，其關(guān)聯(lián)關(guān)系如下：

TensorFlow分布式Cluster由多個(gè)Task組成，每個(gè)Task對(duì)應(yīng)一個(gè)tf.train.Server實(shí)例，作為Cluster的一個(gè)單獨(dú)節(jié)點(diǎn)；

多個(gè)相同作用的Task可以被劃分為一個(gè)Job，在分布式深度學(xué)習(xí)框架中,我們一般把Job劃分為Parameter Server和Worker，Parameter Job是管理參數(shù)的存儲(chǔ)和更新工作，而Worker Job運(yùn)行OPs，作為計(jì)算節(jié)點(diǎn)只執(zhí)行計(jì)算密集型的Graph計(jì)算；

Cluster中的Task會(huì)相對(duì)進(jìn)行通信，以便進(jìn)行狀態(tài)同步、參數(shù)更新等操作，如果參數(shù)的數(shù)量過大，一臺(tái)機(jī)器處理不了，這就要需要多個(gè)Task。

TensorFlow分布式計(jì)算模式

In-graph 模式

In-graph模式，將模型計(jì)算圖的不同部分放在不同的機(jī)器上執(zhí)行。把計(jì)算從單機(jī)多GPU擴(kuò)展到了多機(jī)多GPU， 不過數(shù)據(jù)分發(fā)還是在一個(gè)節(jié)點(diǎn)。這樣配置簡(jiǎn)單， 多機(jī)多GPU的計(jì)算節(jié)點(diǎn)只需進(jìn)行join操作， 對(duì)外提供一個(gè)網(wǎng)絡(luò)接口來接受任務(wù)。訓(xùn)練數(shù)據(jù)的分發(fā)依然在一個(gè)節(jié)點(diǎn)上， 把訓(xùn)練數(shù)據(jù)分發(fā)到不同的機(jī)器上， 將會(huì)影響并發(fā)訓(xùn)練速度。在大數(shù)據(jù)訓(xùn)練的情況下， 不推薦使用這種模式。

Between-graph 模式

Between-graph模式下，數(shù)據(jù)并行，每臺(tái)機(jī)器使用完全相同的計(jì)算圖。訓(xùn)練的參數(shù)保存在參數(shù)服務(wù)器，數(shù)據(jù)不用分發(fā)，而是分布在各個(gè)計(jì)算節(jié)點(diǎn)自行計(jì)算， 把要更新的參數(shù)通知參數(shù)服務(wù)器進(jìn)行更新。這種模式不需要再練數(shù)據(jù)的分發(fā)， 數(shù)據(jù)量在TB級(jí)時(shí)可以節(jié)省大量時(shí)間，目前主流的分布式訓(xùn)練模式以 Between-graph為主。

參數(shù)更新方式

同步更新

各個(gè)用于并行計(jì)算的節(jié)點(diǎn)，計(jì)算完各自的batch 后，求取梯度值，把梯度值統(tǒng)一送到PS參數(shù)服務(wù)機(jī)器中，并等待PS更新模型參數(shù)。PS參數(shù)服務(wù)器在收集到一定數(shù)量計(jì)算節(jié)點(diǎn)的梯度后，求取梯度平均值，更新PS參數(shù)服務(wù)器上的參數(shù)，同時(shí)將參數(shù)推送到各個(gè)worker節(jié)點(diǎn)。

異步更新

PS參數(shù)服務(wù)器只要收到一臺(tái)機(jī)器的梯度值，就直接進(jìn)行參數(shù)更新，無需等待其它機(jī)器。這種迭代方法比較不穩(wěn)定，因?yàn)楫?dāng)A機(jī)器計(jì)算完更新了PS參數(shù)服務(wù)器中的參數(shù)，可能B機(jī)器還是在用上一次迭代的舊版參數(shù)值。

分布式訓(xùn)練步驟

命令行參數(shù)解析，獲取集群的信息ps_hosts和worker_hosts，以及當(dāng)前節(jié)點(diǎn)的角色信息job_name和task_index

創(chuàng)建當(dāng)前Task結(jié)點(diǎn)的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)

如果當(dāng)前節(jié)點(diǎn)是Parameter Server，則調(diào)用server.join()無休止等待；如果是Worker，則執(zhí)行下一步

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

構(gòu)建要訓(xùn)練的模型

# build tensorflow graph model

創(chuàng)建tf.train.Supervisor來管理模型的訓(xùn)練過程

# 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 downwhile not sv.should_stop() # train model

UAITrain分布式訓(xùn)練部署

UCloud AI 訓(xùn)練服務(wù)（UCloud AI Train）是面向AI訓(xùn)練任務(wù)的大規(guī)模分布式計(jì)算平臺(tái)，基于高性能GPU計(jì)算節(jié)點(diǎn)提供一站式托管AI訓(xùn)練任務(wù)服務(wù)。用戶在提交AI訓(xùn)練任務(wù)后，無需擔(dān)心計(jì)算節(jié)點(diǎn)調(diào)度、訓(xùn)練環(huán)境準(zhǔn)備、數(shù)據(jù)上傳下載以及容災(zāi)等問題。

目前，UAI--Train平臺(tái)支持TensorFlow 和 MXNet 框架的分布式訓(xùn)練。需要將PS代碼和Worker代碼實(shí)現(xiàn)在同一個(gè)代碼入口中，執(zhí)行過程中，PS 和 Worker 將使用相同的Docker容器鏡像和相同的python代碼入口進(jìn)行執(zhí)行，系統(tǒng)將自動(dòng)生成PS和Worker的env環(huán)境參數(shù)。TensorFlow 分布式訓(xùn)練采用PS-Worker的分布式格式，并提供python的接口運(yùn)行分布式訓(xùn)練。

UAI--Train分布式訓(xùn)練采用Parameter Server和Worker Server混合部署的方法，所有計(jì)算節(jié)點(diǎn)均由GPU物理云主機(jī)組成。PS 僅使用CPU進(jìn)行計(jì)算，Worker Server則同時(shí)使用GPU和CPU進(jìn)行計(jì)算，PS 和 Worker的比例為1:1。

數(shù)據(jù)存儲(chǔ)

分布式訓(xùn)練所使用的輸入數(shù)據(jù)可以來自不同的數(shù)據(jù)源，目前UAI--Train僅支持UFS作為數(shù)據(jù)的存儲(chǔ)。

Input 數(shù)據(jù)存儲(chǔ)

指定一個(gè)UFS網(wǎng)盤作為Input數(shù)據(jù)源，UAI--Train平臺(tái)在訓(xùn)練執(zhí)行過程中會(huì)將對(duì)應(yīng)的UFS數(shù)據(jù)映射到訓(xùn)練執(zhí)行的Worker容器的 /data/data 目錄下，系統(tǒng)會(huì)自動(dòng)將數(shù)據(jù)映射到執(zhí)行的容器中，如 ip:/xxx/data/imagenet/tf → /data/data/。

Output 數(shù)據(jù)存儲(chǔ)

指定一個(gè)UFS網(wǎng)盤作為output數(shù)據(jù)源，UAI--Train平臺(tái)在訓(xùn)練執(zhí)行過程中會(huì)將對(duì)應(yīng)的UFS數(shù)據(jù)映射到訓(xùn)練執(zhí)行的每一個(gè)PS容器和Worker容器的 /data/output 目錄下，并以共享的方式訪問同一份數(shù)據(jù)。同時(shí)，在訓(xùn)練過程可以通過其云主機(jī)實(shí)時(shí)訪問訓(xùn)練保存的模型checkpoint。

案例：通過CIFAR-10進(jìn)行圖像識(shí)別

CIFAR-10是機(jī)器學(xué)習(xí)中常見的圖像識(shí)別數(shù)據(jù)集，該數(shù)據(jù)集共有60000張彩色圖像。這些圖像，分為10個(gè)類，每類6000張圖，有50000張用于訓(xùn)練，另外10000用于測(cè)試。

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

調(diào)整訓(xùn)練代碼

為了在UAI平臺(tái)上進(jìn)行訓(xùn)練，首先下載源代碼，并對(duì)cifar10_main.py做如下修改：

添加相關(guān)參數(shù)：--data_dir, --output_dir, --work_dir, --log_dir, --num_gpus，UAITrain平臺(tái)將會(huì)自動(dòng)生成這些參數(shù)；

在代碼中增加UAI參數(shù)：使用data_dir配置輸入文件夾、使用output_dir配置輸出文件夾。

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

在UAI--Train平臺(tái)執(zhí)行訓(xùn)練

據(jù)https://github.com/tensorflow/models/tree/master/tutorials/image/cifar10_estimator 的說明生成CIFAR-10的tfrecords；

使用UAI-SDK提供的tf_tools.py 生成CIFAR-10樣例的Docker鏡像；

確保Docker鏡像已經(jīng)上傳至UHub，在UAI--Train平臺(tái)上執(zhí)行。

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

在UAI平臺(tái)上的分布式訓(xùn)練

CIFAR-10樣例代碼使用tf.estimator.Estimator API，只需一個(gè)分布式環(huán)境和分布式環(huán)境配置便可直接進(jìn)行分布式訓(xùn)練，該配置需要適用于tf.estimator.Estimator API的標(biāo)準(zhǔn)，即定義一個(gè)TF_CONFIG 配置。

TF_CONFIG = {

"cluster":{

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

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

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

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

"environment":"cloud"

}

UAITrain平臺(tái)的分布式訓(xùn)練功能可以自動(dòng)生成TensorFlow分布式訓(xùn)練的GPU集群環(huán)境，同時(shí)為每個(gè)訓(xùn)練節(jié)點(diǎn)自動(dòng)生成TF_CONFIG。因此，在UAITrain平臺(tái)上執(zhí)行CIFAR-10的分布式訓(xùn)練和單機(jī)訓(xùn)練一樣，僅需要指定input/output的UFS地址并執(zhí)行如下指令即可：

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

總結(jié)

UAI--Train TensorFlow的分布式訓(xùn)練環(huán)境實(shí)現(xiàn)基于TensorFlow 的分布式訓(xùn)練系統(tǒng)實(shí)現(xiàn)，采用默認(rèn)的grpc協(xié)議進(jìn)行數(shù)據(jù)交換。PS和Worker采用混合部署的方式部署，PS使用純CPU計(jì)算，Worker使用GPU+CPU計(jì)算。

在UAI--Train平臺(tái)中可以非常方便的開展分布式計(jì)算，提高效率、壓縮訓(xùn)練時(shí)間。最后通過CIFAR-10 案例解析在UAITrain平臺(tái)上訓(xùn)練所需做出的修改，并在UAITrain平臺(tái)上進(jìn)行分布式UAI--Train平臺(tái)訓(xùn)練。