分布式TF入门学习整理(一)

2017年,Facebook在ResNet50上采用了分布式训练,开启了分布式深度学习的的大规模应用时代。通过使用分布在32个服务器上的256个GPU进行训练,大大提升了模型训练的速度和精度。在面对海量数据和复杂模型时,计算资源很容易称为深度学习的瓶颈,分布式学习的应用则很好地解决了这一难题。在推荐广告领域,算法工程师面对的往往是数以亿计的大规模稀疏数据,分布式学习就显得尤为重要。作为工业界的宠儿(雾),tensorflow提供了一整套分布式训练的代码框架,许多公司都会直接采用或进行二次开发。本文简要介绍了分布式训练的基本知识,仅做入门参考。

1. 参数服务器(Parameter Server, PS)架构
在推荐系统领域,PS是目前业界分布式TF训练最常用的结构。
分布式TF中所有设备被总称为一个cluster(集群),每个cluster会包含若干server,每个server会执行一个task,因此server和task是一一对应的,一个cluster可以看作由server或task组成。
此外,TF又将一系列相似的task集合称为一个job,因此cluster也可以看做job的集合。在PS架构中有两种job:ps和worker,其中ps负责存储网络参数和更新梯度,worker负责读入训练数据并计算梯度。
在分布式训练过程中,神经网络的参数(embedding,DNN等)存储在PS上。PS-worker结构中TF的训练流程如下图所示
分布式TF入门学习整理(一)
图中的训练流程分为5步:
(1). worker读取训练数据
(2). 在worker内经过一次前向传播并计算梯度
(3). 将计算好的梯度上传至PS
(4). 在PS内更新网络的权重
(5). 将更新后的参数重新传到worker上.

2. 参数的更新方式:同步更新和异步更新
同步更新是指所有的worker将梯度计算好之后,在PS中进行合并,一次性更新网络的参数。这种方式能够保证模型的参数是同时更新的,保证了训练的稳定性。但在实际应用中会出现worker计算速度不同,导致整个训练被个别worker拖慢的情况。
异步更新则是每个worker将计算好的梯度异步传给PS,每次PS只更新改梯度对应的网络参数。这种方式训练速度更快,但由于参数是分批更新的,会影响到算法的收敛速度。
在实际应用中可以采取折中的方案,即在异步更新的基础上设置一个延迟阈值,当有worker延迟超过该阈值时,就强制所有worker停下等待其他worker迭代完成。

3. 分布式TF代码示例
3.1 使用tf.train.ClusterSpec来创建一个cluster:

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

在TF中,job由job name(string)来标识,而task用task index(int)来标识,则cluster中的每个task 就可以用job name + task index来唯一标识。(这里的的job name即PS和worker,task index即各自的编号,如PS0, PS1,worker0,worker1…)
上述cluster中包含2种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

3.2 使用tf.train.Server创建各个task的server
如创建上述第一个worker的server:

server = tf.train.Server(cluster, job_name="worker", task_index=0)

3.3 使用tf.device调用cluster中的各个server
创建完server后,现在需要开始构建TF的graph。这里使用tf.device指定cluster中的各个server,来完成graph的创建。

with tf.device("/job:ps/task:0"):
  weights_1 = tf.Variable(...)
  biases_1 = tf.Variable(...)
with tf.device("/job:ps/task:1"):
  weights_2 = tf.Variable(...)
  biases_2 = tf.Variable(...)

这里注意需要在tf.Session()中指定target参数。
grpc是Google rpc的缩写,是google开源的rpc框架,分布式tensorflow底层通信的方式。
4. Replicated Training
Replicated training指分布式TF中的数据并行策略,即多个worker使用不同的mini-batch来训练,然后更新ps上存放的参数。在了解Replicated training之前,首先要了解client的概念。
Client
在前述的创建cluster和server步骤中,搭建了TF训练所需要的分布式环境。当真正进行训练时,需要创建一个client程序, 来创建TF的计算图并建立一个session来与cluster中的设备进行交互。
目前主流的replicated Training做法是使用Between-graph replication方式,这种模式下的参数放在ps上,每个worker包含一个client,它们构建相同的Graph。TF中的tf.train.replica_device_setter函数专门用来方便Graph的构建。

#创建cluster
cluster_spec = {
    "ps": ["ps0:2222", "ps1:2222"],
    "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]}
cluster = tf.train.ClusterSpec(cluster_spec)
with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:%d"%FLAGS.task_index, cluster=cluster)):
      #构建计算图
      v1 = tf.variable(...)
     v2 = tf.variable(...)
     v3 = tf.variable(...)

使用tf.train.replica_device_setter可以自动将Graph中的参数放在ps上,并将Graph的计算部分放在当前worker上。
在使用Between-graph replication进行训练时,一般会指定一个chief worker,协调各个worker之间的训练,并完成模型初始化,模型checkpoint保存和恢复等公共操作。这里可以使用tf.train.MonitoredTrainingSession来创建client 的Session,并指定哪个worker为chief worker。

5. 总结
分布式训练可以做到数据并行和训练并行。数据并行指DNN的网络参数分片地存储在不同的ps上,训练并行是指不同的worker独立计算梯度,并更新PS中的网络参数。分布式训练可以大大提高复杂神经网络的训练效率,也是推荐、广告等算法工程师的必备技能。

上一篇:十、Redis 有了主从复制, 为什么还需要搭建集群?以及如何搭建Redis集群详细图解


下一篇:Order, Sort, Cluster, and Distribute By