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

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

Cluster、Job、Task

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

TensorFlow分布式計算模式

In-graph 模式

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

Between-graph 模式

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

參數(shù)更新方式

同步更新

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

異步更新

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

分布式訓練步驟

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

2.創(chuàng)建當前Task結(jié)點的Server



    

  1. cluster?=?tf.train.ClusterSpec({“ps”:?ps_hosts,?“worker”:?worker_hosts})
    2. 

  2. 

  3. server?=?tf.train.Server(cluster,?job_name=FLAGS.job_name,?task_index=FLAGS.task_index)
    4. 



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




    

  1. if?FLAGS.job_name?==?“ps”:
    2. 

  2. 

  3. ????server.join()
    4. 



4.則構(gòu)建要訓練的模型




    

  1. #?build?tensorflow?graph?model
    2. 



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




    

  1. #?Create?a?“supervisor”,?which?oversees?the?training?process.
    2. 

  2. sv?=?tf.train.Supervisor(is_chief=(FLAGS.task_index?==?0),?logdir=“/tmp/train_logs”)
    3. 

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

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

  5. #?Loop?until?the?supervisor?shuts?down
    6. 

  6. while?not?sv.should_stop()
    7. 

  7. ?????#?train?model
    8. 



UAI Train 分布式訓練部署


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


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




UCloud AI Train分布式訓練采用Parameter Server和Worker Server混合部署的方法，所有計算節(jié)點均由GPU物理云主機組成。PS 僅使用CPU進行計算，Worker Server則同時使用GPU和CPU進行計算，PS 和 Worker的比例為1:1。




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


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


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


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


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


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


案例研究：通過CIFAR-10進行圖像識別


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


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




調(diào)整訓練代碼


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


1.添加相關參數(shù)：–data_dir, –output_dir, –work_dir, –log_dir, –num_gpus，UAI Train平臺將會自動生成這些參數(shù)；


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


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


在UAI Train平臺執(zhí)行訓練


1.根據(jù)?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鏡像已經(jīng)上傳至UHub，在UAI Train平臺上執(zhí)行。




    

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



在UAI平臺上的分布式訓練


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




    

  1. TF_CONFIG?=?{
    2. 

  2. ????“cluster”:{
    3. 

  3. ????????“master”:[“ip0:2222”],
    4. 

  4. ????????“ps”:[“ip0:2223”,“ip1:2223”],
    5. 

  5. ????????“worker”:[“ip1:2222”]},
    6. 

  6. ????“task”:{“type”:“worker”,“index”:0},
    7. 

  7. ????“environment”:“cloud”
    8. 

  8. }
    9. 



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




    

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



總結(jié)


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


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

					
										
		
分享到

				

		
				

						
songjy

					

		
		            
        
				
相關推薦