TensorFlow分散式實踐

個推2018發表於2019-01-16
分散式

大資料時代，基於單機的建模很難滿足企業不斷增長的資料量級的需求，開發者需要使用分散式的開發方式，在叢集上進行建模。而單機和分散式的開發程式碼有一定的區別，本文就將為開發者們介紹，基於TensorFlow進行分散式開發的兩種方式，幫助開發者在實踐的過程中，更好地選擇模組的開發方向。

基於TensorFlow原生的分散式開發


分散式開發會涉及到更新梯度的方式，有同步和非同步的兩個方案，同步更新的方式在模型的表現上能更快地進行收斂，而非同步更新時，迭代的速度則會更加快。兩種更新方式的圖示如下：

同步更新流程

（圖片來源：TensorFlow:Large-Scale Machine Learning on Heterogeneous Distributed Systems）

非同步更新流程

（圖片來源：TensorFlow:Large-Scale Machine Learning on Heterogeneous Distributed Systems）


TensorFlow是基於ps、work 兩種伺服器進行分散式的開發。ps伺服器可以只用於引數的彙總更新，讓各個work進行梯度的計算。

基於TensorFlow原生的分散式開發的具體流程如下：

首先指定ps 伺服器啟動引數 –job_name=ps: 

python distribute.py --ps_hosts=192.168.100.42:2222 --worker_hosts=192.168.100.42:2224,192.168.100.253:2225 --job_name=ps --task_index=0

接著指定work伺服器引數(啟動兩個work 節點) –job_name=work2: 



python

 distribute.py --ps_hosts=

192.168.100.42:2222

 --worker_hosts=

192.168.100.42:2224

,

192.168.100.253:2225

 --job_name=worker --task_index=




python distribute.py --ps_hosts=

192.168.100.42:2222

 --worker_hosts=

192.168.100.42:2224

,

192.168.100.253:2225

 --job_name=worker --task_index=

1

之後，上述指定的引數 worker_hosts ps_hosts job_name task_index 都需要在py檔案中接受使用： 

tf.app.flags.DEFINE_string("worker_hosts", "預設值", "描述說明")

接收引數後，需要分別註冊ps、work，使他們各司其職： 


ps_hosts = FLAGS.ps_hosts.

split

(

","

)


worker_hosts = FLAGS.worker_hosts.

split

(

","

)


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)




issync = FLAGS.issync



if

 FLAGS.job_name == 

"ps"

:


   server.

join

()


elif FLAGS.job_name == 

"worker"

:


   with 

tf

.device(

tf

.train.replica_device_setter(


                   worker_device=

"/job:worker/task:%d"

 % FLAGS.task_index,


                   cluster=cluster)):

繼而更新梯度。

（1）同步更新梯度： 


rep_op = 

tf

.train.SyncReplicasOptimizer(optimizer,


                                               replicas_to_aggregate=

len

(worker_hosts),


                                               replica_id=FLAGS.task_index,


                                               total_num_replicas=

len

(worker_hosts),


                                               use_locking=True)


train_op = rep_op.apply_gradients(grads_and_vars,global_step=global_step)


init_token_op = rep_op.get_init_tokens_op()


chief_queue_runner = rep_op.get_chief_queue_runner()

（2）非同步更新梯度： 

train_op = optimizer.apply_gradients(grads_and_vars,global_step=global_step)

最後，使用tf.train.Supervisor 進行真的迭代

另外，開發者還要注意，如果是同步更新梯度，則還需要加入如下程式碼： 



sv


.start_queue_runners

(

sess

, 

[chief_queue_runner]

)



sess


.run

(

init_token_op

)

需要注意的是，上述非同步的方式需要自行指定叢集IP和埠，不過，開發者們也可以藉助TensorFlowOnSpark，使用Yarn進行管理。

基於TensorFlowOnSpark的分散式開發

作為個推面向開發者服務的移動APP資料統計分析產品，個數所具有的使用者行為預測功能模組，便是基於TensorFlowOnSpark這種分散式來實現的。基於TensorFlowOnSpark的分散式開發使其可以在遮蔽了埠和機器IP的情況下，也能夠做到較好的資源申請和分配。而在多個千萬級應用同時建模的情況下，叢集也有良好的表現，在sparkUI中也能看到相對應的資源和程式的情況。最關鍵的是，TensorFlowOnSpark可以在單機過度到分散式的情況下，使程式碼方便修改，且容易部署。

基於TensorFlowOnSpark的分散式開發的具體流程如下：

首先，需要使用spark-submit來提交任務，同時指定spark需要執行的引數（–num-executors 6等）、模型程式碼、模型超參等，同樣需要接受外部引數： 


parser = argparse.ArgumentParser()


parser.add_argument(

"-i"

, 

"--tracks"

, 

help

=

"資料集路徑"

)  



args

 = parser.parse_args()

之後，準備好引數和訓練資料(DataFrame)，呼叫模型的API進行啟動。

其中，soft_dist.map_fun是要調起的方法，後面均是模型訓練的引數。 


estimator = TFEstimator(soft_dist.map_fun, 

args

) \


     .setInputMapping({

'tracks'

: 

'tracks'

, 

'label'

: 

'label'

}) \


     .setModelDir(

args

.model) \


     .setExportDir(

args

.serving) \


     .setClusterSize(

args

.cluster_size) \


     .setNumPS(num_ps) \


     .setEpochs(

args

.epochs) \


     .setBatchSize(

args

.batch_size) \


     .setSteps(

args

.max_steps)


   model = estimator.fit(df)

接下來是soft_dist定義一個 map_fun(args, ctx)的方法： 




def

 

map_fun


(args, ctx)

:



...


worker_num = ctx.worker_num  

# worker數量



job_name = ctx.job_name  

# job名



task_index = ctx.task_index  

# 任務索引




if

 job_name == 

"ps"

:  

# ps節點(主節點)



  time.sleep((worker_num + 

1

) * 

5

)


  cluster, server = TFNode.start_cluster_server(ctx, 

1

, args.rdma)


  num_workers = len(cluster.as_dict()[

'worker'

])


  

if

 job_name == 

"ps"

:


       server.join()


  

elif

 job_name == 

"worker"

:


       

with

 tf.device(tf.train.replica_device_setter(worker_device=

"/job:worker/task:%d"

 % task_index, cluster=cluster)):

之後，可以使用tf.train.MonitoredTrainingSession高階API，進行模型訓練和預測。

總結

基於TensorFlow的分散式開發大致就是本文中介紹的兩種情況，第二種方式可以用於實際的生產環境，穩定性會更高。

在執行結束的時候，開發者們也可透過設定郵件的通知，及時地瞭解到模型執行的情況。

同時，如果開發者使用SessionRunHook來儲存最後輸出的模型，也需要了解到，框架程式碼中的一個BUG，即它只能在規定的時間內儲存，超出規定時間，即使執行沒有結束，程式也會被強制結束。如果開發者使用的版本是未修復BUG的版本，則要自行處理，放寬執行時間。


來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/31556026/viewspace-2563692/，如需轉載，請註明出處，否則將追究法律責任。

相關文章