TensorFlow分布式采坑记

前言

为什么需要分布式训练？简而言之，就是数据量很大、模型参数很大，大到在一台机器上无法执行前向和后向计算，这时候就需要分布式系统来训练更大更复杂的模型。这种模型往往是深度学习模型（废话），而稳坐深度学习框架头把交椅的 TensorFlow 自然会想到这一点，TensorFlow官方从0.8版本开始支持模型的分布式训练，后续版本也不断在优化，现在的TensorFlow支持 单机多卡、多机多卡、同步、异步等训练方式。在这篇文章里，将简单介绍下TensorFlow分布式的基础知识，然后用几个实例（采坑）来演示下在TensorFlow中如何分布式地训练模型，不足之处，还望各位 TF Boys 们轻捶。

单机单卡

在 TensorFlow 中，变量是可以复用的，变量通过变量名唯一确定，变量和计算图都可以和设备绑定。如果一个图计算时需要用到变量a，而变量a不在该设备上，则会自动生成相应的通信代码，将变量a加载到该设备上。
因此变量存放到哪个设备上对于程序的正确性没有影响，但会导致通信开销有所差异。下图展示的代码是一个单机单卡的例子，将参数 w, b 放在 CPU 上，然后将计算过程放在 GPU 上。

def one_gpu():
    '''
    单机单卡
    对于单机单卡，可以把参数和定义都放在GPU上，不过如果模型较大参数较多，全放在GPU上显存不够时
    可以把参数定义放在CPU上，如下所示
    '''
    with tf.device("/cpu:0"):
        w = tf.Variable(tf.constant([[1.0, 2.0], [4.0, 5.0]]), name="w")
        b = tf.Variable(tf.constant([[1.0], [2.0]]), name="b")

    with tf.device("/gpu:0"):
        addwb = tf.add(w, b)
        mulwb = tf.matmul(w, b)

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        val1, val2 = sess.run([addwb, mulwb])
        print(val1)
        print(val2)

单机多卡

我现在的机器上只有单卡GPU，没办法演示多卡，不过可以看看 tensorflow 官网提供的 cifar10 单机多卡的例子，重点在170行左右。这里我也提供了一个单机多卡的示例，代码也是参考的 cifar10。

这里只讲下cifar10是如何实现单机多卡的，大体思路是每块GPU分别从 data_queue 里去 dequeue 一个 batch 的数据，然后在当前GPU上执行前向计算的过程，并收集这次的所有梯度。当循环到下一个GPU时，该次训练所使用的参数和上个GPU是一样的，代码中体现在tf.get_variable_scope().reuse_variables()这一句。也就是说一次梯度更新前，所有GPU之间是共享模型参数的。然后等所有GPU都收集完梯度后，CPU会统一做一次梯度的平均(grads = average_gradients(tower_grads))，然后用平均梯度去更新参数(opt.apply_gradients(grads,global_step)。需要注意的是这个收集梯度的过程是同步的，必须等所有 GPU 结束后CPU才开始平均梯度的操作，很明显整个模型的训练速度取决于最慢的那块 GPU 卡。整个过程，可以用下图一图以蔽之。
解释下从哪看出 cpu 负责收集梯度然后更新的，从代码开始出可以看到，整个 train 的计算图默认是在 cpu 上的(with tf.Graph().as_default(), tf.device('/cpu:0'))，只有前向计算 loss 和计算梯度的过程是被带有 gpu 的 with 块包起来的，这与上图的描述也是一样的，感兴趣的同学可以把代码拉下来，自己跑一遍。土豪可以用多块 GPU 感受下飞一般的速度，中产阶级可以用多 CPU 代替多 GPU，我等屌丝只能望其项背。。

with tf.Graph().as_default(), tf.device('/cpu:0'):
    # Create a variable to count the number of train() calls. This equals the
    # number of batches processed * FLAGS.num_gpus.
    global_step = tf.get_variable(
        'global_step', [],
        initializer=tf.constant_initializer(0), trainable=False)

    # Calculate the learning rate schedule.
    num_batches_per_epoch = (cifar10.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN /
                             FLAGS.batch_size / FLAGS.num_gpus)
    decay_steps = int(num_batches_per_epoch * cifar10.NUM_EPOCHS_PER_DECAY)

    # Decay the learning rate exponentially based on the number of steps.
    lr = tf.train.exponential_decay(cifar10.INITIAL_LEARNING_RATE,
                                    global_step,
                                    decay_steps,
                                    cifar10.LEARNING_RATE_DECAY_FACTOR,
                                    staircase=True)

    # Create an optimizer that performs gradient descent.
    opt = tf.train.GradientDescentOptimizer(lr)

    # Get images and labels for CIFAR-10.
    images, labels = cifar10.distorted_inputs()
    batch_queue = tf.contrib.slim.prefetch_queue.prefetch_queue(
          [images, labels], capacity=2 * FLAGS.num_gpus)
    # Calculate the gradients for each model tower.
    tower_grads = []
    with tf.variable_scope(tf.get_variable_scope()):
        ......

从上面单机单卡和单机多卡的例子中，可以知道模型参数或者计算图是可以拆开放到不同的设备上的，不同的设备通过变量名可以共享参数。

分布式训练

重头戏来了，当单机多卡也无法满足训练的速度需求的话，那就需要上分布式，搞多台机器一起训练了，也可以说是多机多卡吧。谈到分布式，必不可少的一个概念就是集群，先来看下TensorFlow官方对集群的定义：

A TensorFlow “cluster” is a set of “tasks” that participate in the distributed execution of a TensorFlow graph. Each task is associated with a TensorFlow “server”, which contains a “master” that can be used to create sessions, and a “worker” that executes operations in the graph. A cluster can also be divided into one or more “jobs”, where each job contains one or more tasks.

翻译一下就是，Tensorflow 集群是一系列分布式执行计算图的tasks, 每一个 task 与一个 server 相对应，一个server 包含 master service 和 worker service。Master service 负责创建 session，worker 负责执行图中的计算操作。一个集群也可以被切分成多个 jobs，每个 job 包含一系列相同功能的 tasks。

为了创建集群，需要在每个task上开启 TensorFlow 的 server，每个 task 通常对应一台机器。然后在每个task上都需要做如下的事情：

1. 创建一个 tf.train.ClusterSpec ，来描述集群中的所有tasks。这部分对于所有tasks都是一样的
2. 在每个task上创建一个 tf.train.Server ，在构造时候传入tf.train.ClusterSpec ，并且用 job_name 和 task_index 来标识当前task ，这样每个server在启动时就加入了集群的信息，可以和集群中的其他 server通信。需要注意的是，sever的创建需要在自己所在 host 上，一旦所有的 server 在各自的 host 上创建好了，整个集群就搭建好了。

用代码描述一下就是：

# 创建一个tf集群，包含2个job，job1有3个tasks, job2有2个tasks
# cluster接收的其实就是一个字典，字典里面包含了各个task所在host的主机地址
cluster = tf.train.