分布式TensorFlow

通过使用多个GPU服务器，减少神经网络的实验时间和训练时间。

作者：Jim Dowling

说明：可以在这里找到示例的完整源代码。

2017年6月8日，分布式深度学习的时代开始了。在那一天，Facebook发表了一篇paper，展示了他们将卷积神经网络（ImageNet上的RESNET-50）的训练时间从两周减少到一小时的方法，该方法使用32个服务器的256个GPU。在软件中，他们引入了一种具有非常大的mini-batch大小的技术来训练卷积神经网络（ConvNets）：使学习率与mini-batch大小成比例。这意味着任何人现在都可以使用TensorFlow将分布式训练扩展到数百个GPU。但这不是分布式TensorFlow的唯一优势：通过在许多GPU上并行运行许多实验，可以大幅缩短训练时间，这减少了为神经网络找到最优超参数所需的时间。

随着计算而扩展的方法是AI的未来。
-Rich Sutton，强化学习之父

在本教程中，我们将探索使用TensorFlow的两种不同的分布式方法：

1. 在许多GPU（和服务器）上运行并行实验来搜索好的超参数
2. 通过许多GPU（和服务器）分配单个网络的训练，减少训练时间

我们将在这篇文章中提供方法（1）和（2）的代码示例，但首先，我们需要阐明我们将要讨论的分布式深度学习的类型。

模型并行与数据并行

一些神经网络模型非常庞大，无法适合单个设备（GPU）的内存，Google的神经机器翻译系统就是这种网络的一个例子。这些模型需要拆分到许多设备上（TensorFlow文档中的workers），并行地在设备上进行训练，例如，网络中的不同层可以在不同的GPU上并行训练。这种训练过程通常被称为“模型并行”（或TensorFlow文档中的“in-graphreplication”）。获得良好性能是一项挑战，文章后面不会进一步说明这种方法。

在“数据并行”（或TensorFlow文档中的“between-graphreplication”）中，每个设备都使用相同的模型，但使用不同的训练样本在每个设备中训练模型。这与模型并行形成了对比，模型并行为每个设备使用相同的数据，但在设备之间切分模型。而数据并行中，每个设备将独立地计算其训练样本的预测值与标记的输出（这些训练样本的正确值）之间的误差，由于每个设备都训练不同的样本，因此它会计算模型的不同更新（“梯度”）。然而，该算法依赖于对每次新迭代使用所有处理的聚合结果，就像该算法在单个处理器上运行一样。因此，每个设备必须发送所有更新到其它所有设备中的所有模型。

在本文中，我们将重点放在数据并行上。图1显示了典型的数据并行，将32个不同的图像分配给运行单个模型的256个GPU中的每一个，一个迭代的mini-batch大小总共为8,092个图像（32 x 256）。

图1.在数据并行中，设备使用不同的训练数据子集进行训练。图片由Jim Dowling提供。

同步与异步分布式训练

随机梯度下降（SGD）是用于寻找最优值的迭代算法，是AI中最受欢迎的训练算法之一，它涉及多轮训练，每轮的结果都更新到模型中，以备下一轮训练，每轮训练可以在多个设备上同步或异步运行。

每次SGD迭代运行一个mini-batch的训练样本（Facebook拥有8,092张图像的大尺寸mini-batch）。在同步训练中，所有设备都使用单个（大尺寸）mini-batch数据的不同部分来训练其本地模型，然后将他们本地计算的梯度（直接或间接）传送给其它所有设备，只有在所有设备成功计算并发送了梯度后，模型才会更新。然后将更新后的模型和下一个mini-batch的拆分一起发送到所有节点。也就是说，设备在mini-batch的非重叠分割的子集上进行训练。

虽然并行有很大的加速训练的潜力，但它自然会引入开销，大型模型和慢速网络会增加训练时间。如果有失速（慢速设备或网络连接），训练可能会失速。我们还希望减少训练模型所需的迭代次数，因为每次迭代都需要将更新的模型广播到所有节点。实际上，这意味着尽可能增加mini-batch的尺寸，以免降低训练模型的准确性。

在他们的论文中，Facebook介绍了针对学习率的线性缩放规则，可以用大尺寸mini-batch进行训练，该规则规定“当mini-batch大小乘以k时，将学习率也同样乘以k”，但条件是在达到目标学习率之前，学习率应该在几个epochs内缓慢增加。

在异步训练中，没有设备等待来自任何其他设备的模型更新。这些设备可以独立运行并与对等设备共享结果，或通过一个或多个称为“参数”服务器的中央服务器进行通信。在对等架构中，每个设备运行一个循环，读取数据，计算梯度，将它们（直接或间接）发送到所有设备，并将模型更新为最新版本。在更中心化的架构中，设备以梯度的形式将其输出发送到参数服务器，这些服务器收集和聚合梯度。在同步训练中，参数服务器计算模型最近的最新版本，并将其发送回设备。在异步训练中，参数服务器将梯度发送到本地计算新模型的设备。在这两种架构中，循环重复直到训练结束。图2说明了异步和同步训练之间的区别。

图2.随机梯度下降（SGD）的异步和同步训练。图片由Jim Dowling提供。

参数服务器架构

当并行SGD使用参数服务器时，算法首先将模型广播给workers（设备）。在每次训练迭代中，每个worker从mini-batch中读取自己的部分，计算其自己的梯度，并将这些梯度发送到一个或多个参数服务器。参数服务器汇总来自worker的所有梯度，并等到所有workers完成之后，才计算下一次迭代的新模型，然后将其广播给所有workers。数据流如图3所示。

图3.同步随机梯度下降的参数服务器架构 图片由Jim Dowling提供。

Ring – allreduce架构

在ring-allreduce架构中，没有中央服务器负责聚合来自workers的梯度。相反，在训练迭代中，每个worker读取它自己的mini-batch部分，计算其梯度，将梯度发送到环上的后继邻居，并从环上的前一个邻居接收梯度。对于具有N个worker的环，所有workers将在每个worker发送和接收N-1个梯度消息之后收到计算更新模型所需的梯度。

Ring-allreduce是带宽最优化的，因为它可以确保每个主机上可用的上传和下载的网络带宽得到充分利用（与参数服务器方式不同）。Ring-allreduce还可以将深层神经网络中较低层的梯度计算与高层梯度的传输重叠，从而进一步缩短训练时间。数据流如图4所示。

图4.同步随机梯度下降的ring-allreduce架构 图片由JimDowling提供。

并行实验

到目前为止，我们已经讲解了分布式训练。但是，许多GPU也可用于并行化超参数优化。也就是说，当我们想要建立合适的学习率或mini-batch时，我们可以使用不同的超参数组合并行运行多个实验。在所有实验完成后，我们可以使用结果来确定是否需要更多实验，或者当前超参数值是否足够好。如果超参数是可接受的，则可以在许多GPU上训练模型时使用它们。

TensorFlow中分布式GPU的两种用途

以下部分说明如何使用TensorFlow进行并行实验和分布式训练。

并行实验

在许多GPU上并行扫描参数很容易，因为我们只需要一个中心点来安排实验。TensorFlow不提供启动和停止TensorFlow服务器的内置支持，因此我们将使用ApacheSpark在PySpark mapper函数中运行每个TensorFlow Python程序。下面，我们定义一个启动函数，它将参数（1）作为Spark会话对象，（2）一个指定要在每个Spark executor上执行的TensorFlow函数map_fun，以及（3）一个包含超参数的字典args_dict。Spark可以通过在Spark executor中执行它们来并行运行许多Tensorflow服务器，Spark executor是执行任务的分布式服务。在这个例子中，每个executor将会使用在args_dict中根据索引executor_num获得的param_val来计算超参数，然后使用这些超参数运行提供的训练函数。

def launch(spark_session, map_fun, args_dict):
   """ Execute a ‘map_fun’ for each hyperparameter combination from the dictionary ‘args_dict’
    Args:
      :spark_session: SparkSession object
      :map_fun: The TensorFlow function to run (wrapped inside a Spark mapper function)
      :args_dict: hyperparameters to insert as arguments for each TensorFlow function
    """
    sc = spark_session.sparkContext
    # Length of the list of the first list of arguments represents the number of Spark tasks
    num_tasks = len(args_dict.values()[0])
    # Create a number of partitions (tasks)
    nodeRDD = sc.parallelize(range(num_tasks), num_tasks)
    # Execute each of the hyperparameter arguments as a task
    nodeRDD.foreachPartition(_do_search(map_fun, args_dict))

def _do_search(map_fun, args_dict):
    def _wrapper_fun(iter):
        for i in iter:
            executor_num = i
        arg_count = map_fun.func_code.co_argcount
        names = map_fun.func_code.co_varnames
        args = []
        arg_index = 0
        while arg_count > 0:
            # Get arguments for hyperparameter combination
            param_name = names[arg_index]
            param_val = args_dict[param_name][executor_num]
            args.append(param_val)
            arg_count -= 1
            arg_index += 1
        map_fun(*args)
    return _wrapper_fun

TensorFlow训练函数mnist，现在可以从spark中调用。请注意，我们只启动一次，但对于每个超参数组合，任务将在不同的执行程序（共四个）上执行：

args_dict = {'learning_rate': [0.001, 0.0001], 'dropout': [0.45, 0.7]}
def mnist(learning_rate, dropout):
   """
     An implementation of FashionMNIST should go here
  """
launch(spark, mnist, args_dict):

分布式训练

我们将简要介绍三种TensorFlow分布式训练框架：原生分布式TensorFlow，TensorFlowOnSpark和Horovod。

分布式TensorFlow

分布式TensorFlow应用程序由包含一个或多个参数服务器和workers的集群组成，由于workers在训练期间计算梯度，因此通常将其放置在GPU上。参数服务器只需要聚合梯度和广播更新，因此它们通常放置在CPU上，而不是GPU上。其中一个worker，主worker协调模型训练、初始化模型、统计完成的训练步骤数、监控会话、保存TensorBoard的日志、保存和恢复模型检查点以从故障中恢复。主worker还管理故障，确保workers或参数服务器出现故障时的容错。如果主worker自己死掉，则需要从最近的检查点重新开始训练。

分布式TensorFlow作为TensorFlow核心的一部分的一个缺点是您必须明确地管理服务器的启动和停止。这意味着要跟踪程序中所有TensorFlow服务器的IP地址和端口，并手动启动和停止这些服务器。通常，这会导致代码中有很多开关语句来确定哪些语句应该在当前服务器上执行。因此，通过使用集群管理器和Spark，我们将使管理它们更轻松。希望你永远不必像这样编写代码，手动定义ClusterSpec：

tf.train.ClusterSpec({"local": ["localhost:2222", "localhost:2223"]})
tf.train.ClusterSpec({
    "worker": [
        "worker0.example.com:2222",
        "worker1.example.com:2222",
        "worker2.example.com:2222"
    ],
    "ps": [
        "ps0.example.com:2222",
        "ps1.example.com:2222"
    ]})
…
if FLAGS.job_name == "ps":
    server.join()
elif FLAGS.job_name == "worker":
    ….

使用主机端点（IP地址和端口号）创建ClusterSpec是很容易出错和不切实际的。相反，您应该使用诸如YARN、Kubernetes或Mesos之类的集群管理器来降低配置和启动TensorFlow应用程序的复杂性。主要选项是云管理解决方案（如GoogleCloud ML或Databrick的Deep LearningPipelines）或通用资源管理器（如Mesos或YARN）。

TensorFlowOnSpark

TensorFlowOnSpark是一个允许从Spark程序启动分布式TensorFlow应用程序的框架，它可以在独立的Spark群集或YARN群集上运行。下面的TensorFlowOnSpark程序使用ImageNet数据集执行Inception的分布式训练。

它引入的新概念是用于启动集群的TFCluster对象，以及执行训练和推理。集群可以以SPARK模式或TENSORFLOW模式启动。SPARK模式使用RDD向TensorFlow workers提供数据，这对于构建从Spark到TensorFlow的集成管道非常有用，但这存在性能瓶颈，因为只有一个Python线程为TensorFlow worker将RDD序列化为feed_dict。TENSORFLOW输入模式通常是首选，因为数据可以使用更高效的多线程输入队列从分布式文件系统（如HDFS）中读取。当一个集群启动时，它启动TensorFlowworkers和参数服务器（可能在不同的主机上）。参数服务器只执行server.join()命令，而workers读取ImageNet数据并执行分布式训练。主workertask_id 为‘0’。

以下程序收集使用Spark启动和管理Spark上的参数服务器和workers所需的信息。

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from pyspark.context import SparkContext
from pyspark.conf import SparkConf
from tensorflowonspark import TFCluster, TFNode
from datetime import datetime

import os
import sys
import tensorflow as tf
import time

def main_fun(argv, ctx):

  # extract node metadata from ctx
  worker_num = ctx.worker_num
  job_name = ctx.job_name
  task_index = ctx.task_index

  assert job_name in ['ps', 'worker'], 'job_name must be ps or worker'

  from inception import inception_distributed_train
  from inception.imagenet_data import ImagenetData
  import tensorflow as tf

  # instantiate FLAGS on workers using argv from driver and add job_name and task_id
  print("argv:", argv)
  sys.argv = argv

  FLAGS = tf.app.flags.FLAGS
  FLAGS.job_name = job_name
  FLAGS.task_id = task_index
  print("FLAGS:", FLAGS.__dict__['__flags'])

  # Get TF cluster and server instances
  cluster_spec, server = TFNode.start_cluster_server(ctx, 4, FLAGS.rdma)

  if FLAGS.job_name == 'ps':
    # `ps` jobs wait for incoming connections from the workers.
    server.join()
  else:
    # `worker` jobs will actually do the work.
    dataset = ImagenetData(subset=FLAGS.subset)
    assert dataset.data_files()
    # Only the chief checks for or creates train_dir.
    if FLAGS.task_id == 0:
      if not tf.gfile.Exists(FLAGS.train_dir):
        tf.gfile.MakeDirs(FLAGS.train_dir)
    inception_distributed_train.train(server.target, dataset, cluster_spec, ctx)


# parse arguments needed by the Spark driver
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--epochs", help="number of epochs", type=int, default=5)
parser.add_argument("--steps", help="number of steps", type=int, default=500000)
parser.add_argument("--input_mode", help="method to ingest data: (spark|tf)", choices=["spark","tf"], default="tf")
parser.add_argument("--tensorboard", help="launch tensorboard process", action="store_true")

(args,rem) = parser.parse_known_args()
input_mode = TFCluster.InputMode.SPARK if args.input_mode == 'spark' else TFCluster.InputMode.TENSORFLOW

print("{0} ===== Start".format(datetime.now().isoformat()))
sc = spark.sparkContext
num_executors = int(sc._conf.get("spark.executor.instances"))
num_ps = int(sc._conf.get("spark.tensorflow.num.ps"))

cluster = TFCluster.run(sc, main_fun, sys.argv, num_executors, num_ps, args.tensorboard, input_mode)
if input_mode == TFCluster.InputMode.SPARK:
    dataRDD = sc.newAPIHadoopFile(args.input_data,
    "org.tensorflow.hadoop.io.TFRecordFileInputFormat",
    keyClass="org.apache.hadoop.io.BytesWritable",
    valueClass="org.apache.hadoop.io.NullWritable")
    cluster.train(dataRDD, args.epochs)
cluster.shutdown()

请注意，ApacheYARN尚不支持GPU作为资源，TensorFlowOnSpark使用YARN节点标签来调度主机上具有GPU的TensorFlow workers。前面的例子也可以在确实支持GPU作为资源的Hops YARN上运行，从而实现CPU和GPU资源的更精细共享。

容错

可以创建一个MonitoredTrainingSession对象，在发生故障时自动从最新检查点中恢复会话的训练状态。

saver = tf.train.Saver(sharded=True)
is_chief = True if FLAGS.task_id == 0 else False
with tf.Session(server.target) as sess :
   # sess.run(init_op) 
   # re-initialze from checkpoint, if there is one.
    saver.restore(sess, ...)
    while True:
        if is_chief and step % 1000 == 0 :
             saver.save(sess, "hdfs://....")

with tf.train.MonitoredTrainingSession(server.target, is_chief) as sess:
    while not sess.should_stop():
         sess.run(train_op)

Spark将重启失败的executor。如果executor不是主worker，它将联系参数服务器，并像以前一样继续下去，因为worker实际上是无状态的。如果参数服务器死掉，则在新参数服务器加入系统后，主worker可以从最后一个检查点恢复。主worker每1000步就保存一份模型副本作为检查点，如果主worker本身出故障，训练失败，必须开始新的训练工作，但它可以从最新的完整检查点恢复训练。

Horovod

TensorFlow有两个ring-allreduce框架：tensorflow.contrib.mpi_collectives（由百度贡献）和Uber的Horovod，它们建立在Nvidia的NCCL 2库上。我们将研究Horovod，因为它在Nvidia GPU上具有更简单的API和良好的性能，如图5所示。Horovod使用pip进行安装，并且需要事先安装Open MPI和NCCL-2库。Horovod比TensorFlow或TensorFlowOnSpark需要对TensorFlow程序的更改更少，它引入了必须初始化的hvd对象，并且必须包装优化器（hvd使用allreduce或allgather来平均梯度）。GPU使用其本地rank绑定到此进程，并且在初始化期间将rank 0的变量广播到所有其他进程。

使用mpirun命令启动Horovod Python程序，它将每台服务器的主机名称以及每台服务器上要使用的GPU数量作为参数。另一种mpirun方法是使用Hops Hadoop平台从Spark应用程序中运行Horovod，该平台使用HopsYARN自动管理GPU的分配给Horovod进程。目前，Horovod不支持容错操作，应该定期检查模型，以便在失败后，训练可以从最新的检查点恢复。

import horovod.tensorflow as hvd ; import tensorflow as tf
def main(_):
   hvd.init()
   loss = ...
   tf.ConfigProto().gpu_options.visible_device_list = str(hvd.local_rank())
   opt = tf.train.AdagradOptimizer(0.01)
   opt = hvd.DistributedOptimizer(opt)
   hooks = [hvd.BroadcastGlobalVariablesHook(0)]
   train_op = opt.minimize(loss)

图5.在ImageNet数据集上使用ResNet-101进行训练，在DeepLearning11服务器上，Horovod / TensorFlow在DeepLearning11服务器上线性扩展至10个GPU （成本：15,000美元）。图片由Jim Dowling提供。

可伸缩的深度学习层次

在看过许多TensorFlow和大尺寸mini-batch随机梯度下降（SGD）的分布式训练架构之后，我们现在可以定义下面的可伸缩的层次结构。金字塔的顶端是当前TensorFlow算法（包括ring-allreduce）的allreduce系列中最可伸缩的方法，最底层是可扩展性最低（因此也是训练网络最慢的方法）。尽管并行实验与分布式训练是互补的，但正如我们已经表明的那样，它们是普通方式并行的（具有较弱的缩放比例），因此在金字塔中的较低位置。

图6.同步SGD的深度学习层次结构。图片由Jim Dowling提供。

结论

很好！您现在知道了分布式TensorFlow能够做什么，以及如何修改您的TensorFlow程序以进行分布式训练或运行并行实验。这些例子的完整源代码可以在这里找到。

这篇文章是O'Reilly和TensorFlow合作的一部分。 请参阅我们的编辑独立声明。