实战Google深度学习框架：TensorFlow计算加速

作者：才云科技Caicloud，郑泽宇，顾思宇

要将深度学习应用到实际问题中，一个非常大的问题在于训练深度学习模型需要的计算量太大。比如Inception-v3模型在单机上训练到78%的正确率需要将近半年的时间 ，这样的训练速度是完全无法应用到实际生产中的。为了加速训练过程，本章将介绍如何通过TensorFlow利用GPU或/和分布式计算进行模型训练。本文节选自《TensorFlow：实战Google深度学习框架》第十章。

本文将介绍如何在TensorFlow中使用单个GPU进行计算加速，也将介绍生成TensorFlow会话（tf.Session）时的一些常用参数。通过这些参数可以使调试更加方便而且程序的可扩展性更好。然而，在很多情况下，单个GPU的加速效率无法满足训练大型深度学习模型的计算量需求，这时将需要利用更多的计算资源。为了同时利用多个GPU或者多台机器，10.2节中将介绍训练深度学习模型的并行方式。然后，10.3节将介绍如何在一台机器的多个GPU上并行化地训练深度学习模型。在这一节中也将给出具体的TensorFlow样例程序来使用多GPU训练模型，并比较并行化效率提升的比率。最后在10.4节中将介绍分布式TensorFlow，以及如何通过分布式TensorFlow训练深度学习模型。在这一节中将给出具体的TensorFlow样例程序来实现不同的分布式深度学习训练模式。虽然TensorFlow可以支持分布式深度学习模型训练，但是它并不提供集群创建、管理等功能。为了更方便地使用分布式TensorFlow，10.4节中将介绍才云科技基于Kubernetes容器云平台搭建的分布式TensorFlow系统。

1. TensorFlow使用GPU

TensorFlow程序可以通过tf.device函数来指定运行每一个操作的设备，这个设备可以是本地的CPU或者GPU，也可以是某一台远程的服务器。但在本节中只关心本地的设备。TensorFlow会给每一个可用的设备一个名称，tf.device函数可以通过设备的名称来指定执行运算的设备。比如CPU在TensorFlow中的名称为/cpu:0。在默认情况下，即使机器有多个CPU，TensorFlow也不会区分它们，所有的CPU都使用/cpu:0作为名称。而一台机器上不同GPU的名称是不同的，第n个GPU在TensorFlow中的名称为/gpu:n。比如第一个GPU的名称为/gpu:0，第二个GPU名称为/gpu:1，以此类推。

TensorFlow提供了一个快捷的方式来查看运行每一个运算的设备。在生成会话时，可以通过设置log_device_placement参数来打印运行每一个运算的设备。下面的程序展示了如何使用log_device_placement这个参数。

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
c = a + b
# 通过log_device_placement参数来输出运行每一个运算的设备。
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在没有GPU的机器上运行以上代码可以得到以下输出：
Device mapping: no known devices.

add: /job:localhost/replica:0/task:0/cpu:0
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中，TensorFlow程序生成会话时加入了参数log_device_placement=True，所以程序会将运行每一个操作的设备输出到屏幕。于是除了可以看到最后的计算结果之外，还可以看到类似“add: /job:localhost/replica:0/task:0/cpu:0”这样的输出。这些输出显示了执行每一个运算的设备。比如加法操作add是通过CPU来运行的，因为它的设备名称中包含了/cpu:0。

在配置好GPU环境的TensorFlow中 ，如果操作没有明确地指定运行设备，那么TensorFlow会优先选择GPU。比如将以上代码在亚马逊（Amazon Web Services, AWS）的 g2.8xlarge实例上运行时，会得到以下运行结果。

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
[ 2.  4.  6.]

从上面的输出可以看到在配置好GPU环境的TensorFlow中，TensorFlow会自动优先将运算放置在GPU上。不过，尽管g2.8xlarge实例有4个GPU，在默认情况下，TensorFlow只会将运算优先放到/gpu:0上。于是可以看见在上面的程序中，所有的运算都被放在了/gpu:0上。如果需要将某些运算放到不同的GPU或者CPU上，就需要通过tf.device来手工指定。下面的程序给出了一个通过tf.device手工指定运行设备的样例。

import tensorflow as tf

# 通过tf.device将运算指定到特定的设备上。
with tf.device('/cpu:0'):
   a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
   b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
with tf.device('/gpu:1'):
    c = a + b

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在AWS g2.8xlarge实例上运行上述代码可以得到以下结果：
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:1
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中可以看到生成常量a和b的操作被加载到了CPU上，而加法操作被放到了第二个GPU“/gpu:1”上。在TensorFlow中，不是所有的操作都可以被放在GPU上，如果强行将无法放在GPU上的操作指定到GPU上，那么程序将会报错。以下代码给出了一个报错的样例。

import tensorflow as tf

# 在CPU上运行tf.Variable
a_cpu = tf.Variable(0, name="a_cpu")

with tf.device('/gpu:0'):
    # 将tf.Variable强制放在GPU上。
    a_gpu = tf.Variable(0, name="a_gpu")

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
sess.run(tf.initialize_all_variables())

'''
运行上面的程序将会报出以下错误：
tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'a_gpu': Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and devices: 
Identity: CPU 
Assign: CPU 
Variable: CPU 
[[Node: a_gpu = Variable[container="", dtype=DT_INT32, shape=[], shared_ name="", _device="/device:GPU:0"]()]]
'''

不同版本的TensorFlow对GPU的支持不一样，如果程序中全部使用强制指定设备的方式会降低程序的可移植性。在TensorFlow的kernel 中定义了哪些操作可以跑在GPU上。比如可以在variable_ops.cc程序中找到以下定义。

# define REGISTER_GPU_KERNELS(type)                                       \
    REGISTER_KERNEL_BUILDER(                                                 \
      Name("Variable").Device(DEVICE_GPU).TypeConstraint("dtype"),\
      VariableOp);                                                                 \
   …
TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_KERNELS);

在这段定义中可以看到GPU只在部分数据类型上支持tf.Variable操作。如果在TensorFlow代码库中搜索调用这段代码的宏TF_CALL_GPU_NUMBER_TYPES，可以发现在GPU上，tf.Variable操作只支持实数型（float16、float32和double）的参数。而在报错的样例代码中给定的参数是整数型的，所以不支持在GPU上运行。为避免这个问题，TensorFlow在生成会话时可以指定allow_soft_placement参数。当allow_soft_placement参数设置为True时，如果运算无法由GPU执行，那么TensorFlow会自动将它放到CPU上执行。以下代码给出了一个使用allow_soft_placement参数的样例。

import tensorflow as tf

a_cpu = tf.Variable(0, name="a_cpu")
with tf.device('/gpu:0'):
    a_gpu = tf.Variable(0, name="a_gpu")

# 通过allow_soft_placement参数自动将无法放在GPU上的操作放回CPU上。
sess = tf.Session(config=tf.ConfigProto(
    allow_soft_placement=True, log_device_ placement=True))
sess.run(tf.initialize_all_variables())

'''
运行上面这段程序可以得到下面的结果：
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0
a_gpu: /job:localhost/replica:0/task:0/cpu:0
a_gpu/read: /job:localhost/replica:0/task:0/cpu:0
a_gpu/Assign: /job:localhost/replica:0/task:0/cpu:0
init/NoOp_1: /job:localhost/replica:0/task:0/gpu:0
a_cpu: /job:localhost/replica:0/task:0/cpu:0
a_cpu/read: /job:localhost/replica:0/task:0/cpu:0
a_cpu/Assign: /job:localhost/replica:0/task:0/cpu:0
init/NoOp: /job:localhost/replica:0/task:0/gpu:0
init: /job:localhost/replica:0/task:0/gpu:0
a_gpu/initial_value: /job:localhost/replica:0/task:0/gpu:0
a_cpu/initial_value: /job:localhost/replica:0/task:0/cpu:0
从输出的日志中可以看到在生成变量a_gpu时，无法放到GPU上的运算被自动调整到了CPU上（比如a_gpu和a_gpu/read），而可以被GPU执行的命令（比如a_gpu/initial_value）依旧由GPU执行。
'''

虽然GPU可以加速TensorFlow的计算，但一般来说不会把所有的操作全部放在GPU上。一个比较好的实践是将计算密集型的运算放在GPU上，而把其他操作放到CPU上。GPU是机器中相对独立的资源，将计算放入或者转出GPU都需要额外的时间。而且GPU需要将计算时用到的数据从内存复制到GPU设备上，这也需要额外的时间。TensorFlow可以自动完成这些操作而不需要用户特别处理，但为了提高程序运行的速度，用户也需要尽量将相关的运算放在同一个设备上。

2. 深度学习训练并行模式

TensorFlow可以很容易地利用单个GPU加速深度学习模型的训练过程，但要利用更多的GPU或者机器，需要了解如何并行化地训练深度学习模型。常用的并行化深度学习模型训练方式有两种，同步模式和异步模式。本节中将介绍这两种模式的工作方式及其优劣。

为帮助读者理解这两种训练模式，本节首先简单回顾一下如何训练深度学习模型。图10-1展示了深度学习模型的训练流程图。深度学习模型的训练是一个迭代的过程。在每一轮迭代中，前向传播算法会根据当前参数的取值计算出在一小部分训练数据上的预测值，然后反向传播算法再根据损失函数计算参数的梯度并更新参数。在并行化地训练深度学习模型时，不同设备（GPU或CPU）可以在不同训练数据上运行这个迭代的过程，而不同并行模式的区别在于不同的参数更新方式。

图10-2展示了异步模式的训练流程图。从图10-2中可以看到，在每一轮迭代时，不同设备会读取参数最新的取值，但因为不同设备读取参数取值的时间不一样，所以得到的值也有可能不一样。根据当前参数的取值和随机获取的一小部分训练数据，不同设备各自运行反向传播的过程并独立地更新参数。可以简单地认为异步模式就是单机模式复制了多份，每一份使用不同的训练数据进行训练。在异步模式下，不同设备之间是完全独立的。

图10-1 深度学习模型训练流程图

图10-2 异步模式深度学习模型训练流程图

然而使用异步模式训练的深度学习模型有可能无法达到较优的训练结果。图10-3中给出了一个具体的样例来说明异步模式的问题。其中黑色曲线展示了模型的损失函数，黑色小球表示了在t0时刻参数所对应的损失函数的大小。假设两个设备d0和d1在时间t0同时读取了参数的取值，那么设备d0和d1计算出来的梯度都会将小黑球向左移动。假设在时间t1设备d0已经完成了反向传播的计算并更新了参数，修改后的参数处于图10-3中小灰球的位置。然而这时的设备d1并不知道参数已经被更新了，所以在时间t2时，设备d1会继续将小球向左移动，使得小球的位置达到图10-3中小白球的地方。从图10-3中可以看到，当参数被调整到小白球的位置时，将无法达到最优点。

图10-3 异步模式训练深度学习模型存在的问题示意图

图10-4 同步模式深度学习模型训练流程图

为了避免更新不同步的问题，可以使用同步模式。在同步模式下，所有的设备同时读取参数的取值，并且当反向传播算法完成之后同步更新参数的取值。单个设备不会单独对参数进行更新，而会等待所有设备都完成反向传播之后再统一更新参数 。图10-4展示了同步模式的训练过程。从图10-4中可以看到，在每一轮迭代时，不同设备首先统一读取当前参数的取值，并随机获取一小部分数据。然后在不同设备上运行反向传播过程得到在各自训练数据上参数的梯度。注意虽然所有设备使用的参数是一致的，但是因为训练数据不同，所以得到参数的梯度就可能不一样。当所有设备完成反向传播的计算之后，需要计算出不同设备上参数梯度的平均值，最后再根据平均值对参数进行更新。

同步模式解决了异步模式中存在的参数更新问题，然而同步模式的效率却低于异步模式。在同步模式下，每一轮迭代都需要设备统一开始、统一结束。如果设备的运行速度不一致，那么每一轮训练都需要等待最慢的设备结束才能开始更新参数，于是很多时间将被花在等待上。虽然理论上异步模式存在缺陷，但因为训练深度学习模型时使用的随机梯度下降本身就是梯度下降的一个近似解法，而且即使是梯度下降也无法保证达到全局最优值，所以在实际应用中，在相同时间内，使用异步模式训练的模型不一定比同步模式差。所以这两种训练模式在实践中都有非常广泛的应用。

3. 多GPU并行

在10.2节中介绍了常用的分布式深度学习模型训练模式。这一节将给出具体的TensorFlow代码，在一台机器的多个GPU上并行训练深度学习模型。因为一般来说一台机器上的多个GPU性能相似，所以在这种设置下会更多地采用同步模式训练深度学习模型。下面将给出具体的代码，在多GPU上训练深度学习模型解决MNIST问题。本节的样例代码将沿用5.5节中使用的代码框架，并且使用5.5节中给出的mnist_inference.py程序来完成神经网络的前向传播过程。以下代码给出了新的神经网络训练程序mnist_multi_ gpu_train.py。

# -*- coding: utf-8 -*-

from datetime import datetime
import os
import time

import tensorflow as tf
import mnist_inference

# 定义训练神经网络时需要用到的配置。这些配置与5.5节中定义的配置类似。
BATCH_SIZE = 100 
LEARNING_RATE_BASE = 0.001
LEARNING_RATE_DECAY = 0.99
REGULARAZTION_RATE = 0.0001
TRAINING_STEPS = 1000
MOVING_AVERAGE_DECAY = 0.99 
N_GPU = 4

# 定义日志和模型输出的路径。
MODEL_SAVE_PATH = "/path/to/logs_and_models/"
MODEL_NAME = "model.ckpt"

# 定义数据存储的路径。因为需要为不同的GPU提供不同的训练数据，所以通过placerholder
# 的方式就需要手动准备多份数据。为了方便训练数据的获取过程，可以采用第7章中介绍的输
# 入队列的方式从TFRecord中读取数据。于是在这里提供的数据文件路径为将MNIST训练数据
# 转化为TFRecords格式之后的路径。如何将MNIST数据转化为TFRecord格式在第7章中有
# 详细介绍，这里不再赘述。
DATA_PATH = "/path/to/data.tfrecords" 

# 定义输入队列得到训练数据，具体细节可以参考第7章。
def get_input():
    filename_queue = tf.train.string_input_producer([DATA_PATH]) 
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    # 定义数据解析格式。
    features = tf.parse_single_example(
        serialized_example,
        features={
            'image_raw': tf.FixedLenFeature([], tf.string),
            'pixels': tf.FixedLenFeature([], tf.int64),
            'label': tf.FixedLenFeature([], tf.int64),
        })        
    # 解析图片和标签信息。
    decoded_image = tf.decode_raw(features['image_raw'], tf.uint8)
    reshaped_image = tf.reshape(decoded_image, [784])
    retyped_image = tf.cast(reshaped_image, tf.float32)
    label = tf.cast(features['label'], tf.int32)

    # 定义输入队列并返回。
    min_after_dequeue = 10000
    capacity = min_after_dequeue + 3 * BATCH_SIZE
    return tf.train.shuffle_batch(
        [retyped_image, label], 
        batch_size=BATCH_SIZE, 
        capacity=capacity, 
        min_after_dequeue=min_after_dequeue)

# 定义损失函数。对于给定的训练数据、正则化损失计算规则和命名空间，计算在这个命名空间
# 下的总损失。之所以需要给定命名空间是因为不同的GPU上计算得出的正则化损失都会加入名为
# loss的集合，如果不通过命名空间就会将不同GPU上的正则化损失都加进来。
def get_loss(x, y_, regularizer, scope):
    # 沿用5.5节中定义的函数来计算神经网络的前向传播结果。
    y = mnist_inference.inference(x, regularizer)
    # 计算交叉熵损失。
    cross_entropy = tf.reduce_mean(
        tf.nn.sparse_softmax_cross_entropy_with_logits(y, y_))
    # 计算当前GPU上计算得到的正则化损失。
    regularization_loss = tf.add_n(tf.get_collection('losses', scope))
    # 计算最终的总损失。
    loss = cross_entropy + regularization_loss
    return loss

    # 计算每一个变量梯度的平均值。
    def average_gradients(tower_grads):
      average_grads = []
      # 枚举所有的变量和变量在不同GPU上计算得出的梯度。
      for grad_and_vars in zip(*tower_grads):
          # 计算所有GPU上的梯度平均值。
          grads = []
          for g, _ in grad_and_vars:
              expanded_g = tf.expand_dims(g, 0)
              grads.append(expanded_g)
          grad = tf.concat(0, grads)
          grad = tf.reduce_mean(grad, 0)

          v = grad_and_vars[0][1]
          grad_and_var = (grad, v)
          # 将变量和它的平均梯度对应起来。
          average_grads.append(grad_and_var)
    # 返回所有变量的平均梯度，这将被用于变量更新。
    return average_grads

# 主训练过程。
def main(argv=None): 
    # 将简单的运算放在CPU上，只有神经网络的训练过程放在GPU上。
    with tf.Graph().as_default(), tf.device('/cpu:0'):
        # 获取训练batch。
        x, y_ = get_input()
        regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)

        # 定义训练轮数和指数衰减的学习率。
        global_step = tf.get_variable(
            'global_step', [], initializer=tf.constant_initializer(0), 
             trainable=False)
        learning_rate = tf.train.exponential_decay(
            LEARNING_RATE_BASE, global_step, 60000 / BATCH_SIZE, 
            LEARNING_ RATE_DECAY)       

        # 定义优化方法。
        opt = tf.train.GradientDescentOptimizer(learning_rate)

        tower_grads = []
        # 将神经网络的优化过程跑在不同的GPU上。
        for i in range(N_GPU):
            # 将优化过程指定在一个GPU上。
            with tf.device('/gpu:%d' % i):
                with tf.name_scope('GPU_%d' % i) as scope:
                    cur_loss = get_loss(x, y_, regularizer, scope)
                    # 在第一次声明变量之后，将控制变量重用的参数设置为True。这样可以
                    # 让不同的GPU更新同一组参数。注意tf.name_scope函数并不会影响
                    # tf.get_ variable的命名空间。
                    tf.get_variable_scope().reuse_variables()

                    # 使用当前GPU计算所有变量的梯度。
                    grads = opt.compute_gradients(cur_loss)
                    tower_grads.append(grads)

        # 计算变量的平均梯度，并输出到TensorBoard日志中。
        grads = average_gradients(tower_grads)
        for grad, var in grads:
            if grad is not None:
                tf.histogram_summary(
                    'gradients_on_average/%s' % var.op.name, grad)

        # 使用平均梯度更新参数。
        apply_gradient_op = opt.apply_gradients(
            grads, global_step=global_ step)
        for var in tf.trainable_variables():
            tf.histogram_summary(var.op.name, var)

        # 计算变量的滑动平均值。
        variable_averages = tf.train.ExponentialMovingAverage(
            MOVING_AVERAGE_DECAY, global_step)
        variables_averages_op = variable_averages.apply(
            tf.trainable_variables())

        # 每一轮迭代需要更新变量的取值并更新变量的滑动平均值。
        train_op = tf.group(apply_gradient_op, variables_averages_op)

        saver = tf.train.Saver(tf.all_variables())
        summary_op = tf.merge_all_summaries()        
        init = tf.initialize_all_variables()

        # 训练过程。
        with tf.Session(config=tf.ConfigProto(
            allow_soft_placement=True, 
            log_device_placement=True)) as sess:
            # 初始化所有变量并启动队列。
            init.run()
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess=sess, coord=coord)
            summary_writer = tf.train.SummaryWriter(
                MODEL_SAVE_PATH, sess.graph)

            for step in range(TRAINING_STEPS):
                # 执行神经网络训练操作，并记录训练操作的运行时间。
                start_time = time.time()
                _, loss_value = sess.run([train_op, cur_loss])
                duration = time.time() - start_time

                # 每隔一段时间展示当前的训练进度，并统计训练速度。
                if step != 0 and step % 10 == 0:
                    # 计算使用过的训练数据个数。因为在每一次运行训练操作时，每一个GPU 
                    # 都会使用一个batch的训练数据，所以总共用到的训练数据个数为
                    # batch大小×GPU个数。
                    num_examples_per_step = BATCH_SIZE * N_GPU

                    # num_examples_per_step为本次迭代使用到的训练数据个数， 
                    # duration为运行当前训练过程使用的时间，于是平均每秒可以处理的训
                    # 练数据个数为num_examples_per_step / duration。
                    examples_per_sec = num_examples_per_step / duration

                    # duration为运行当前训练过程使用的时间，因为在每一个训练过程中， 
                    # 每一个GPU都会使用一个batch的训练数据，所以在单个batch上的训
                    # 练所需要时间为duration / GPU个数。
                    sec_per_batch = duration / N_GPU

                    # 输出训练信息。
                    format_str = ('step %d, loss = %.2f (%.1f examples/ '
                                    ' sec; %.3f sec/batch)')
                    print(format_str % (step, loss_value, 
                                            examples_per_sec, sec_per_batch))

                    # 通过TensorBoard可视化训练过程。
                    summary = sess.run(summary_op)
                    summary_writer.add_summary(summary, step)

                # 每隔一段时间保存当前的模型。
                if step % 1000 == 0 or (step + 1) == TRAINING_STEPS:
                    checkpoint_path = os.path.join(
                        MODEL_SAVE_PATH, MODEL_ NAME)
                    saver.save(sess, checkpoint_path, global_step=step)

            coord.request_stop()
            coord.join(threads)

if __name__ == '__main__':
tf.app.run()

'''
在AWS的g2.8xlarge实例上运行上面这段程序可以得到类似下面的结果：
step 10, loss = 71.90 (15292.3 examples/sec; 0.007 sec/batch)
step 20, loss = 37.97 (18758.3 examples/sec; 0.005 sec/batch)
step 30, loss = 9.54 (16313.3 examples/sec; 0.006 sec/batch)
step 40, loss = 11.84 (14199.0 examples/sec; 0.007 sec/batch)
...
step 980, loss = 0.66 (15034.7 examples/sec; 0.007 sec/batch)
step 990, loss = 1.56 (16134.1 examples/sec; 0.006 sec/batch)
'''

在AWS的g2.8xlarge实例上运行以上代码可以同时使用4个GPU训练神经网络。图10-5显示了运行样例代码时不同GPU的使用情况。

图10-5 在AWS的g2.8xlarge实例上运行MNIST样例程序时GPU的使用情况

因为在5.5节中定义的神经网络规模比较小，所以在图10-5中显示的GPU使用率不高。如果训练大型的神经网络模型，TensorFlow将会占满所有用到的GPU。

图10-6展示了通过TensorBoard 可视化得到的样例代码TensorFlow计算图，其中节点上的颜色代表了不同的设备，比如黑色代表CPU、白色代表第一个GPU，等等。从图10-5中可以看出，训练神经网络的主要过程被放到了GPU_0、GPU_1、GPU_2和GPU_3这4个模块中，而且每一个模块运行在一个GPU上。对比图10-5中的TensorFlow计算图可视化结果和图10-4中介绍的同步模式分布式深度学习训练流程图可以发现，这两个图的结构是非常接近的。

图10-6 使用了4个GPU的TensorFlow计算图可视化结果

通过调整参数N_GPU，可以实验同步模式下随着GPU个数的增加训练速度的加速比率。图10-7展示了在给出的MNIST样例代码上，训练速度随着GPU数量增加的变化趋势。从图10-7中可以看出，当使用两个GPU时，模型的训练速度是使用一个GPU的1.92倍。也就是说当GPU数量较小时，训练速度基本可以随着GPU的数量线性增长。图10-8展示了当GPU数量增多时，训练速度随着GPU数量增加的加速比。从图10-8中可以看出，当GPU数量增加时，虽然加速比不再是线性，但TensorFlow仍然可以通过增加GPU数量有效地加速深度学习模型的训练过程。

图10-7 训练速度随着GPU数量增加的变化趋势
（此数据是通过MNIST样例代码在AWS的g2.8xlarge实例上测试得到的）

图10-8 训练速度随着GPU数量增加的变化趋势，此数据来自谷歌官方的测试结果

4. 分布式TensorFlow

通过多GPU并行的方式可以达到很好的加速效果。然而一台机器上能够安装的GPU有限，要进一步提升深度学习模型的训练速度，就需要将TensorFlow分布式运行在多台机器上。本节将介绍如何编写以及运行分布式TensorFlow的程序。首先，在10.4.1小节中将介绍分布式TensorFlow的工作原理，并给出最简单的分布式TensorFlow样例程序。在这一小节中也将介绍不同的TensorFlow分布式方式。然后，在10.4.2小节中将给出两个完整的TensorFlow样例程序来同步或者异步地训练深度学习模型。最后，在10.4.3小节中将总结目前原生态分布式TensorFlow中的不足，并介绍才云科技（Caicloud.io）提供的分布式TensorFlow解决方案（TensorFlow as a Service，TaaS）。

分布式TensorFlow原理

在10.2节中介绍了分布式TensorFlow训练深度学习模型的理论。本小节将具体介绍如何使用TensorFlow在分布式集群中训练深度学习模型。以下代码展示了如何创建一个最简单的TensorFlow集群：

import tensorflow as tf
c = tf.constant("Hello, distributed TensorFlow!")
# 创建一个本地TensorFlow集群
server = tf.train.Server.create_local_server()  
# 在集群上创建一个会话。
sess = tf.Session(server.target)      
# 输出Hello, distributed TensorFlow!
print sess.run(c)　　

在以上代码中，首先通过 tf.train.Server.create_local_server函数在本地建立了一个只有一台机器的TensorFlow集群。然后在该集群上生成了一个会话，并通过生成的会话将运算运行在创建的TensorFlow集群上。虽然这只是一个单机集群，但它大致反应了TensorFlow集群的工作流程。TensorFlow集群通过一系列的任务（tasks）来执行TensorFlow计算图中的运算。一般来说，不同任务跑在不同机器上。最主要的例外是使用GPU时，不同任务可以使用同一台机器上的不同GPU。TensorFlow集群中的任务也会被聚合成工作（jobs），每个工作可以包含一个或者多个任务。比如在训练深度学习模型时，一台运行反向传播的机器是一个任务，而所有运行反向传播机器的集合是一种工作。

上面的样例代码是只有一个任务的集群。当一个TensorFlow集群有多个任务时，需要使用tf.train.ClusterSpec来指定运行每一个任务的机器。比如以下代码展示了在本地运行有两个任务的TensorFlow集群。第一个任务的代码如下：

import tensorflow as tf
c = tf.constant("Hello from server1!")

# 生成一个有两个任务的集群，一个任务跑在本地2222端口，另外一个跑在本地2223端口。
cluster = tf.train.ClusterSpec(
    {"local": ["localhost:2222", "localhost: 2223"]})
# 通过上面生成的集群配置生成Server，并通过job_name和task_index指定当前所启动
# 的任务。因为该任务是第一个任务，所以task_index为0。
server = tf.train.Server(cluster, job_name="local", task_index=0)

# 通过server.target生成会话来使用TensorFlow集群中的资源。通过设置
# log_device_placement可以看到执行每一个操作的任务。
sess = tf.Session(
    server.target, config=tf.ConfigProto(log_device_placement=True)) 
print sess.run(c)
server.join()

下面给出了第二个任务的代码：

import tensorflow as tf
c = tf.constant("Hello from server2!")

# 和第一个程序一样的集群配置。集群中的每一个任务需要采用相同的配置。
cluster = tf.train.ClusterSpec(
    {"local": ["localhost:2222", "localhost: 2223"]})
# 指定task_index为1，所以这个程序将在localhost:2223启动服务。
server = tf.train.Server(cluster, job_name="local", task_index=1)
# 剩下的代码都和第一个任务的代码一致。
...

启动第一个任务后，可以得到类似下面的输出：

I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:206] Initialize HostPortsGrpcChannelCache for job local -> {localhost:2222, localhost:2223}
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:202] Started server with target: grpc://localhost:2222
E1123 08:26:06.824503525   12232 tcp_client_posix.c:173]     failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused
E1123 08:26:08.825022513   12232 tcp_client_posix.c:173]     failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused
I tensorflow/core/common_runtime/simple_placer.cc:818] Const: /job:local/ replica:0/task:0/cpu:0
Const: /job:local/replica:0/task:0/cpu:0
Hello from server1!

从第一个任务的输出中可以看到，当只启动第一个任务时，程序会停下来等待第二个任务启动，而且持续输出failed to connect to ‘ipv4:127.0.0.1:2223’: socket error: connection refused。当第二个任务启动后，可以看到从第一个任务中会输出Hello from server1!的结果。第二个任务的输出如下：

I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:206] Initialize HostPortsGrpcChannelCache for job local -> {localhost:2222, localhost:2223} 
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:202] Started server with target: grpc://localhost:2223 
Const: /job:local/replica:0/task:0/cpu:0 
I tensorflow/core/common_runtime/simple_placer.cc:818] Const: /job:local/ replica:0/task:0/cpu:0 
Hello from server2!

值得注意的是第二个任务中定义的计算也被放在了设备/job:local/replica:0/task:0/cpu:0上。也就是说这个计算将由第一个任务来执行。从上面这个样例可以看到，通过tf.train.Server.target生成的会话可以统一管理整个TensorFlow集群中的资源。

和使用多GPU类似，TensorFlow支持通过tf.device来指定操作运行在哪个任务上。比如将第二个任务中定义计算的语句改为以下代码，就可以看到这个计算将被调度到/job:local/replica:0/task:1/cpu:0上面。

with tf.device("/job:local/task:1"):
    c = tf.constant("Hello from server2!")

在上面的样例中只定义了一个工作“local”。但一般在训练深度学习模型时，会定义两个工作。一个工作专门负责存储、获取以及更新变量的取值，这个工作所包含的任务统称为参数服务器（parameter server，ps）。另外一个工作负责运行反向传播算法来获取参数梯度，这个工作所包含的任务统称为计算服务器（worker）。下面给出了一个比较常见的用于训练深度学习模型的TensorFlow集群配置方法。

tf.train.ClusterSpec({
    "worker": [
        "tf-worker0:2222",
        "tf-worker1:2222",
        "tf-worker2:2222"
    ],
    "ps": [
        "tf-ps0:2222",
        "tf-ps1:2222"
    ]}) 

使用分布式TensorFlow训练深度学习模型一般有两种方式。一种方式叫做计算图内分布式（in-graph replication）。使用这种分布式训练方式时，所有的任务都会使用一个TensorFlow计算图中的变量（也就是深度学习模型中的参数），而只是将计算部分发布到不同的计算服务器上。10.3节中给出的使用多GPU样例程序就是这种方式。多GPU样例程序将计算复制了多份，每一份放到一个GPU上进行运算。但不同的GPU使用的参数都是在一个TensorFlow计算图中的。因为参数都是存在同一个计算图中，所以同步更新参数比较容易控制。在10.3节中给出的代码也实现了参数的同步更新。然而因为计算图内分布式需要有一个中心节点来生成这个计算图并分配计算任务，所以当数据量太大时，这个中心节点容易造成性能瓶颈。

另外一种分布式TensorFlow训练深度学习模型的方式叫计算图之间分布式（between-graph replication）。使用这种分布式方式时，在每一个计算服务器上都会创建一个独立的TensorFlow计算图，但不同计算图中的相同参数需要以一种固定的方式放到同一个参数服务器上。TensorFlow提供了tf.train.replica_device_setter函数来帮助完成这一个过程，在10.4.2小节中将给出具体的样例。因为每个计算服务器的TensorFlow计算图是独立的，所以这种方式的并行度要更高。但在计算图之间分布式下进行参数的同步更新比较困难。为了解决这个问题，TensorFlow提供了tf.train.SyncReplicasOptimizer函数来帮助实现参数的同步更新。这让计算图之间分布式方式被更加广泛地使用。在10.4.2小节中将给出使用计算图之间分布式的样例程序来实现异步模式和同步模式的并行化深度学习模型训练过程。

分布式TensorFlow模型训练

本小节中将给出两个样例程序分别实现使用计算图之间分布式（Between-graph replication）完成分布式深度学习模型训练的异步更新和同步更新。第一部分将给出使用计算图之间分布式实现异步更新的TensorFlow程序。这一部分也会给出具体的命令行将该程序分布式的运行在一个参数服务器和两个计算服务器上，并通过TensorBoard可视化在第一个计算服务器上的TensorFlow计算图。第二部分将给出计算图之间分布式实现同步参数更新的TensorFlow程序。同步参数更新的代码大部分和异步更新相似，所以在这一部分中将重点介绍它们之间的不同之处。

• 异步模式样例程序

下面的样例代码将仍然采用5.5节中给出的模式，并复用5.5节mnist_inference.py程序中定义的前向传播算法。以下代码实现了异步模式的分布式神经网络训练过程。

# -*- coding: utf-8 -*-

import time
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

import mnist_inference

# 和5.5节中类似的配置神经网络的设置。
BATCH_SIZE = 100 
LEARNING_RATE_BASE = 0.01
LEARNING_RATE_DECAY = 0.99
REGULARAZTION_RATE = 0.0001
TRAINING_STEPS = 10000
# 模型保存的路径。
MODEL_SAVE_PATH = "/path/to/model"
# MNIST数据路径。
DATA_PATH = "/path/to/data"

# 通过flags指定运行的参数。在10.4.1小节中对于不同的任务（task）给出了不同的程序，
# 但这不是一种可扩展的方式。在这一小节中将使用运行程序时给出的参数来配置在不同
# 任务中运行的程序。
FLAGS = tf.app.flags.FLAGS

# 指定当前运行的是参数服务器还是计算服务器。参数服务器只负责TensorFlow中变量的维护
# 和管理，计算服务器负责每一轮迭代时运行反向传播过程。
tf.app.flags.DEFINE_string('job_name', 'worker', ' "ps" or "worker" ')

# 指定集群中的参数服务器地址。
tf.app.flags.DEFINE_string(
    'ps_hosts', ' tf-ps0:2222,tf-ps1:1111',
    'Comma-separated list of hostname:port for the parameter server jobs. '
    'e.g. "tf-ps0:2222,tf-ps1:1111" ')

# 指定集群中的计算服务器地址。
tf.app.flags.DEFINE_string(
    'worker_hosts', ' tf-worker0:2222,tf-worker1:1111',
    'Comma-separated list of hostname:port for the worker jobs. '
    'e.g. "tf-worker0:2222,tf-worker1:1111" ')

# 指定当前程序的任务ID。TensorFlow会自动根据参数服务器/计算服务器列表中的端口号
# 来启动服务。注意参数服务器和计算服务器的编号都是从0开始的。
tf.app.flags.DEFINE_integer(
    'task_id', 0, 'Task ID of the worker/replica running the training.')

# 定义TensorFlow的计算图，并返回每一轮迭代时需要运行的操作。这个过程和5.5节中的主
# 函数基本一致，但为了使处理分布式计算的部分更加突出，本小节将此过程整理为一个函数。
def build_model(x, y_, is_chief):
   regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)
   # 通过和5.5节给出的mnist_inference.py代码计算神经网络前向传播的结果。
   y = mnist_inference.inference(x, regularizer)
   global_step = tf.Variable(0, trainable=False)

   # 计算损失函数并定义反向传播过程。
   cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        y, tf.argmax(y_, 1))
   cross_entropy_mean = tf.reduce_mean(cross_entropy)
   loss = cross_entropy_mean + tf.add_n(tf.get_collection('losses'))
   learning_rate = tf.train.exponential_decay(
       LEARNING_RATE_BASE, global_step, 60000 / BATCH_SIZE, 
        LEARNING_RATE_ DECAY)

   # 定义每一轮迭代需要运行的操作。
   train_op = tf.train.GradientDescentOptimizer(learning_rate)\
                   .minimize(loss, global_step=global_step)   
   return global_step, loss, train_op

# 训练分布式深度学习模型的主过程。
def main(argv=None): 
   # 解析flags并通过tf.train.ClusterSpec配置TensorFlow集群。
   ps_hosts = FLAGS.ps_hosts.split(',')
   worker_hosts = FLAGS.worker_hosts.split(',')
   cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
   # 通过ClusterSpec以及当前任务创建Server。
   server = tf.train.Server(
       cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_id)

   # 参数服务器只需要管理TensorFlow中的变量，不需要执行训练的过程。server.join()
   # 会一直停在这条语句上。
   if FLAGS.job_name == 'ps':
       server.join()   

   # 定义计算服务器需要运行的操作。在所有的计算服务器中有一个是主计算服务器，它除了负责
   # 计算反向传播的结果，它还负责输出日志和保存模型。
   is_chief = (FLAGS.task_id == 0)
   mnist = input_data.read_data_sets(DATA_PATH, one_hot=True)

   # 通过tf.train.replica_device_setter函数来指定执行每一个运算的设备。
   # tf.train.replica_device_setter函数会自动将所有的参数分配到参数服务器上，而
   # 计算分配到当前的计算服务器上。图10-9展示了通过TensorBoard可视化得到的第一个计
   # 算服务器上运算分配的结果。
   with tf.device(tf.train.replica_device_setter(
           worker_device="/job:worker/task:%d" % FLAGS.task_id, 
            cluster=cluster)):
       x = tf.placeholder(
            tf.float32, [None, mnist_inference.INPUT_NODE], 
            name='x-input')
       y_ = tf.placeholder(
            tf.float32, [None, mnist_inference.OUTPUT_NODE],
            name='y-input')
        # 定义训练模型需要运行的操作。
       global_step, loss, train_op = build_model(x, y_)

        # 定义用于保存模型的saver。
       saver = tf.train.Saver()
        # 定义日志输出操作。
       summary_op = tf.merge_all_summaries()
        # 定义变量初始化操作。
       init_op = tf.initialize_all_variables()
        # 通过tf.train.Supervisor管理训练深度学习模型的通用功能。
        # tf.train. Supervisor能统一管理队列操作、模型保存、日志输出以及会话的生成。
       sv = tf.train.Supervisor(
           is_chief=is_chief,           # 定义当前计算服务器是否为主计算服务器，只有
                                            # 主计算服务器会保存模型以及输出日志。
         logdir=MODEL_SAVE_PATH,    # 指定保存模型和输出日志的地址。
         init_op=init_op,               # 指定初始化操作。
         summary_op=summary_op,     # 指定日志生成操作。
         saver = saver,                 # 指定用于保存模型的saver。
         global_step=global_step,   # 指定当前迭代的轮数，这个会用于生成保存模
                                            # 型文件的文件名。
         save_model_secs=60,            # 指定保存模型的时间间隔。
         save_summaries_secs=60)    # 指定日志输出的时间间隔。

       sess_config = tf.ConfigProto(allow_soft_placement=True, 
                                           log_device_placement=False)
        # 通过tf.train.Supervisor生成会话。
       sess = sv.prepare_or_wait_for_session(
             server.target, config=sess_config)   

       step = 0
       start_time = time.time()
        # 执行迭代过程。在迭代过程中tf.train.Supervisor会帮助输出日志并保存模型， 
        # 所以不需要直接调用这些过程。
       while not sv.should_stop():
           xs, ys = mnist.train.next_batch(BATCH_SIZE)
           _, loss_value, global_step_value = sess.run(
                 [train_op, loss, global_step], feed_dict={x: xs, y_: ys})
           if global_step_value >= TRAINING_STEPS: break

             # 每隔一段时间输出训练信息。
           if step > 0 and step % 100 == 0:
               duration = time.time() - start_time
                  # 不同的计算服务器都会更新全局的训练轮数，所以这里使用
                  # global_step_value可以直接得到在训练中使用过的batch的总数。
               sec_per_batch = duration / global_step_value
               format_str = ("After %d training steps (%d global steps), "
                                  "loss on training batch is %g.  " 
                                  "(%.3f sec/batch)")
               print(format_str % (step, global_step_value, 
                                          loss_value, sec_per_batch))
           step += 1
   sv.stop()

if __name__ == "__main__":
   tf.app.run()

假设上面代码的文件名为dist_tf_mnist_async.py，那么要启动一个拥有一个参数服务器、两个计算服务器的集群，需要先在运行参数服务器的机器上启动以下命令：

python dist_tf_mnist_async.py \
--job_name='ps' \
--task_id=0 \
--ps_hosts='tf-ps0:2222' \
--worker_hosts='tf-worker0:2222,tf-worker1:2222'

然后在运行第一个计算服务器的机器上启动以下命令：

python dist_tf_mnist_async.py \
--job_name='worker' \
--task_id=0 \
--ps_hosts='tf-ps0:2222' \
--worker_hosts='tf-worker0:2222,tf-worker1:2222'

最后在运行第二个计算服务器的机器上启动以下命令：

python dist_tf_mnist_async.py \
--job_name='worker' \
--task_id=1 \
--ps_hosts='tf-ps0:2222' \
--worker_hosts='tf-worker0:2222,tf-worker1:2222'

在启动第一个计算服务器之后，这个计算服务器就会尝试连接其他的服务器（包括计算服务器和参数服务器）。如果其他服务器还没有启动，则被启动的计算服务器会报连接出错的问题。下面展示了一个出错信息。

E1201 01:26:04.166203632   21402 tcp_client_posix.c:173]     failed to connect to 'ipv4:tf-worker1:2222': socket error: connection refused

不过这不会影响TensorFlow集群的启动。当TensorFlow集群中所有服务器都被启动之后，每一个计算服务器将不再报错。在TensorFlow集群完全启动之后，训练过程将被执行。图10-9展示了第一个计算服务器的TensorFlow计算图。从图10-9中可以看出，神经网络中定义的参数被放在了参数服务器上（图中浅灰色节点），而反向传播的计算过程则放在了当前的计算服务器上（图中的深灰色节点）。

图10-9 通过TensorBoard可视化的分布式TensorFlow计算图

在计算服务器训练神经网络的过程中，第一个计算服务器会输出类似下面的信息。

After 100 training steps (100 global steps), loss on training batch is 0.302718.  (0.039 sec/batch)
After 200 training steps (200 global steps), loss on training batch is 0.269476.  (0.037 sec/batch)
After 300 training steps (300 global steps), loss on training batch is 0.286755.  (0.037 sec/batch)
After 400 training steps (463 global steps), loss on training batch is 0.349983.  (0.033 sec/batch)
After 500 training steps (666 global steps), loss on training batch is 0.229955.  (0.029 sec/batch)
After 600 training steps (873 global steps), loss on training batch is 0.245588.  (0.027 sec/batch)

第二个计算服务器会输出类似下面的信息。

After 100 training steps (537 global steps), loss on training batch is 0.223165.  (0.007 sec/batch)
After 200 training steps (732 global steps), loss on training batch is 0.186126.  (0.010 sec/batch)
After 300 training steps (925 global steps), loss on training batch is 0.228191.  (0.012 sec/batch)

从输出的信息中可以看到，在第二个计算服务器启动之前，第一个计算服务器已经运行了很多轮迭代了。在异步模式下，即使有计算服务器没有正常工作，参数更新的过程仍可继续，而且全局的迭代轮数是所有计算服务器迭代轮数的和。

• 同步模式样例程序

和异步模式类似，下面给出的代码同样也是基于5.5小节中给出的框架。该代码实现了同步模式的分布式神经网络训练过程。

# -*- coding: utf-8 -*-

import time
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

import mnist_inference

BATCH_SIZE = 100 
LEARNING_RATE_BASE = 0.8
LEARNING_RATE_DECAY = 0.99
REGULARAZTION_RATE = 0.0001
TRAINING_STEPS = 10000
MODEL_SAVE_PATH = "/path/to/model"
DATA_PATH = "/path/to/data"

# 和异步模式类似的设置flags。
FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string('job_name', 'worker', ' "ps" or "worker" ')
tf.app.flags.DEFINE_string(
    'ps_hosts', ' tf-ps0:2222,tf-ps1:1111',
   'Comma-separated list of hostname:port for the parameter server jobs. '
    'e.g. "tf-ps0:2222,tf-ps1:1111" ')

tf.app.flags.DEFINE_string(
    'worker_hosts', ' tf-worker0:2222,tf-worker1:1111',
    'Comma-separated list of hostname:port for the worker jobs. '
    'e.g. "tf-worker0:2222,tf-worker1:1111" ')

tf.app.flags.DEFINE_integer(
    'task_id', 0, 'Task ID of the worker/replica running the training.')

# 和异步模式类似地定义TensorFlow的计算图。唯一的区别在于使用
# tf.train.SyncReplicasOptimizer函数处理同步更新。
def build_model(x, y_, n_workers, is_chief):
   regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)
   y = mnist_inference.inference(x, regularizer)
   global_step = tf.Variable(0, trainable=False)

   variable_averages = tf.train.ExponentialMovingAverage(
        MOVING_AVERAGE_ DECAY, global_step)
   variables_averages_op = variable_averages.apply(tf.trainable_variables())

   cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        y, tf.argmax(y_, 1))
   cross_entropy_mean = tf.reduce_mean(cross_entropy)
   loss = cross_entropy_mean + tf.add_n(tf.get_collection('losses'))
   learning_rate = tf.train.exponential_decay(
       LEARNING_RATE_BASE, global_step, 60000 / BATCH_SIZE, 
        LEARNING_RATE_ DECAY)

   # 通过tf.train.SyncReplicasOptimizer函数实现同步更新。
   opt = tf.train.SyncReplicasOptimizer(
           # 定义基础的优化方法。
         tf.train.GradientDescentOptimizer(learning_rate),        
           # 定义每一轮更新需要多少个计算服务器得出的梯度。
         replicas_to_aggregate=n_workers,                                 
           # 指定总共有多少个计算服务器。
            total_num_replicas=n_workers,                                     
           # 指定当前计算服务器的编号。
         replica_id=FLAGS.task_id)                                           

   train_op = opt.minimize(loss, global_step=global_step)     
   return global_step, loss, train_op, opt

def main(argv=None): 
   # 和异步模式类似地创建TensorFlow集群。
   ps_hosts = FLAGS.ps_hosts.split(',')
   worker_hosts = FLAGS.worker_hosts.split(',')
   n_workers = len(worker_hosts)

   cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
   server = tf.train.Server(
       cluster, job_name = FLAGS.job_name, task_index=FLAGS.task_id)

   if FLAGS.job_name == 'ps':
       server.join()

   is_chief = (FLAGS.task_id == 0)
   mnist = input_data.read_data_sets(DATA_PATH, one_hot=True)

   with tf.device(tf.train.replica_device_setter(
           worker_device="/job:worker/task:%d" % FLAGS.task_id,
             cluster=cluster)):
       x = tf.placeholder(
            tf.float32, [None, mnist_inference.INPUT_NODE], 
            name='x-input')
       y_ = tf.placeholder(
            tf.float32, [None, mnist_inference.OUTPUT_NODE], 
            name='y-input')
       global_step, loss, train_op, opt = build_model(
            x, y_, n_workers, is_chief)

       saver = tf.train.Saver()
       summary_op = tf.merge_all_summaries()
       init_op = tf.initialize_all_variables()

        # 在同步模式下，主计算服务器需要协调不同计算服务器计算得到的参数梯度并最终更新
        # 参数。这需要主计算服务器完成一些额外的初始化工作。
       if is_chief:
             # 定义协调不同计算服务器的队列并定义初始化操作。
           chief_queue_runner = opt.get_chief_queue_runner()
           init_tokens_op = opt.get_init_tokens_op(0)

        # 和异步模式类似的声明tf.train.Supervisor。
       sv = tf.train.Supervisor(is_chief=is_chief,
                                logdir=MODEL_SAVE_PATH,
                                init_op=init_op,
                                summary_op=summary_op,
                                saver=saver,
                                global_step=global_step,
                                save_model_secs=60,
                                save_summaries_secs=60)
       sess_config = tf.ConfigProto(allow_soft_placement=True, 
                                           log_device_placement=False)
       sess = sv.prepare_or_wait_for_session(
            server.target, config=sess_config)        

        # 在开始训练模型之前，主计算服务器需要启动协调同步更新的队列并执行初始化操作。
       if is_chief:
           sv.start_queue_runners(sess, [chief_queue_runner])
           sess.run(init_tokens_op)

        # 和异步模式类似的运行迭代的训练过程。
       step = 0
       start_time = time.time()
       while not sv.should_stop():
           xs, ys = mnist.train.next_batch(BATCH_SIZE)
           _, loss_value, global_step_value = sess.run(
                 [train_op, loss, global_step], feed_dict={x: xs, y_: ys})
           if global_step_value >= TRAINING_STEPS: break

           if step > 0 and step % 100 == 0:
               duration = time.time() - start_time
               sec_per_batch = duration / (global_step_value * n_workers)
               format_str = ("After %d training steps (%d global steps), "
                                  "loss on training batch is %g.  " 
                                  "(%.3f sec/batch)")
               print(format_str % (step, global_step_value, 
                                          loss_value, sec_per_batch))
           step += 1
   sv.stop()

if __name__ == "__main__":
   tf.app.run()

和异步模式类似，在不同机器上运行以上代码就可以启动TensorFlow集群。但和异步模式不同的是，当第一台计算服务器初始化完毕之后，它并不能直接更新参数。这是因为在程序中要求每一次参数更新都需要来自两个计算服务器的梯度。在第一个计算服务器上，可以看到与下面类似的输出。

E1201 01:26:04.166203632   21402 tcp_client_posix.c:173]     failed to connect to 'ipv4:10.57.60.76:2222': socket error: connection refused
After 100 training steps (100 global steps), loss on training batch is 1.88782.  (0.176 sec/batch)
After 200 training steps (200 global steps), loss on training batch is 0.834916.  (0.101 sec/batch)
…
After 800 training steps (800 global steps), loss on training batch is 0.524181.  (0.045 sec/batch)
After 900 training steps (900 global steps), loss on training batch is 0.384861.  (0.042 sec/batch)

第二个计算服务器的输出如下：

After 100 training steps (100 global steps), loss on training batch is 1.88782.  (0.028 sec/batch)
After 200 training steps (200 global steps), loss on training batch is 0.834916.  (0.027 sec/batch)
…
After 800 training steps (800 global steps), loss on training batch is 0.474765.  (0.026 sec/batch)
After 900 training steps (900 global steps), loss on training batch is 0.420769.  (0.026 sec/batch)

在第一个计算服务器的第一行输出中可以看到，前100轮迭代的平均速度为0.176 sec/batch，要远远慢于最后的平均速度0.042 sec/batch。这是因为在第一迭代轮开始之前，第一个计算服务器需要等待第二个计算服务器执行初始化的过程，于是导致前100轮迭代的平均速度是最慢的。这也反应了同步更新的一个问题。当一个计算服务器被卡住时，其他所有的计算服务器都需要等待这个最慢的计算服务器。

为了解决这个问题，可以调整tf.train.SyncReplicasOptimizer函数中的replicas_to_ aggregate参数。当replicas_to_aggregate小于计算服务器总数时，每一轮迭代就不需要收集所有的梯度，从而避免被最慢的计算服务器卡住。TensorFlow也支持通过调整同步队列初始化操作tf.train.SyncReplicasOptimizer.get_init_tokens_op中的参数来控制对不同计算服务器之间的同步要求。当提供给初始化函数get_init_tokens_op的参数大于0时，TensorFlow支持多次使用由同一个计算服务器得到的梯度，于是也可以缓解计算服务器性能瓶颈的问题。

使用Caicloud运行分布式TensorFlow

从10.4.2小节中给出的样例程序可以看出，每次运行分布式TensorFlow都需要登录不同的机器来启动集群。这使得使用起来非常不方便。当需要使用100台机器运行分布式TensorFlow时，需要手动登录到每一台机器并启动TensorFlow服务，这个过程十分烦琐。而且，当某个服务器上的程序死掉之后，TensorFlow并不能自动重启，这给监控工作带来了巨大的难度。如果类比TensorFlow与Hadoop ，可以发现TensorFlow只实现了相当于Hadoop中MapReduce的计算框架，而没有提供类似Yarn的集群管理工具以及HDFS的存储系统。为了降低分布式TensorFlow的使用门槛，才云科技(Caicloud.io)基于Kubernetes 容器云平台提供了一个分布式TensorFlow平台TensorFlow as a Service（TaaS） 。本节中将大致介绍如何使用Caicloud提供的TaaS平台运行分布式TensorFlow。

从10.4.2小节中给出的代码可以看出，编写分布式TensorFlow程序需要指定很多与模型训练无关的代码来完成TensorFlow集群的设置工作。为了降低分布式TensorFlow的学习成本，Caicloud的TensorFlow as a Service（TaaS）平台首先对TensorFlow集群进行了更高层的封装，屏蔽了其中与模型训练无关的底层细节。其次，TaaS平台结合了谷歌开源的容器云平台管理工具Kubernetes来实现对分布式TensorFlow任务的管理和监控，并支持通过UI设置分布式 TensorFlow任务的节点个数、是否使用GPU等信息。

Caicloud的TaaS平台提供了一个抽象基类CaicloudDistTensorflowBase，该类封装了分布式TensorFlow集群的配置与启动、模型参数共享与更新逻辑处理、计算节点之间的协同交互以及训练得到的模型和日志的保存等与模型训练过程无关的操作。用户只需要继承该基类，并实现与模型训练相关的函数即可。其代码结构如下：

import caicloud_dist_tensorflow_base as caicloud
class MyDistTfModel(caicloud.CaicloudDistTensorflowBase):
    """基于自身业务来定义训练模型、执行训练操作等"""

用户继承的类需要选择性地实现CaicloudDistTensorflowBase基类中的4个函数，它们分别是build_model、get_init_fn、train和after_train。build_model给出了定义TensorFlow计算图的接口。在这个函数中，用户需要处理输入数据、定义深度学习模型以及定义训练模型的过程。get_init_fn函数中可以定义在会话（tf.Session）生成之后需要额外完成的初始化工作，当没有特殊的初始化操作时，用户可以不用定义这个函数。这个函数可以完成模型预加载的过程。train函数中定义的是每一轮训练中需要运行的操作。TaaS会自动完成迭代的过程，但用户需要定义每一轮迭代中需要执行的操作。最后在训练结束之后，用户可以通过定义after_train来评测以及保存最后得到的模型。

图10-10 Caicloud提供的TaaS分布式TensorFlow任务详情页面

5. 小结

在本章中介绍了TensorFlow如何通过GPU或/和分布式集群的方式加速深度学习模型的训练过程。首先，10.1节介绍了在TensorFlow中使用单个GPU加速计算的过程。TensorFlow对于单个GPU的支持是非常方便的，几乎不需要任何的额外设置。TensorFlow可以自动将计算优先分配到GPU上。这一节也介绍了如何使用tf.device函数来手动配置计算运行的设备。然后，在10.2节中详细介绍了训练深度学习模型的并行模式。这一节中介绍了同步模式和异步模型两种并行模式，并介绍了这两种模式各自的优缺点。接着，在10.3节中给出了通过TensorFlow实现了在一台机器的多个GPU上并行地训练深度学习模型。

最后，在10.4节中介绍了如何通过TensorFlow集群进一步加大训练深度学习模型的并行化程度。10.4.1小节介绍了TensorFlow集群的运行机制，并给出了启动简单TensorFlow集群的样例程序。这个小节中也介绍了TensorFlow集群的计算图内分布式方式和计算图之间分布式方式，并指出在海量数据下，使用计算图之间分布式方式的可扩展性更强。10.4.2小节给出了具体的TensorFlow代码，该代码通过计算图之间分布式方式实现了并行化深度学习模型训练的同步模式和异步模式。10.4.3小节中指出了原生的TensorFlow在支持分布式中的不足，并介绍了如何使用才云科技（caicloud.io）提供的TensorFlow as a Service平台来更加高效的运行分布式TensorFlow模型训练过程。

点击订购：《TensorFlow：实战Google深度学习框架》

---