Distributed TensorFlow

最新推荐文章于 2022-01-28 20:28:58 发布

Omni-Space

最新推荐文章于 2022-01-28 20:28:58 发布

阅读量727

点赞数

TensorFlow 同时被 2 个专栏收录

71 篇文章 0 订阅

订阅专栏

Distributed Deep Learning

3 篇文章 0 订阅

订阅专栏

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

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

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

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

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

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

模型并行性与数据并行性

一些神经网络模型非常庞大，无法适应单个设备（GPU）的内存。 Google的神经机器翻译系统就是这种网络的一个例子。这些模型需要在许多设备上分开（TensorFlow文档中的工作人员），并行地在设备上进行培训。例如，网络中的不同层可以在不同的GPU上并行训练。这种训练过程通常被称为“模型并行”（或TensorFlow文档中的“图中复制”）。获得良好性能是一项挑战，我们不会进一步涵盖这种方法。

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

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

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

同步与异步分布式培训

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

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

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

在他们的论文中，Facebook引入了针对学习率的线性缩放规则，可以用大批量小批量进行培训。该规则规定“当小批量大小乘以k时，将学习速率乘以k”，但条件是在达到目标学习速率之前，学习速率应该在几个时期内缓慢增加。

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

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

参数服务器架构

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

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

环 - allreduce体系结构

在ring-allreduce体系结构中，没有集中来自工作者的梯度的中央服务器。相反，在训练迭代中，每个工作人员读取它自己的最小批次拆分，计算其梯度，将梯度发送到环上的后继邻居，并从环上的前一个邻居接收梯度。对于具有N个工人的环，所有工人将在每个工人发送和接收N-1个梯度消息之后收到计算更新模型所需的梯度。

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

图4.同步随机梯度下降的ring-allreduce体系结构图片由Jim Dowling提供。

平行实验

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

TensorFlow中分布式GPU的两种用途

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

平行实验

在许多GPU上并行扫描参数很容易，因为我们只需要一个中心点来安排实验。 TensorFlow不提供启动和停止TensorFlow服务器的内置支持，因此我们将使用Apache Spark在PySpark映射器函数中运行每个TensorFlow Python程序。在下面，我们定义了一个启动函数，该函数需要参数（1）Spark会话对象，（2）一个map_fun命名将在每个Spark执行器上执行的TensorFlow函数，以及（3）包含超参数的args_dict字典。 Spark可以通过在Spark执行程序中运行它们来并行运行许多Tensorflow服务器。 Spark执行程序是执行任务的分布式服务。在这个例子中，每个执行程序都会使用它的executor_num来计算它应该从args_dict使用的超参数， args_dict将其索引到正确的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 ()[ 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_val args_dict ][ executor_num ]
             args .  append ( param_val )
             arg_count -= 1
             arg_index += 1
         map_fun ( * args )
     return _wrapper_fun

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

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

分布式培训

我们将简要介绍三种TensorFlow分布式培训框架：本地分布式TensorFlow，TensorFlowOnSpark和Horovod。

分布式TensorFlow

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

分布式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应用程序的复杂性。主要选项是云管理解决方案（如Google Cloud ML或Databrick的Deep Learning Pipelines）或通用资源管理器（如Mesos或YARN）。

TensorFlowOnSpark

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

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

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

 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 import 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

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

   from inception import inception_distributed_train
   from inception.imagenet_data import ImagenetData
   import tensorflow import 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 '__flags' [ '__flags' ])

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

   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 = ImagenetData )
     assert dataset .  data_files ()
    # Only the chief checks for or creates train_dir.
     if FLAGS .  task_id == 0 :
       if not tf .  gfile .  Exists ( train_dir ):
         tf .  gfile .  MakeDirs ( 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' TFCluster .  InputMode . TENSORFLOW

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

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

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

容错

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

 saver = tf .  train .  Saver ( sharded = True )
 is_chief = if FLAGS is_chief True .  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 ( is_chief , is_chief ) as sess :
     while not sess .  should_stop ():
          sess .  run ( train_op )

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

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使用其本地等级绑定到此进程，并且在初始化期间将等级0的变量广播到所有其他进程。

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

 import horovod.tensorflow as hvd ;  import tensorflow import as tf
 def main ( _ ):
    hvd .  init ()
    loss = ...
    tf .  ConfigProto () .  gpu_options .  visible_device_list = str ( 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和大型小批量随机梯度下降（SGD）的分布式训练架构之后，我们现在可以定义下面的比例层次结构。金字塔的顶端是当前TensorFlow算法（包括ring-allreduce）的allreduce系列中最可伸缩的方法，最底层是可扩展性最低（因此也是训练网络最慢的方法）。尽管平行实验与分布式训练是互补的，但正如我们已经表明的那样，它们是平凡并行的（具有较弱的缩放比例），因此在金字塔中被发现较低。