分布式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的训练流程如下图所示
分布式训练流程

图中的训练流程分为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(...)
3.4 创建Session来执行计算图
with tf.Session("grpc://worker7.example.com:2222") as sess:
  for _ in range(10000):
    sess.run(train_op)

这里注意需要在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中的网络参数。分布式训练可以大大提高复杂神经网络的训练效率,也是推荐、广告等算法工程师的必备技能。

参考资料

  1. https://zhuanlan.zhihu.com/p/35083779
  2. https://www.tensorflow.org/deploy/distributed
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 213,992评论 6 493
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 91,212评论 3 388
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 159,535评论 0 349
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 57,197评论 1 287
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 66,310评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 50,383评论 1 292
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,409评论 3 412
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,191评论 0 269
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,621评论 1 306
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 36,910评论 2 328
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,084评论 1 342
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,763评论 4 337
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,403评论 3 322
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,083评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,318评论 1 267
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 46,946评论 2 365
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 43,967评论 2 351