本篇博客主要介绍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()