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.ClusterSpec({
    "job1": [
        "job1_task0.example.com:2222",
        "job1_task1.example.com:2222",
        "job1_task2.example.com:2222"
    ],
    "job2": [
        "job2_task0.example.com:2222",
        "job2_task1.example.com:2222"
    ]})
# 在机器 job1_task0.example.com:2222 上启动server
server = tf.train.Server(cluster, job_name="job1", task_index=0)

''''''

# 在机器 job2_task1.example.com:2222 上启动server
server = tf.train.Server(cluster, job_name="job1", task_index=1)

好,现在为止出现了几个概念,先来汇总一下,以免大家绕晕:

  • cluster,TensorFlow 集群定义,包含一个或多个 jobs,每个 job 被切分成一个或多个具有相同功能的 tasks
  • job,一个 job 由一系列相同目的的tasks组成,一个 job 中的 task 通常运行在不同的机器上
  • task,一个t ask 完成一个具体的任务,被 job 中的 task_index 唯一标识,通常对应一台机器 ,与一个具体的TensorFlow Server相关联
  • TensorFlow Server,运行 tf.train.Server 实例的进程,是集群中的一个成员,可以和集群中的其他server 通信,通常包括 master service和 worker service
  • Master service,一个GRPC service,提供到多个分布式设备的远程访问,表现形式是session.target,master service对应master_service.proto,可以看到里面有 createSession, closeSession等接口。所有的TensorFlow server都实现了master service
  • Worker service,一个GRPC service,用指定设备来执行部分计算图操作。worker service对应worker_service.proto,主要作用是在当前设备上执行计算图。所有的TensorFlow server也都实现了worker service
  • client,client概念上面没有提到,但是对于下面的内容至关重要,所以放在这里一并介绍了。一个client通常是一段构造TensorFlow计算图的程序,同时构造了一个 TensorFlow session 来与集群交互。说白了前面创建的cluster与server只是搭建分布式环境,真正要执行计算需要创建client。那么创建session就很简单了,sess=tf.Session() 对吧。But,在分布式环境下可不是这么随意的,对于tf.Session这个类,其第一个参数是target,一般情况下大家确实用不到,因为不指定这个参数的话,Session 就默认调用本地设备,但是在分布式环境就需要指定了,指定的这个值就是server里面的master(server.target提供这个参数)

分布式训练架构:parameter server

上面介绍了,如何创建分布式的环境。下面来介绍下TensorFlow的分布式训练时所使用的架构。我们知道,一般机器学习的训练过程可以分为以下两步:

  1. 根据已有参数计算梯度
  2. 根据梯度来更新参数

对于超大规模的训练,数据量大,模型参数多,不仅计算梯度的过程要放到多个CPU或GPU上执行,参数的存储和更新也需要放到多个CPU或GPU上执行,这就需要多台机器来进行计算能力的扩展。在parameter server (参数服务器)架构中,集群中的节点被分为两类:parameter server 和 worker。
其中 parameter server 负责存放模型的参数和更新梯度的操作,而worker负责计算参数的梯度。在每个迭代过程中,worke 节点从 parameter sever 节点中获得参数,再将计算的梯度回传给 parameter server,parameter server 聚合从所有 worker 传回的梯度,然后开始更新参数。整个架构如下图所示,这里只简单介绍下,想要深入理解参数服务器架构,可以看下沐神的这两篇文章,Parameter Server for Distributed Machine Learning
Scaling Distributed Machine Learning with Parameter Server,这里只抛砖引玉。而且TensorFlow的分布式本身只是采用了 parameter server 的思想,与沐神的PS架构并不完全对等。
在这里插入图片描述
在tensorflow的分布式中,同样也将集群中的节点分为 ps 和 worker 两类:

  1. ps节点就是参数服务器,简称 ps(parameter server), 负责参数的存储和更新,ps 节点可以不止一个,也可以拿多台机器来充当ps节点。
  2. worker节点,从 ps 节点上 pull 下最新参数,加载数据,执行计算图,获得梯度,然后将梯度 push 给 PS 节点。

现在再回想下刚才创建 cluster 的过程,不是按 job 划分的吗?那么现在就可以划分为两类 job,ps job 和 worker job,

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"
    ]})

可以看出,cluster接收的其实就是一个字典,字典里面包含了各个task所在host的主机地址,这个 cluster 共包含两类job:ps和worker,共5个task:

/job:worker/task:0
/job:worker/task:1
/job:worker/task:2
/job:ps/task:0
/job:ps/task:1

再加上刚才的client概念,就有了下面这两张图,来自TensorFlow Architecture, 上面有一系列这样的图,对理解分布式的概念很有帮助。再来梳理下,client负责创建计算图,然后要跑计算时,先把计算图以及要执行的节点(Graph中的Node)发给mastermaster负责资源调度(就是这个计算该怎么执行,在哪些设备执行),最终的执行需要各个worker进程去做(使用本地设备执行计算),所以每个 server 会包含masterworker两个部分。
在这里插入图片描述
在这里插入图片描述
上面简单介绍了下如何创建分布式环境、server和client等一些基本概念、还有深度学习中常用的ps架构。在单机单卡和单机多卡的例子中,还知道了可以指定具体的设备来存储参数,那么如何在分布式的环境中把这些整合起来呢?由于创建分布式集群的过程对于每个节点都是一样的,然后每个 server 要到各自的机器上去启动,所以一般通过命令行的方式传入 ps 的 hosts 和 worker 对应的 hosts,看下面的示例代码 (完整代码在此):

import tensorflow as tf

tf.app.flags.DEFINE_string("ps_hosts", "localhost:2222", "ps hosts")
tf.app.flags.DEFINE_string("worker_hosts", "localhost:2223,localhost:2224", "worker hosts")
tf.app.flags.DEFINE_string("job_name", "worker", "'ps' or'worker'")
tf.app.flags.DEFINE_integer("task_index", 0, "Index of task within the job")

FLAGS = tf.app.flags.FLAGS

def main():
    '''
    这些代码对与所有server都是一样的
    '''
    ps_hosts = FLAGS.ps_hosts.split(",")
    worker_hosts = FLAGS.worker_hosts.split(",")
    # create cluster
    cluster = tf.train.ClusterSpec({"ps": ps_hosts, "worker": worker_hosts})
    # create the server
    server = tf.train.Server(cluster, job_name=FLAGS.job_name,
                             task_index=FLAGS.task_index)

    server.join()

if __name__ == "__main__":
    main()

这里使用 localhost 环境模拟多机环境,假设上述代码在 example.py 文件中,在本地分别执行下面三个命令启动三个 server

nohup python example.py --job_name=ps --task_index=0 > ps0.log &
nohup python example.py --job_name=worker --task_index=0 > worker0.log &
nohup python example.py --job_name=worker --task_index=1 > worker1.log &

现在分布式环境已经成功运行了,集群中有1个PS,2个worker在跑。且三个server都处在等待状态,现在创建一个 client 来执行一个计算图。上面已经说过了,创建 session 的时候,默认是使用本地设备,但是在分布式环境下,可以使用某个 server 对应的 target 来创建 session, 这里使用worker0这个server对应的target,即 “grpc://localhost:2223” 来创建 Session,如下所示,完整代码在example_distributed_client.py

import tensorflow as tf

if __name__ == "__main__":
    # 指定设备:格式,/job:job_name/task:task_index/gpu:gpu_index
    # 参数存储在ps节点上
    with tf.device('/job:ps/task:0/cpu:0'):
        input_data = tf.Variable(
                [[1., 2., 3.], [4., 5., 6.], [7., 8., 9.], [10., 11., 12.]],
                name="input_data")
        b = tf.Variable([[1.], [1.], [2.]], name="w")

    inputs = tf.split(input_data, 2)
    outputs = []
    
    # 使用worker0对应的session创建计算图
    with tf.Session("grpc://localhost:2223") as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        # 2 means 2 workers
        for i in range(2):
            # 把计算任务放到GPU上执行,对应实际中计算梯度的过程
            with tf.device("/job:worker/task:%d/gpu:0" % i):
                print(sess.run(inputs[i]))
                outputs.append(tf.matmul(inputs[i], b))
        # 每个worker计算完,都将结果收集到outputs中,然后在ps节点上汇总处理,这里只concat模拟一下
        with tf.device('/job:ps/task:0/cpu:0'):
            output = tf.concat(outputs, axis=0)
            print(sess.run(output))

执行这段程序的进程就是一个 client,但是在计算时需要依靠 cluster 中的 worker0 和 worker1 的设备来计算,参数存储在PS上,遵循了标准的 ps 架构。

分布式训练策略:模型并行与数据并行

模型并行

所谓模型并行指的是对于同一批样本数据,当神经网络模型很大时,由于显存限制,它是难以在跑在单个GPU上,这个时候就需要模型并行。做法是把模型拆分成不同部分,然后分散到多个设备上进行并行训练,如下图所示:
模型并行

但是深度学习模型一般包含很多层,如果要采用模型并行策略,一般需要将不同的层运行在不同的设备上,但是实际上层与层之间的运行是存在约束的:前向运算时,后面的层需要等待前面层的输出作为输入,而在反向传播时,前面的层又要受限于后面层的计算结果。所以除非模型本身很大,一般不会采用模型并行。但是如果模型本身存在一些可以并行的单元,那么也是可以利用模型并行策略来提升训练速度,比如GoogLeNet的几个Inception模块。

数据并行

数据并行是分布式中比较常用的一种并行策略,因为实际中往往是数据量很大。数据并行就是在不同机器上放置相同的模型,然后每个机器喂给不同的训练样本来计算梯度。上面说的 cifar10 单机多卡例子就是数据并行,其实TensorFlow中的分布式将的也是数据并行。数据并行根据其参数更新的方式又可以分为同步更新和异步更新。同步更新很简单,如下图所示,就是等待所有设备都计算完才开始更新参数,迭代速度取决于速度最慢的那个机器。异步更新就是各个机器自己干自己的,计算完梯度就去更新。
数据并行

复制训练: Replicated training

TensorFlow中的分布式策略是数据并行,并且把这种策略换了个名字叫复制训练 (Replicated training),就是说多个 worker 使用不同的 batch 数据来训练相同的模型,计算出的梯度放到ps上用于更新参数。TensorFlow也增加了一些库函数来简化复制训练的实现。在 TensorFlow 中共有四种不同的方式来实现复制训练:

  1. in-graph replication:只构建一个client,这个 client 构建一个 Graph,Graph中包含一套模型参数,放置在ps上,同时Graph中包含模型计算部分的多个副本,每个副本都放置在一个worker上,这样多个worker可以同时训练复制的模型。上面单机多卡的 cifar10分类模型example_distributed_client 的那段代码,都属于这个类型。这种方式的缺点也很明显,一旦唯一创建 client 的那个 worker 挂了,整个系统就全崩溃了,容错能力差,因此在实际中使用较少。
  2. Between-graph replication:每个worker都创建一个 client,这个 client 一般还与 task 的主程序在同一进程中。各个 client 构建相同的 Graph,但是参数还是放置在ps上。这种方式就比较好,一个 worker 的client挂掉了,系统还可以继续跑。
  3. Asynchronous training:异步方式训练,各个 worker 自己干自己的,不需要与其它worker来协调,前面也已经详细介绍了异步训练,上面两种方式都可以采用异步训练。
  4. Synchronous training:同步训练,各个 worker 要统一步伐,计算出的梯度要先聚合才可以执行一次模型更新,对于 In-graph replication 方法,由于各个 worker 的计算子图属于同一个 Graph,很容易实现同步训练。但是对于 Between-graph replication 方式,各个 worker 都有自己的 client,这就需要系统上的设计了,TensorFlow 提供了tf.train.SyncReplicasOptimizer来实现 Between-graph replication 的同步训练。

下面来详细讲述下这四种方式

in-graph replication

还是以上面example_distributed_client的例子来说明,不过这次改下代码。上面的代码中每个server启动后都会join(),然后一直 block,并没有人告诉它结束。但worker节点做完了理应结束的,所以改下刚才的代码,判断下如果是ps节点,就 block,这样ps节点一直处于监听状态,等待着worker节点传入参数,而worker节点会把一些参数定义在ps节点上。这样改完后,worker干完整个计算图就结束了,而ps节点进程其实还一直存在。

为了直观地感受创建client及计算图执行的过程,这里用 timeline 来记录最终的 output op的执行过程,并用TensorBoard可视化计算图,代码在example_in_graph.py中,执行这个脚本分别传入 ps_hosts, worker_hosts和对应的 task_index

nohup python example_in_graph.py --job_name=ps --task_index=0 > ps0.log &
nohup python example_in_graph.py --job_name=worker --task_index=0 > worker0.log &
nohup python example_in_graph.py --job_name=worker --task_index=1 > worker.log &

过一会woker0, worker1执行结束。查看两个worker的log可以发现,worker0.log中多了一条tensorflow/core/distributed_runtime/master_session.cc:998] Start master session 5eeddba53ed488e2 with config:, 而对应的worke1的log,并没有start master session的输出信息,说明只在worker0的节点上创建了session,用来执行计算,证明这是图内复制模式。

保存的timeline_client.json文件记录了计算图的执行轨迹,在Chrome中输入chrome://tracing/, 然后load 这个json文件,可以看到可视化的轨迹。可以看到有两个Mat操作分别在worker0上和worker1上,concat操作发生在ps上,这与代码一致。
在这里插入图片描述
通过tensorBoard的可视化,也可以看到整个图的计算过程,及每一步的device.
在这里插入图片描述

between-graph replication

图间复制是TensorFlow比较常用的一种方式,着重介绍下。

分布式参数配置

图间复制方式,每个worker都各自包含一个client,它们构建相同的计算图,然后把模型参数推送到PS节点上,多个worker共享ps上的参数。上面我们介绍的都是只有一个ps的情况,那么如果有多个PS节点时,变量存储和更新怎么分配呢?万能的TensorFlow提供了一个神奇的设备函数,可以用不同的设备函数来创建不同的参数放置策略。这个神奇的设备函数就是 tf.train.replica_device_setter 来做,它会循环地分配变量,自动把把Graph中的参数放到PS上,而同时将计算部分放在当前worker节点上。只需要你在实现时,把整个模型构建的代码用这个函数包起来即可。先看一个应用tf.train.replica_device_setter的例子,来自TF官网

# To build a cluster with two ps jobs on hosts ps0 and ps1, and 3 worker
# jobs on hosts worker0, worker1 and worker2.
cluster_spec = {
    "ps": ["ps0:2222", "ps1:2222"],
    "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]}
with tf.device(tf.train.replica_device_setter(cluster=cluster_spec)):
  # Build your graph
  v1 = tf.Variable(...)  # assigned to /job:ps/task:0
  v2 = tf.Variable(...)  # assigned to /job:ps/task:1
  v3 = tf.Variable(...)  # assigned to /job:ps/task:0
# Run compute

先看下这个函数的API:

  • ps_tasks: ps job中的task数量,如果提供了cluster参数,则可以忽略
  • Ps_device: ps job的字符串名字,默认就是”ps"
  • worker_device: worker job的字符串名字,默认就是”worker"
  • cluster: ClusterDef proto或者ClusterSpec
  • ps_ops: 需要放到ps节点上的op列表,默认是STANDARD_PS_OPS,在源码中的定义如下:可见所有的Variable变量都会被默认放到PS节点上
# This is a tuple of PS ops used by tf.estimator.Estimator which should work in
# almost all of cases.
STANDARD_PS_OPS = ("Variable", "VariableV2", "AutoReloadVariable",
                   "MutableHashTable", "MutableHashTableV2",
                   "MutableHashTableOfTensors", "MutableHashTableOfTensorsV2",
                   "MutableDenseHashTable", "MutableDenseHashTableV2",
                   "VarHandleOp", "BoostedTreesEnsembleResourceHandleOp")
  • ps_strategy:既然模型参数都会被放到PS节点上,那如果有多个PS节点,是如何决定哪个参数放到哪个节点上的呢?就是这个参数来控制的,默认是round-robin 策略,就是按出现次序将参数挨个放到各个PS上,由上图中也可以看出一共2个PS节点,v1先出现就被放到ps0上,接着v1放到ps1上,然后v2又被放到ps0上。图下图所示:
    -在这里插入图片描述

    但是这种方式可能不能使ps负载均衡,在2个PS的情况下可能会出现,所有的W都在ps0,所有的偏置b都在ps1上,这显然会给ps0带来更大的负载压力。如果需要更加合理的分配,可以使用tf.contrib.training.GreedyLoadBalancingStrategy策略, 这是一个简单的贪婪策略,它可以根据参数的内存字节大小来放置到内存合适的ps节点上,从而带来更好的负载均衡,如下图所示:
    在这里插入图片描述
    以上讨论的都还是小字节的参数,每个PS都还可以单独处理一个变量。但当遇到超大字节的变量,比如可能 是千万甚至亿级别的特征,该如何处理?TensorFlow提供了一个分割变量的方法,对于这种超大字节的变量,可 以使用一个分隔符,把这个变量分成多个部分,分发到不同的ps节点上去,如下图所示:
    在这里插入图片描述
    现在解决了参数如何放置的问题,按照惯例下一步就是创建Session了,上面图内复制模式下是这样创建session的: with tf.Session(“grpc://localhost:2223”) as sess, 这里传入一个url: grpc://localhost:2223, 后面的地址对应着worker0的 host:ip ,字符串 grpc://localhost:2223 等价于server0.target, server0就是 server = tf.train.Server(cluster, job_name=“worker” task_index=0), 如果打印此时server.target,发现其值就是 grpc://localhost:2223, 那么问题来了?worker0有target, worker1也有target,那么在各自的节点上用各自的server创建session,不就实现单独创建client了吗?实际上,图间复制模式也是这样做的,下面来总结下图间复制模式下创建session的步骤:

    1 定义一个ClusterSpec;

    2 然后创建一个tensorflow Server,它代表集群里的一个特定任务

    3 如果是ps节点就 server.join(), block在这里,等待worker节点的接入

    4 如果是worker节点,就先用tf.train.replica_device_setter 包起来,作用是自动分配参数到ps节点,而把计算部分留给当前worker 。然后用当前的server创建session,with tf.sessioon(server.target) as sess:

MonitoredTrainingSession

好,现在解决了参数如何放置、如何单独创建session的问题,那么问题又来了,由于各个worker都拥有独立的client,但是都拥有相同的计算图,对于图中的一些公共操作,如模型参数初始化、checkpoint文件保存、summary等,如果每个client都独立进行,显示会出现重复操作而且也是对集群资源的浪费。为了解决这个问题,一般会设定一个worker为chief workr,这个chief worker就作为各个worker的管家,协调它们之间的训练,并且完成一些公共操作如模型初始化、模型保存和恢复、summary等。一般来说使用worker0作为chief worker, 因为worker是从0开始计数的,所以总会有worker0存在。

万能的TensorFlow也替我们做好了这个事情,你只需要告诉它哪个worker是chief workr即可。刚才我们创建session都是用tf.Session()来创建,tf还提供了一个更强大的tf.train.MonitoredTrainingSession,在创建时可以指定哪个worker是chief worker。(有的文章里写的是tf.train.Supervisor, 这个API即将被弃用,TF官方建议使用新的tf.train.MonitoredTrainingSession), 然后MonitoredTrainingSession的代码会自动帮你初始化参数,连sess.run(tf.global_variables_initializer())这一显示的初始化操作都给省了,而且传入的chief worker会帮你完成上述的公共操作。

异步训练

下面用一段示例程序来演示图间复制下的异步训练,完整代码在example_between_graph.py, 这里模拟一个线性运算,采用图间复制模式+异步训练+chief_worker, 从运行log中可以看出在worker0和worker1上各自启动了一个master session,这与代码里的逻辑也吻合,因为创建session时用的是server.target, 这样每个worker都会创建一个master session, 继而创建创建client ,构造计算图.
每个worker拥有独立的client, 下图中左边是worker0,右边是worker1,可以看到step是交叉打印的,说明是异步训练,每个worker自己干自己的。这个模拟过程发现会有重复step出现,可能是这个模拟数据太小,worker训练很快,导致同时开始一个step?之后用稍大数据再逐个step测试。。。
在这里插入图片描述

同步训练

如果想用同步训练怎么办呢?tensorflow提供了对同步训练的支持,使用tf.train.SyncReplicasOptimizer可以实现同步训练,官网上有使用教程,清晰明了。现简单介绍下几个初始化参数:

  • opt, 实际的优化器函数,如 tf.train.GradientDescentOptimizer

  • replicas_to_aggregate, 每次梯度更新时要聚合 replicas_to_aggregate 个steps的计算结果,注意这replicas_to_aggregate个结果没有要求必须是来自全部worker的,假如有一个worker计算特别慢,而其他worker已经计算replicas_to_aggregate步,那么也会收集这些梯度执行一次梯度更新。

  • total_num_replicas,所有的worker 副本数,这个与replicas_to_aggregate不同,如果total_num_replicas > replicas_to_aggregate, 那么剩下的total_num_replicas - replicas_to_aggregate个将成为备份复制。如果total_num_replicas < replicas_to_aggregate,则一个worker可能会执行多个mini_batches,然后ps节点才开始收集梯度。一般replicas_to_aggregate = num_workers,在各个worker的计算能力相同的情况下,每个worker计算一步后开始收集梯度进行更新。、

  • use_lockint, 如果True,用自己设置的locks来执行梯度更新操作

    同步训练的代码也在example_between_graph.py, 只需要在启动脚本时设定is_sync=True即可。同步训练结果如下:
    在这里插入图片描述
    左边是worker0,右边是worker1,可以看到刚开始worker1速度比较快,导致开始时并没有实现真正的同步,worker1训练2个step然后就更新一次梯度(2个step是因为tf.train.SyncReplicasOptimizer(opt, replicas_to_aggregate=2,total_num_replicas=2)),直到step10开始,两个worker基本同时打印相同的step,对于同一个step,可以看到weight和biase是一样的,也证明了这是同步模式,但是loss不一样,因为batch数据不同。这里是worker1开始比较快,实验过程中还出现过worker0 刚开始快,后面同步的情况,还有worker0或者worker1直接从step0干到step2000结束, 然后另一个worker一直没开始的情况,所以同步有风险,使用需谨慎

几个需要注意的地方:

  • 两个worker同步训练,如果一个worker启动较慢,那么刚开始就是另一个worker一直在更新,会发现global step有重复出现,重复出现次数就是replicas_to_aggregate的值。然后到某个step开始,两个worker基本同步,同时打印一个global step.
  • 启动worker或者ps必须在其对应的机器上启动, 如启动worker执行task0, 需要在worker_0对应的host上启动,否则会报Could not start gRPC server错误, 从日志中可以看出会自动把当前角色的host当成localhost,但是创建session时是按照集群中的host去找的
  • 如果ps节点有多个且不在同一个机器上,那么checkpoint_dir必须用hdfs 路径,这样各个ps节点才可以共享读写

Distributed_mnist

好了,上面介绍了复制训练的几种模式,其中图间复制模式比较常用,常见的配置是图间复制模式+异步训练。这里用TensorFlow官网提供的分布式训练mnist数据的例子,来介绍一个细节。完整代码再此, 这里做了一些修改,主要是tf.train.Supervisor替换为tf.train.MonitoredTrainingSession, 使用2个ps,2个worker进行训练。

实验时发现的问题:tf.train.StopAtStepHook是设定停止条件,从文档上来看(下图是部分源码截图), 是总共训练10000步之后就停止session,last_step是训练最大global_step >= last_step 时停止session。使用tf.train.StopAtStepHook(num_step=10000)作为停止条件时发现,如果一个worker训练较快,global先到10000然后就停止了,这时候另外一个worker还会一直训练到某个step停止,而停止的这个step是不受控制的。仔细看StopAtStepHook源码会发现,请求停止的关键语句是:global_step >= self._last_step

  def begin(self):
    self._global_step_tensor = training_util.get_global_step()
    if self._global_step_tensor is None:
      raise RuntimeError("Global step should be created to use StopAtStepHook.")

  def after_create_session(self, session, coord):
    if self._last_step is None:
      global_step = session.run(self._global_step_tensor)
      self._last_step = global_step + self._num_steps

  def before_run(self, run_context):  # pylint: disable=unused-argument
    return SessionRunArgs(self._global_step_tensor)

  def after_run(self, run_context, run_values):
    global_step = run_values.results
    if global_step >= self._last_step:
      run_context.request_stop()

也就是内部是已last_step为停止条件的,如果指定last_step=10000, 那么任何一个worker判断出global_step >=10000,就会停止当前session,如果指定的是num_steps=10000, 那么每个worker内部会计算出last_step = global_step + num_steps, 就是这个global_step,每个worker拿到的值是不一样的。2个worker的情况下,较快的worker从global_step=0开始,它计算出的last_step=10000, 可能它训练了1000个step了,然后另一个慢的work开始了,慢的worker此时也会计算last_step=global_step + num_steps=1000+10000=11000,慢的worker就会一直训练直到global_step > 11000然后停止当前session。为了验证这个判断,自己写了个Hook打印下after_create-session和after_run时的step值:

class MyStopAtStepHook(tf.train.StopAtStepHook):

    def after_create_session(self, session, coord):
        if self._last_step is None:
            global_step = session.run(self._global_step_tensor)
            self._last_step = global_step + self._num_steps
            print("now global_step is %d after create session, num_steps: %d, last_step:%d :"
                  % (global_step, self._num_steps, self._last_step))

    def after_run(self, run_context, run_values):
        global_step = run_values.results
        if global_step >= self._last_step:
            print("global_step is %d when stop." % global_step)
            run_context.request_stop()

再次训练,查看运行log,证明和猜测一致, 所以以后尽量使用last-step来控制停止条件,如果是从checkpoint开始训练的,记得调整这个值,加上要训练的num_steps即可。下图中上边是worker0,下边是worker1,可以看到worker0开始训练前,打印global-step=0, num-steps=10000, last_step=10000, 最后在gloabal_step=10000停止,而worker1速度较慢,在global_step=8390时才开始创建session, 这时计算到的last_step=10000+8390=18390, 最后停止训练时的global_step也是19390。
在这里插入图片描述在这里插入图片描述

加载模型预测

TensorFlow提供了分布式训练的支持,训练的结果同样也会保存checkpoint文件等信息,如果要加载分布式训练出的模型进行预测,过程和加载单机模型是一样的。这里也提供了一个示例,这个没什么可说的。

容错性

上面的内容介绍了分布式训练的过程,但是这样就完了吗?万一跑一个很长时间的training任务时,worker突然挂了怎么办,难道之前的training都付之东流了吗?当然不,伟大的谷歌爸爸在设计时肯定会考虑这个问题, 且听我娓娓道来。

定时saver

在长时间的训练任务中,推荐使用saver来不断把checkpoint保存到磁盘上,这样万一某个节点挂了,还可以从最近的checkpoint恢复, 在tf.train.MonitoredTrainingSession中提供了checkpoint_dir参数可以设置保存路径,需要注意的是如果是多个ps,需要写入HDFS,这样多个PS才可以共同访问。
其中save_checkpoint_secs参数设置了默认的保存间隔,默认是600秒,可以根据自己需求来设置。

此外TensorFlow提供了一个默认的saver来做保存的事情, 这个saver封装在CheckpointSaverHook中,如果没有指定saver,会创建一个默认的saver,代码如下:

def _get_saver_or_default():
  """Returns the saver from SAVERS collection, or creates a default one.

  This method is used by other members of the training module, such as
  `Scaffold`, or `CheckpointSaverHook`.

  Returns:
    `Saver`.

  Raises:
    RuntimeError: If the SAVERS collection already has more than one items.
  """
  collection_key = ops.GraphKeys.SAVERS
  savers = ops.get_collection(collection_key)
  if savers:
    if len(savers) > 1:
      raise RuntimeError(
          "More than one item in collection {}. "
          "Please indicate which one to use by passing it to the constructor.".
          format(collection_key))
    return savers[0]
  saver = Saver(sharded=True, allow_empty=True)
  if saver is not None:
    ops.add_to_collection(collection_key, saver)
  return saver

重点是里面的saver = Saver(sharded=True, allow_empty=True)这句话,其中sharded=True这个参数是设置分片数据,也就是说假如有3个PS任务,则会将checkpoint分片为3份来并行写,提高写入速度。看到一些教程中说要主动构造一个saver,并设置sharded=True,其实不用这么做,因为tf.train.MonitoredTrainingSession中默认就是这样设置的。

节点挂了

节点挂了可以分为几种情况:

  1. 最好的情况就是非Chief的worker task出错了,因为这些task实际上是无状态的。那么当遇到了这种错误,当这样的一个worker task恢复的时候,它会重新与它的PS task建立连接,并且重新开始之前崩溃过的进程。
    在这里插入图片描述

  2. 比较差的一种情况就是PS task失败了,那么就有一个麻烦,因为PS task是有状态的,所有的worker task需要依赖他们来发送他们的梯度并且取得新的参数值。所以这种情况下,他们的chief worker task负责监测这种错误,如果发生了这种错误,chief worker task就打断整个训练,并从上一个检查点恢复所有的PS tasks。

在这里插入图片描述

  1. 最糟糕的情况就是chief worker task失败了,因为我们让它负责其他所有的任务,我们就要确保它失败了,能够让集群里所有的任务都回到一个很好的状态。所以我们做的就是打断训练,并在当它恢复了时候从上一个好的检查点恢复。这样的处理方式很简单,但也依赖于机器的健壮性。这里也抛出了一个想法,或许可以使用一个配置比如Hadoop ZooKeeper 或Etcd来选举chief worker task而不是简单地定义为task0。
    在这里插入图片描述

:以上容错性所有截图和参数分配策略部分的截图均来自2017TensorFlow开发者峰会,TensorFlow主管研发Derek Murray做了半小时左右的关于TensorFlow分布式设计的演讲,干货满满,本篇博客也有部分文字是根据这个视频翻译过来的,感兴趣可以看下。这里先放张 Derek Murray的帅照,大家意淫一下?
在这里插入图片描述

后记

感谢大家听我扯皮,坚持看到最后。好了,现在到了最关键的时刻,大家有钱的捧个钱场,没钱的出个脑场(这是什么词,我自己也不知道o(╯□╰)o),只能来清空我的TODO List 了?。

  • 如何读取数据并行
    上面的例子中,每个worker都是从HDFS读取全量数据,然后在自己机器上next_batch去训练。如何实现多个worker共用一份数据?
  • 如何优雅地结束PS
    官网很多例子,都没有讲如何优雅地结束ps节点?ps从一开始就block在那里,worker进程干完就结束了,但是ps一直还在,只能暴力kill掉(据说MXNet更粗暴,PS运行完import mxnet就block了)。 这里贴一个美团云的做法:
采用队列通信的方式同步。
1. 为PS和每个worker创建一个队列
2. PS在server创建后,对队列进行worker数量的dequeue操作
3. 每个worker在完成后,对队列进行一次enqueue操作
  • ps/worker数量的设置
    PS和worker的数量如何设置?这个需要大量的实验,一般的策略是模型参数多,就开PS多一点,数据量很大,就开worker多一点。但是没有绝对的平衡点,当PS太多时variable的分片也多,网络额外开销会大,当PS太少时worker拉参数会发生瓶颈。我自己实验时,在上面distributed_mnist的例子中,2个PS2个worker的配置反倒比1个PS1个worker还慢。

一个人可以跑的很快,一群人可以跑的很远 ~

参考资料

  1. 官网 Distributed TensorFlow
  2. TensorFlow Architecture
  3. Distributed TensorFlow(2017 TensorFlow 开发峰会)
  4. MNIST实例:Distributing TensorFlow
  5. TensorFlow白皮书
  6. TensorFlow分布式博客
发布了20 篇原创文章 · 获赞 11 · 访问量 2万+
展开阅读全文

没有更多推荐了,返回首页

©️2019 CSDN 皮肤主题: 大白 设计师: CSDN官方博客

分享到微信朋友圈

×

扫一扫,手机浏览