谈谈TensorFlow的分布式[1]

众所周知，TensorFlow来自于Google第一代深度学习DistBelief的改进。DistBelief基于参数服务器实现了数据并行，而每个数据分片内部可进行模型并行，如下图所示。

相比之下，TensorFlow并没有实现参数服务器本身，但可以由用户来实现类似的逻辑。在下边的代码中，创建了一个cluster，包含2个job：ps（Parameter Server）和worker，ps中运行2个task，worker中运行1个task。

当client执行56行的sess.run时，worker进程的job就会进行运算，并且把梯度发送到两个ps，因此这里实际上包含了模型并行的逻辑，如果worker进程中还包含其他job可以对不同的数据做运算，就实现了数据并行。在进一步讲述之前，先来介绍一下TensorFlow的术语：

Worker: TensorFlow由一个master进程和数个worker进程组成，master和worker之间通过接口通信，单机版和分布式均包含这种抽象，因此更改接口就实现了单机版和分布式的切换。每个worker负责管理一个或者多个Device（包含CPU和GPU），Device是TensorFlow具体计算的硬件抽象。

Cluster: cluster定义为TensorFlow分布式执行的任务集，cluster由一个或者多个job组成，每个job由多个task组成。每个task由一个worker来执行，

在前边的用TensorFlow手动实现参数服务器的代码中，在ps job里，并没有指定task运行的device：代码用的是with tf.device("/job:ps/task:0")而不是with tf.device("/job:ps/task:0/gpu:0")，此时TensorFlow会自动给task分配device。根据TensorFlow白皮书，TensorFlow提供了一个代价模型来给计算流图中各节点分配device的自动分配算法，然而根据代码，尽管代价模型存在，但目前TensorFlow(2016年7月版)并没有采用，而是用了很简单的自动分配算法，这部分对应的逻辑放在tensorflow/core/common_runtime/simple_placer.cc中。该算法优先分配GPU而非CPU，只有三种情况除外：一种是代码要求运行在CPU之上，另外两种是不同于TensorFlow白皮书的启发式算法，逻辑很简单，例如generator节点，就是那些在计算流图中没有入口只有一个出口的节点，会跟出口节点分配到同样的device；而一些元数据操作节点，例如reshape，也会跟执行数据操作的节点运行在同样的device之上。

因此通用流程为：给定计算流图，simple_placer总是给sink，source之外的节点分配GPU:0，如果想采用多个GPU，或者在分布式环境下，那么用户需要在他们的代码中手动指定device。因此，在目前看来，TensorFlow依赖程序员自身来实现数据并行和模型并行逻辑，因此如果想有效的利用硬件，需要仔细地设计代码。

关于TensorFlow的cost model，它的实现在tensorflow/core/graph/costmodel.h以及tensorflow/core/common_runtime/costmodel_manage.cc中，根据白皮书的定义，该模型会在计算流图生成的时候模拟每个device上的负载，并利用启发式策略估计device上的完成时间。启发式策略会根据如下数据调整device的分配：节点任务执行的总时间；单个节点任务执行的累计时间；单个节点输出数据的尺寸。但目前cost model还没有生效：在costmodel.cc里，AddToCostGraphDef函数"cnode->set_device(n->assigned_device_name());"，目前该model包含的的都是已经分配的节点，因此还没生效。

简单地评价下TensorFlow的分布式实现，从0.8版本包含该功能以来，感觉TensorFlow目前还处于发展的初期阶段，这甚至可以类比成Kubernetes这个同为Google力推的项目在2015年的状态，因此从效率，易用性（主要指分布式计算资源的管理和调度）上面还需要很多改进工作，特别是在算法的收敛性上，TensorFlow依赖使用者实现Parameter Server的逻辑，对于SSP等收敛性保证则缺乏考虑，至于说深度学习训练中的带宽节省等方面因素，也几乎没有考虑(即便cost model投入使用，也会只是简单的根据带宽消耗进行调度，缺乏进一步和整体考虑)，因此，在现阶段我们只能说在目前TensorFlow实现了基本的SDK，而留了更多的工作给未来，然而，也正是因为项目出自Google引发的从众效应，可以预计TensorFlow也会跟Kubernetes那样快速发展。不过从Roadmap上来看[1]，对于这些因素目前Google还不认为是大的问题。

本文内容来自韩国人对TensorFlow的研究[2]，互联网上对于TensorFlow的研究遍地开花，然而多数对于内在机理的剖析浮于表面，这也是本号转载的目的之一。本文只代表个人观点，如有偏颇之处敬请指正。

[1] https://github.com/tensorflow/tensorflow/blob/master/tensorflow/g3doc/resources/roadmap.md

[2] http://bettercstomorrow.com/category/studies/tensorflow/