本篇博客主要介绍TF的分布式训练，重点从代码层面进行讲解。理论部分可以参考深度学习分布式训练实战（一)
TF的分布式实现方式
TF的分布式有两种实现方式，一种是图内分布式（In-graph replication）；一种是图间分布式(Between-graph replication)。图内分布式中，计算图只有一个，需要一个中心节点分配计算任务并更新参数，由于中心节点的存在，中心节点容易成为瓶颈。图间分布式中，计算图有多个，但是不同计算图的相同变量通过tf.train.replica_device_setter函数放到同一个服务器上，这种情况下，各个计算图相互独立（参数只有一份，计算图有多个），并行度更高，适合异步更新，同步更新下相对麻烦，不过TF给了接口tf.train.SyncReplicasOptimizer函数来帮助实现参数的同步更新，所以图间分布式应用相对广泛一些。
关于数据并行，模型并行可以参考深度学习分布式训练实战（一)
大部分情况下，我们使用图间分布式，图内分布式一般只会在模型太大的情况下使用。对于图间分布式，其基于gRPC通信框架，模型参数只有一份，计算图有多份，一个master负责创建主session，多个worker执行计算图任务。模型训练过程中，每个计算图计算出各自梯度，然后对参数进行更新。更新方式有两种：同步更新，异步更新。
分布式TF中，TF需要建立一个集群，然后在集群中建立两个job，一个是ps job，负责参数初始化，参数更新，一个job下面可以有多个task（有多个task，说明有多台机器，或者GPU负责参数初始化，更新）。一个是woker job，负责计算图的运算，计算梯度，一个worker job下面也可以有很多个task（有多个task，说明有多台机器，或者GPU负责运行计算图）。
参数异步更新的分布式训练
参数同步更新基本上和这里写的差不多TensorFlow分布式部署
。只不过为了方便在本机上调试，所以改了一点点。（自己的笔记本没有GPU），介绍下面几个重点的语句：
tf.train.ClusterSpec()：创建一个集群对象
tf.train.Server()：在这个集群上面创建一个服务器，根据实际情况，可以是参数服务器，也可以是计算服务器
tf.train.Supervisor()：创建一个监视器，就是用来监控训练过程的，个人感觉主要就是方便恢复模型训练，其logdir参数为训练日志目录，如果里面有模型，则直接恢复训练。所以如果想重新训练，需要删除这个目录。
sv.managed_session()：启动Session，相比于其他启动Session的方式，多了一些功能。可以参考TensorFlow 中三种启动图用法
具体代码如下：
# tensorflow distribute train by asynchronously update
import tensorflow as tf
import numpy as np
tf.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")
tf.app.flags.DEFINE_string("job_name", "", "one of ps or worker")
tf.app.flags.DEFINE_integer("task_index", 0, "0, 1, 2...")
FLAGS = tf.app.flags.FLAGS
def main(_):
ps_hosts = FLAGS.ps_hosts.split(",")
worker_hosts = FLAGS.worker_hosts.split(",")
# Create a cluster from the parameter server and worker server
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)
# 如果是参数服务器，则直接阻塞，等待计算服务器下达参数初始化，参数更新命令就可以了。
# 不过“下达命令”这个是TF内部实现的，没有显式实现
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)):
train_X = np.linspace(-1.0, 1.0, 100)
train_Y = 2.0 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10.0
X = tf.placeholder("float")
Y = tf.placeholder("float")
w = tf.Variable(0.0, name="weight")
b = tf.Variable(0.0, name="bias")
loss = tf.square(Y - tf.multiply(X, w) - b)
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.summary.merge_all()
init_op = tf.global_variables_initializer()
# Create a "supervisor", which oversees the training process.
sv = tf.train.Supervisor(is_chief=(FLAGS.task_index==0),
logdir="~/Downloads/log/",
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, retoring from a
# checkpoint, and closing when done or an error occurs.
with sv.managed_session(server.target) as sess:
step = 0
while step < 1000000:
# Run a training step asynchronously
for (x, y) in zip(train_X, train_Y):
_, step =sess.run([train_op, global_step], feed_dict={X:x, Y:y})
loss_value = sess.run(loss, feed_dict={X:x, Y:y})
print("Step: {}, loss: {}".format(step, loss_value))
# Ask for all the services to stop
sv.stop()
if __name__=="__main__":
tf.app.run()
打开三个终端，分别输入以下三个命令，就可以看到训练结果了：
CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=ps --task_index=0
CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=worker --task_index=0
CUDA_VISIBLE_DEVICES='' python AsynDis.py --ps_hosts=localhost:2222 --worker_hosts=localhost:2224,localhost:2225 --job_name=worker --task_index=1
这里有一点要注意，控制计算图使用哪一块GPU是通过命令行设置CUDA_VISIBLE_DEVICES来实现的，而不是写死在代码里面的。
此外，还有一点不方便的地方，如果有很多台机器，则需要把多份这份代码拷贝多次，在每台机器上分别运行上述命令才可以，还是不太方便的。
参数同步更新的分布式训练
同步更新稍微麻烦了点，需要加几行代码（重点参考了《Tensorflow实战》一书），改动部分已经标明，代码如下：
# tensorflow distribute train by synchronously update
import tensorflow as tf
import numpy as np
tf.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")
tf.app.flags.DEFINE_string("job_name", "", "one of ps or worker")
tf.app.flags.DEFINE_integer("task_index", 0, "0, 1, 2...")
FLAGS = tf.app.flags.FLAGS
def main(_):
ps_hosts = FLAGS.ps_hosts.split(",")
worker_hosts = FLAGS.worker_hosts.split(",")
n_works = len(worker_hosts)
# Create a cluster from the parameter server and worker server
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)):
train_X = np.linspace(-1.0, 1.0, 100)
train_Y = 2.0 * train_X + np.random.randn(*train_X.shape) * 0.33 + 10.0
X = tf.placeholder("float")
Y = tf.placeholder("float")
w = tf.Variable(0.0, name="weight")
b = tf.Variable(0.0, name="bias")
loss = tf.square(Y - tf.multiply(X, w) - b)
global_step = tf.Variable(0)
# for Syncmously updata
# 同步更新模式下，需要等待所有计算图计算出梯度，然后梯度求平均，tf.train.SyncReplicasOptimizer实现了这种封装
opt = tf.train.SyncReplicasOptimizer(
tf.train.AdagradOptimizer(0.01),
replicas_to_aggregate=n_works,
total_num_replicas=n_works,
)
train_op = opt.minimize(loss, global_step=global_step)
saver = tf.train.Saver()
summary_op = tf.summary.merge_all()
init_op = tf.global_variables_initializer()
# for Syncmously updata
# 同步模式下，主计算服务器需要协调不同计算服务器计算得到的梯度，并更新参数。
if FLAGS.task_index==0:
# 定义协调不同计算服务器的队列，并定义初始化操作
chief_queue_runner = opt.get_chief_queue_runner()
init_tokens_op = opt.get_init_tokens_op(0)
# Create a "supervisor", which oversees the training process.
sv = tf.train.Supervisor(is_chief=(FLAGS.task_index==0),
logdir="~/Downloads/log/",
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, retoring from a
# checkpoint, and closing when done or an error occurs.
# for Syncmously updata.
# prepare_or_wait_for_session used by sync. It will wait until main node ok and parameter init over!
# for Syncmously updata.
# 这里用的是prepare_or_wait_for_session。
# 相比于异步更新的managed_session：只要某个计算服务器参数初始化完毕就可以开始，
# prepare_or_wait_for_session：等待所有计算服务器参数初始化完毕（参数只有一份，后续的计算服务器应该不需要初始化了？只需要和参数服务器建立一个关系？），主节点协调工作完毕后，开始。
with sv.prepare_or_wait_for_session(server.target) as sess:
# for Syncmously updata
if FLAGS.task_index==0:
# 开始训练之前，主计算服务器需要启动协调同步更新的队列，并执行初始化操作
sv.start_queue_runners(sess, [chief_queue_runner])
sess.run(init_tokens_op)
step = 0
while step < 100000:
# Run a training step asynchronously
for (x, y) in zip(train_X, train_Y):
_, step =sess.run([train_op, global_step], feed_dict={X:x, Y:y})
loss_value = sess.run(loss, feed_dict={X:x, Y:y})
print("Step: {}, loss: {}".format(step, loss_value))
# Ask for all the services to stop
sv.stop()
if __name__=="__main__":
tf.app.run()
