Tensorflow學習筆記4：分布式Tensorflow

Tensorflow學習筆記4：分布式Tensorflow 簡介

Tensorflow API提供了Cluster、Server以及Supervisor來支援模型的分布式訓練。

關於Tensorflow的分布式訓練介紹可以參考Distributed Tensorflow。簡單的概括說明如下： Tensorflow分布式Cluster由多個Task組成，每個Task對應一個tf.train.Server執行個體，作為Cluster的一個單獨節點； 多個相同作用的Task可以被劃分為一個job，例如ps job作為參數伺服器只儲存Tensorflow model的參數，而worker job則作為計算節點只執行計算密集型的Graph計算。 Cluster中的Task會相對進行通訊，以便進行狀態同步、參數更新等操作。

Tensorflow分布式叢集的所有節點執行的代碼是相同的。分布式任務代碼具有固定的模式：

# 第1步：命令列參數解析，擷取叢集的資訊ps_hosts和worker_hosts，以及當前節點的角色資訊job_name和task_index# 第2步：建立當前task結點的Servercluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})server = tf.train.Server(cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)# 第3步：如果當前節點是ps，則調用server.join()無休止等待；如果是worker，則執行第4步。if FLAGS.job_name == "ps":    server.join()# 第4步：則構建要訓練的模型# build tensorflow graph model# 第5步：建立tf.train.Supervisor來管理模型的訓練過程# 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

  Tensorflow分布式訓練代碼架構

根據上面說到的Tensorflow分布式訓練代碼固定模式，如果要編寫一個分布式的Tensorlfow代碼，其架構如下所示。

import tensorflow as tf# Flags for defining the tf.train.ClusterSpectf.app.flags.DEFINE_string("ps_hosts", "",                           "Comma-separated list of hostname:port pairs")tf.app.flags.DEFINE_string("worker_hosts", "",                           "Comma-separated list of hostname:port pairs")# Flags for defining the tf.train.Servertf.app.flags.DEFINE_string("job_name", "", "One of 'ps', 'worker'")tf.app.flags.DEFINE_integer("task_index", 0, "Index of task within the job")FLAGS = tf.app.flags.FLAGSdef main(_):  ps_hosts = FLAGS.ps_hosts.split(",")  worker_hosts = FLAGS.worker_hosts(",")  # Create a cluster from the parameter server and worker hosts.  cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})  # Create and start a server for the local task.  server = tf.train.Server(cluster,                           job_name=FLAGS.job_name,                           task_index=FLAGS.task_index)  if FLAGS.job_name == "ps":    server.join()  elif FLAGS.job_name == "worker":    # Assigns ops to the local worker by default.    with tf.device(tf.train.replica_device_setter(        worker_device="/job:worker/task:%d" % FLAGS.task_index,        cluster=cluster)):      # Build model...      loss = ...      global_step = tf.Variable(0)      train_op = tf.train.AdagradOptimizer(0.01).minimize(          loss, global_step=global_step)      saver = tf.train.Saver()      summary_op = tf.merge_all_summaries()      init_op = tf.initialize_all_variables()    # Create a "supervisor", which oversees the training process.    sv = tf.train.Supervisor(is_chief=(FLAGS.task_index == 0),                             logdir="/tmp/train_logs",                             init_op=init_op,                             summary_op=summary_op,                             saver=saver,                             global_step=global_step,                             save_model_secs=600)    # The supervisor takes care of session initialization and restoring from    # a checkpoint.    sess = sv.prepare_or_wait_for_session(server.target)    # Start queue runners for the input pipelines (if any).    sv.start_queue_runners(sess)    # Loop until the supervisor shuts down (or 1000000 steps have completed).    step = 0    while not sv.should_stop() and step < 1000000:      # Run a training step asynchronously.      # See `tf.train.SyncReplicasOptimizer` for additional details on how to      # perform *synchronous* training.      _, step = sess.run([train_op, global_step])if __name__ == "__main__":  tf.app.run()

對於所有Tensorflow分布式代碼，可變的只有兩點： 構建tensorflow graph模型代碼； 每一步執行訓練的代碼 分布式MNIST任務

我們通過修改tensorflow/tensorflow提供的mnist_softmax.py來構造分布式的MNIST範例來進行驗證。修改後的代碼請參考mnist_dist.py。

我們同樣通過tensorlfow的Docker image來啟動一個容器來進行驗證。

$ docker run -d -v /path/to/your/code:/tensorflow/mnist --name tensorflow tensorflow/tensorflow

啟動tensorflow之後，啟動4個Terminal，然後通過下面命令進入tensorflow容器，切換到/tensorflow/mnist目錄下

$ docker exec -ti tensorflow /bin/bash$ cd /tensorflow/mnist

然後在四個Terminal中分別執行下面一個命令來啟動Tensorflow cluster的一個task節點，

# Start ps 0python mnist_dist.py --ps_hosts=localhost:2221,localhost:2222 --worker_hosts=localhost:2223,localhost:2224 --job_name=ps --task_index=0# Start ps 1python mnist_dist.py --ps_hosts=localhost:2221,localhost:2222 --worker_hosts=localhost:2223,localhost:2224 --job_name=ps --task_index=1# Start worker 0python mnist_dist.py --ps_hosts=localhost:2221,localhost:2222 --worker_hosts=localhost:2223,localhost:2224 --job_name=worker --task_index=0# Start worker 1python mnist_dist.py --ps_hosts=localhost:2221,localhost:2222 --worker_hosts=localhost:2223,localhost:2224 --job_name=worker --task_index=1

具體效果自己驗證哈。