基於TensorFlow原生的分散式開發
分散式開發會涉及到更新梯度的方式,有同步和非同步的兩個方案,同步更新的方式在模型的表現上能更快地進行收斂,而非同步更新時,迭代的速度則會更加快。兩種更新方式的圖示如下:
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=0
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的版本,則要自行處理,放寬執行時間。