分布式Tensorflow原理初探

因工作缘故最近正在对深度学习框架Tensorflow的分布式原理发起猛攻，官网、白皮书和各博客所述稍有偏差，易对新人造成困惑，经过多天推敲，现将心得总结如下，以记录想法为目的，也望大神不吝赐教。

分布式TF采用PS架构，集群中的节点分为parameter server节点和worker节点。ps和worker代表的是节点所做的task（详见下文），每个循环中ps的任务是分发参数给worker，worker完成计算密集型任务后将梯度回传给ps，ps集中处理各梯度并以此更新参数。

分布式训练策略有模型并行与数据并行之分，TF中一般采用数据并行，即在各worker节点用相同的流图计算不同的数据。具体到分布式TF的训练策略又有图内复制（In-graph）和图间复制（Between-graph）两种复制策略，每种策略又分为同步与异步两种方式。

Repliction

In-graph模式只存在一个Client进程（Client详见下文介绍），这个Client只创建一个Session，这个Session只构建一个图，并将计算密集型的部分复制到各个worker节点。而Between-graph模式中每个承担worker任务的节点中都存在一个Client进程（通常与worker在同一进程中），每个Client都构建一张图，包含所有参数部分与计算部分。这是二者的主要区别。

同步与异步比较好理解，同步是指每次参数更新需要等待所有worker上传梯度再一起更新，异步是指有worker上传就更新参数。同步方式比较稳定但异步方式效率较高且有一定的跳出局部最小的作用。因分布式集群主要想解决的也是训练时长问题，故目前TF中常采用的是异步Between-graph.

目前一些工程师分享的博客中常用下图来介绍分布式TF与单机TF的区别，此处极易造成迷惑。此图源于TF白皮书2015，当时并无In-graph与Between-graph之分，故并未说明右图实为In-graph模式，因为Between-graph模式中各Worker Service中都应有独立的Master Service和Client.

In graph and between graph

TF集群由Server、Client组成，各组成部分说明如下：

Server是服务器进程，一个运行tf.train.Server的实例，是分布式TF集群的重要组成部分，TF Server采用主从模式，一个Server可输出一个Master Service和一个Worker Service。

Server

Master Service是一个RPC（远程过程调用）服务，作用有三个：一是实现集群设备的远程访问；二是作为session的target，它也是tf::Session的接口，三是负责一或多个Worker Service的任务分配。一个Master对对应多个Worker的情况是这样的，在In-graph模式下，比如一机多卡，每个GPU卡都可以启动一个Server并输出Master Service和Worker Service，但只有一个Server输出的Master Service起实质作用，作为Session target并负责与其他Worker沟通工作。

Master

Worker Service也是一个RPC服务，调配本地设备实现子图的计算。每一个图计算任务称为task，task又分为ps（参数服务器）task和worker task（计算密集型任务处理），ps和worker的集合各自组成一个job。

Worker

Client构成TF的计算图并构建tf::Session来与集群会话。在分布式集群中，创建session需要指定target（target默认值在本地故单机模式无需指明），由上文所提的Master Service提供（通过server.target调用）。一个Client可与多个Server联接（In-graph模式）；一个Server也可以服务多个Client，也就是在Client构建session时将target同时指定一个server.

Client

Server与task是一一对应的，在构建节点时时使用tf.train.ClusterSpec方法，通过ip:port定义server所在位置，并指定task类型（所属job集合）。

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.train.Server创建Server，指定集群，jobname和task_index以作为Server节点的身份识别标志。

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

如果是In-graph模式，只需指定一个ip:port来说明在哪个Server的target上创建session，而在Between-graph模式下参数服务器只需要通过join( )方法使进程进入待命状态，而worker task节点需要使用tf.train.replica_device_setter( )方法来创建session，并建立ps task与worker task之间的映射，此方法在Between-graph模式中最为关键，需要提供cluster，worker_device, ps_device等参数。

withtf.device(tf.train.replica_device_setter(

worker_device="/job:worker/task:%d" %FLAGS.task_index,

cluster=cluster)):

replica_device_setter方法

在replica_device_setter方法中会根据ps_strategy设定参数在ps集群的分配方式，而且会判断其下的操作是否属于ps op，从而判断需要放在负责ps task还是负责worker task的节点上，TF源码中有具体的逻辑代码。下图即为ps op，需要放在负责ps task的节点，所有其他操作均为worker op，放在负责worker task的节点.

PS操作集合

Ref:

[1]http://download.tensorflow.org/paper/whitepaper2015.pdf

[2]https://tensorflow.google.cn/deploy/distributed