【其他】Tensorflow分布式使用简介

一、 单机多GPU训练
  深度学习由于存在计算量大，并且需要大量的数据来训练的问题，因而需要采用一些并行机制来加快训练速度，目前常用的并行方法主要有数据并行（data parallel）和模型并行（model parallel）两种。下面主要介绍tensorflow框架采用的数据并行方法 。

1.1 数据并行原理
  数据并行的原理如下图所示，假设有两块显卡（GPU1和GPU2），我们经常使用CPU来负责梯度平均和参数更新，另外两块显卡主要负责训练模型副本（model replica）。

主要步骤：
  ① 首先是构图，由于tensorflow主要采用静态图，所以在构图时可以使用for循环分别在GPU1、GPU2上分别定义模型参数变量，网络结构（一致的）；
  ② 分发两个batch的数据， 每个GPU处理一个batch的数据计算，进行前向传播计算出loss，然后计算出关于当前Variables的gradients；
  ③ 把所有GPU算的梯度汇总到CPU上，CPU对梯度进行平均计算，然后对模型的参数进行更新；
  ④ 重复步骤①-③，直到模型收敛。

1.2 主要代码示例

# 对应①、②步，进行构图，分别在每个GPU上定义模型参数变量，网络结构；
x_input = tf.placeholder(tf.int32, shape=[num_gpus,batch_size,..],name="input")
y_input = tf.placeholder(tf.int64, [num_gpus, batch_size,…],name='targets')

optimizer = tf.train.GradientDescentOptimizer(current_learning_rate)
global_step = tf.Variable(0, name="global_step") tower_grads = []
with tf.variable_scope(tf.get_variable_scope(),initializer=initializer):
for i in xrange(num_gpus):
    with tf.device('/gpu:%d' % i):
       with tf.name_scope('%s_%d' % (TOWER_NAME, i)) as scope:
         target_list = inference_graph(x_input[i])
            logits= y_input[i]
            # share variables
            tf.get_variable_scope().reuse_variables()
            tvars = tf.trainable_variables()
            loss = loss_graph(logits, target_list,scope)
            grads,global_norm=tf.clip_by_global_norm(tf.gradients(
                                 loss, tvars), FLAGS.max_grad_norm)
            tower_grads.append(grads)

grads = average_gradients(tower_grads)
train_op=optimizer.apply_gradients(zip(grads,tvars),global_step=global_step)

注意事项：
  ① 在第一步中定义模型参数时，要考虑到不同模型副本之间要能够共享变量，因此要采用tf.get_variable()函数而不是直接tf.Variables()，tf.get_variable拥有一个变量检查机制，会检测已经存在的变量是否设置为共享变量，如果已经存在的变量没有设置为共享变量，TensorFlow 运行到第二个拥有相同名字的变量的时候，就会报错。
  ② 在tensorflow是基于已经定义好的Graph进行模型迭代式训练的，因此在每次迭代过程中，只会对当前的模型参数进行更新，而不会调用tf.get_variable()函数重新定义模型变量，因此变量共享只是存在于模型定义阶段的一个概念。

# 此处对应步骤中的第三步，所有GPU输出的梯度数据传到CPU上，对梯度取平均操作
def average_gradients(tower_grads):
     average_grads = []
       for i,all_grads in enumerate(zip(*tower_grads)):
          if i ==0:
             average_grads.append(all_grads[0])
             continue
          grads = []
          for g in all_grads:
             expanded_g = tf.expand_dims(g, 0)
            grads.append(expanded_g)
         grad = tf.concat(axis=0, values=grads)
         grad = tf.reduce_mean(grad, 0)      
         average_grads.append(grad)
    return average_grads

注意事项：
  ①在编写计算梯度平均函数average_gradients()时，不同的计算梯度方法计算出来的梯度格式有所区别，这里的示例使用tf.gradients()，并使用tf.clip_by_global_norm()对梯度进行修剪，应注意得到的梯度格式然后进行相应的梯度平均计算。

二、分布式多机多GPU训练
  随着模型越来越复杂，模型参数越来越多，当参数的量级上升到百亿量级甚至更大之后，参数的更新的性能都是问题。假设单机有四块显卡，那么一个step也只能处理四个batch的数据，这对于训练一些超大训练集来说是远远不能满足的，这个时候就需要多机多卡进行训练，即分布式的深度学习训练方法。

2.1 参数服务器
  多机多卡的思想跟单机多卡是类似的，即在GPU上训练，然后汇总到CPU上更新，只不过当模型越来越大，模型的参数越来越多，一台机器的性能都不能满足的时候，那么就可以把参数分开放到不同的机器去存储和更新，于是就有了参数服务器的概念，参数服务器可以是多台机器组成的集群。

2.2 tensorflow的分布式
  tensorflow的分布式有in-graph和between-gragh两种架构模式。
  ①in-graph 模式。在in-graph模式， 把计算从单机多GPU扩展到了多机多GPU了，不过数据分发还是在一个节点。这样的好处是配置简单，其他多机多GPU的计算节点，只要起个join操作，暴露一个网络接口，等在那里接受任务就好了。 这些计算节点暴露出来的网络接口，使用起来就跟本机的一个GPU的使用一样， 只要在操作的时候指定tf.device(“/job:worker/task:n”)，就可以向指定GPU一样，把操作指定到一个计算节点上计算，使用起来和多GPU的类似。但是这样的坏处是训练数据的分发依然在一个节点上，要把训练数据分发到不同的机器上， 严重影响并发训练速度。
  ②between-graph模式。在between-graph模式下，训练的参数保存在参数服务器， 数据不用分发， 数据分片的保存在各个计算节点，各个计算节点自己算自己的， 算完了之后， 把要更新的参数告诉参数服务器，参数服务器更新参数。这种模式的优点是不用训练数据的分发了，在训练大数据量的时候可以节省了大量的时间，所以大数据深度学习还是推荐使用between-graph模式。
  between-graph模式使用可见图2示意图，我们只需要在训练过程中，根据机器的数量启动相应数量的worker，其中每个worker可以根据其task_index读入其负责的数据进行训练。
  上述两种模式都支持同步和异步更新。同步更新即每一次所有分发出去的数据计算完成后，汇总所有返回的梯度求平均，再更新参数，那么处理的速度就取决于最慢的那个分片计算的时间。异步更新是所有的计算节点，各自算自己的，更新参数也是自己更新自己计算的结果，这样就有一个明显的问题就个某个节点可能已经算完更新了参数，可能另一个节点还是用的原来的参数计算计算，即存在很难控制参数的问题。综合来看还是推荐使用同步更新，这样网络更新比较稳定，容易控制。

2.3 主要代码示例

import argparse
import sys

import tensorflow as tf

FLAGS = None

def main(_):
  ps_hosts = FLAGS.ps_hosts.split(",")
  worker_hosts = FLAGS.worker_hosts.split(",")

  # 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.contrib.framework.get_or_create_global_step()

      train_op = tf.train.AdagradOptimizer(0.01).minimize(
                                          loss, global_step=global_step)

    # The StopAtStepHook handles stopping after running given steps.
    hooks=[tf.train.StopAtStepHook(last_step=1000000)]

    # The MonitoredTrainingSession takes care of session initialization,
    # restoring from a checkpoint, saving to a checkpoint, and closing when done
    # or an error occurs.
    with tf.train.MonitoredTrainingSession(master=server.target,
                                          is_chief=(FLAGS.task_index == 0),                                                                 
                                          hooks=hooks) as mon_sess:
      while not mon_sess.should_stop():
        # Run a training step asynchronously.
        # See `tf.train.SyncReplicasOptimizer` for additional details on how to
        # perform *synchronous* training.
        # mon_sess.run handles AbortedError in case of preempted PS.
        mon_sess.run(train_op)

  其中tf.train.ClusterSpec是根据参数服务器和工作主机创建一个集群，其中需要把所有的参数服务器和工作主机的ip和端口的信息都包含进去。
  tf.train.Server创建并启动本地任务的服务器，后面根据执行的命令的参数不同，决定了这个任务是哪个任务。
  可以看到，如果参数是ps，那么就是负责参数更新的工作，如果worker那么就负责计算的工作，tf.train. replica_device_setter为负责在相应worker上构建图的副本，默认情况下，只有变量操作被放置在ps任务上。
  tf.train.MonitoredTrainingSession负责初始化session，可以从checkpoint恢复模型，以及保存模型到checkpoint，这个工作是在所有的计算节点里指定一个主节点完成的，即由参数is_chief指定。需要注意的是这里的global_step的值，是所有计算节点共享的，在执行optimizer的minimize的时候，会自动加1，所以可以通过这个可以知道所有的计算节点一共计算了多少步了。
  上述代码就是一个大体的tensorflow多机框架，那么如何部署多机多卡呢？其实很简单，就是把第一部分的单机多卡构图代码加到build mode部分就可以完成多机多卡的部署了。

启动命令：

PS：CUDA_VISIBLE_DEVICES=' '  python train_distribute.py --ps_hosts=172.28.92.31:2220 --worker_hosts=172.28.92.31:2221, 172.28.92.32:2222 --job_name=ps --task_index=0 
Worker1：python train:distribute.py --ps_hosts=172.28.92.31:2220 --worker_hosts=172.28.92.31:2221, 172.28.92.32:2222 --job_name=ps --task_index=0
Worker1：python train_distribute.py --ps_hosts=172.28.92.31:2220 --worker_hosts=172.28.92.31:2221, 172.28.92.32:2222 --job_name=ps --task_index=1

  可以看到在参数服务器的启动中，使用CUDA_VISIBLE_DEVICES=’ ‘，这可以使其可见的显卡数为空，因参数服务器主要使用CPU，之后可以再在这台机器上启用一个worker，即一台机器可以同时作为一个ps和一个worker。

参考文献：
①https://www.tensorflow.org/deploy/distributed
②http://blog.csdn.net/luodongri/article/details/52596780
③http://blog.csdn.net/diligent_321/article/details/53187166