TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems
摘要：TensorFlow是一组用于描述机器学习算法的接口，以及对接口的实现。用Tensorflow描述的计算过程不必修改（或只需稍微修改）就能在异构芯片（比如小至手机，大至分布式集群）中执行。这套接口非常flexible，能描述很多算法，比如深度神经网络中的training和inference算法等。本文介绍TensorFlow，以及Google对它的一种实现。
1. 引言
Google Brain项目于2011年开始探索very-large-scale深度神经网络。作为早期工作的一部分，我们开发了BistBelief作为第一代scalable分布式training和inference系统。BistBelief已经应用于大量产品中。
基于BistBelief的经验和不足，以及对training需求更深刻的理解，我们开发了第二代系统TensorFlow，其定位是large-scale机器学习模型[注：相比于very-large-scale深度神经网络更加通用]。TensorFlow用类似于dataflow graph的方式描述计算过程，然后把它映射到多种多样的硬件平台（比如Android/IOS上的inference，单卡GPU上的training&inference，千级以上GPU卡上的training），一套系统覆盖这么多平台是很有意义的。TensorFlow计算过程用stateful dataflow graphs（有状态数据流图）描述（详见第二章），其目标是：1.足够flexible从而能快速验证新模型；2.足够high performance及robust从而产品级支持对机器学习模型的training和deployment。为了大规模deployment，TensorFlow提供了复制和并行执行dataflow graph的并发方式，各devices合作更新一组代表state的共享参数。不同并发方式能用少量的代码加以区分。相比于BistBelief，TensorFlow更加flexible，高性能以及支持异构硬件。
本文组织如下：第二章更细节地介绍TensorFlow接口，第三章介绍它在单机及分布式上的实现，第四章介绍对基础接口的一些扩展，第五章介绍对基础实现版本的一些优化，第六章是实验，第七章提供一些实践中我们发现有益的TensorFlow使用风格，第九章介绍一些辅助工具，第十和十一章分别介绍未来展望和相关工作，第12章总结。
2. 接口及概念
在TensorFlow中，计算过程被描述为一张有向图（graph）以及一些扩展，扩展包括：允许图中某些节点维护和更新persistent状态；图中加入分支和循环控制机制（类似Naiad[注：一个及时的数据流系统]）。用户可以用不同语言前端构图。图1是用Python前端编写的示例，它构造的计算图见图2。
Graph中node是operation的实例，可以有0或多个输入以及0或等多个输出。Graph中沿普通边流动的数据是tensors，其element type由client程序指定或构图时通过类型推导得出。Graph中还有特殊边（即control dependencies），代表operation之间的happens-before关系，通过可变的mutable state来实现。控制依赖并不总是为了数据依赖，比如有时为了内存消耗不要太多，也会在无数据依赖的operation间加上控制边，以防它们同时执行内存消耗太多。
Operations and Kernels
Operation代表一个数学计算。它可以有属性，但属性必须在构图时就写入，比如element type是一种常见的属性，使能不同实现版本，即多态。
Kernel是operation在特定硬件上的实现。
Sessions
Session是client跟master、worker processes（第三章）交互的通道。Session接口包括Session.Extend方法用于扩展一个图（从0扩展即新建）、Session.Run方法用于提供输入（不一定是graph的起点，有可能直接是graph中的某个占位符）和待计算nodes之后启动计算过程。一般来说，构图一次运行多次。
Variables
大多数图都会被执行多次，而tensor在单次执行中是不会被保存下来的，除了一种例外，即通过Variable。Variable是一种特殊的operation，它返回一个操作持久可变（persistent mutable）tensor的handle，该handle在图的整个执行期都是存活的。而且该handle可以作为其他operation（比如AssignAdd）的输入用于改变其对应的tensor。在机器学习的程序中，模型的parameters一般被存放在被Variable操作的tensor中，在图的执行过程中被更新。[注意，Variable是Operation，而非Tensor]
3. 实现
TensorFlow接口的实现主要由四部分组成：client，即用户使用各种前端语言编写的程序；master，负责构图和分发执行图任务；worker processes，即执行具体计算的进程；每个worker process负责在一或多个devices（如CPU核、GPU卡）上执行具体计算任务。TensorFlow接口有local和distributed两种模式，local模式用于client、master和worker process都运行在同一个机器上的同一个OS进程中的场景（但是devices可以有多个，比如多卡GUP），distributed模式则能够支持client、master和worker process运行在不同机器的不同OS进程中的场景。在我们的distributed环境中，不同的任务指cluster scheduling system所管理的jobs中的containers。图3展示了两种模式。3.1,3.2介绍两种模式的共同问题，3.3介绍了一些distributed版本的一些特殊问题。
Devices
在我们的实现中，每个device object都有一个type和name。其中name由三个部分组成（/%0/%1:%2) % (worker的job和task的唯一标识，type，在worker’s devices中的index），比如”/job:worker/task:17/device:gpu:3”，”/job:localhost/device:cpu:0”（local模式直接写localhost）。Device object负责分配和释放device内存，以及调用具体的kernel。（目前实现了CPU和GPU版本的device object。）
Tensors
在我们的实现中，tensor是有类型的多维数组（typed multidimensional array）。所谓type指element type，支持signed,unsigned int,float,double,复数,string。Tensor的存储由tensor所运行在的device负责，采用引用计数的方式，0引用时释放。
3.1 （Local & distributed模式）：单device执行
这是最简单的场景：单worker，单device。每个node都记录了它所依赖的前序节点的执行情况，当前序都执行完毕时，该node加入一个ready queue。该队列中的节点乱序执行。
3.2 （Local & distributed模式）：多device执行
多device需要考虑两个问题，一是graph中各node在哪个device上执行[注：TF其实也支持单node内的devices间并行，后文介绍]，二是各devices间的通信。下面分别介绍。
3.2.1 node在哪个device执行
先介绍一个最简单的算法。4.3节介绍对该算法的扩展，能允许用户提供线索以及partial constraint来指导该算法。
该算法以一个代价模型为基础，该模型能预测指定node在指定device上的执行时间。该算法是对graph的一次符号执行，输出是为每个node都选择好device。具体符号执行过程是一个贪心算法的过程，从起始node开始，把node指定到能最快完成该node的device上（需要考虑kernel执行时间以及通信时间），依次向下分析各node。预估可以完全静态预估，也可以基于该图之前的执行情况来进行预估。
3.2.2 Cross-device通信
一旦node的device选择做好之后，整个graph被切割成若干subgraphs，每个device一个subgraph。跨device的边都插入一组send/recv node，如图4。
运行时，所有通信任务都在send/recv nodes完成，其他nodes则不必考虑通信问题。
在插入send/recv nodes时，把同一tensor的sed/recv nodes合并从而保证每个tensor只被传递一次，即device上对该tensor的内存指分配一次。（如图4中的b,c）
通过这种方式，能把各devices上node的调度任务交给worker process，从而在master层面实现了去中心化。即send/recv节点负责worker process间和devices间的通信，master只负责发送RUN请求给worker，从而使得整个架构scalable。
3.3 distributed模式
Distributed执行跟多devices执行很类似，唯一区别是send/recv nodes必须通过TCP/RDMA在workers间通信，因为各worker不在一台硬件上了。
Fault Tolerance
Distributed执行错误有很多检测方式。我们主要实现了两种：(a)send/recv通信异常，(b)master定期对各worker process进行健康检查。
当检测到错误时，整个graph终止并从头开始重新执行。但是注意到，Variable代表整个graph生命周期中都存在的tensor，我们支持了一致的检查点和状态恢复。每个Variable节点都连接到一个Save node，这些Save node定期执行，比如每N次迭代或每N秒一次。当Save node执行时，Variables中的内容从内存写入持久存储器（比如分布式文件系统）。每个Variable node也连接一个Restore node，Restore node只在重启后的第一次迭代启用。关于如何在图形的执行中只启用部分nodes的细节，请参见第4.2节。
4. 扩展
本章介绍对第二章基础接口的一些扩展。
4.1 梯度计算
很多神经网络中的优化算法要计算一个函数对个输入的梯度（偏导数），比如随机梯度下降法。TensorFlow内置了自动计算梯度的支持。如果tensor C经过一系列复杂的operations后依赖于一些输入tensors{Xk}，则TensorFlow提供的内置函数能够通过在图中加入反向子图的方法返回tensors{dC/dXk}。添加计算梯度的子图的过程如下。
当要计算{dC/dI}时，首先找到I至C的子图路径，然后从后之前分析该路径，每遇到一个operation就在图中添加一个新反向operation。注意，要把正向operation的输入输出额外作为反向operation的输入。图5是图2加入梯度计算后的样子，其中灰色的边是实际未使用到的输入[是否需要正向的输入输出作为反向的输入是operation相关的，比如乘法需要，加法就不需要]。只需在图1中加入如下一行代码：
[db, dW, dx] = tf.gradients(C, [b, W, x])
自动梯度计算增加了优化（尤其是内存使用）的复杂程度。因为当执行正向图时，启发式地决定下一个执行哪个node，这意味着输出内存很快被后续node重用，即使启发式算法不够好用户也可以通过手动改变图的构造顺序或者加入控制边的方式来提高内存重用。然而对于反向图，用户不能对它直接修改，所以手动方式就失效了，只能依赖于启发式算法，然而不幸的是启发式算法理论上也没有优秀的解，因为前向开头的内存要在反向最后用到，所以不得不让它们一直停留在内存中，这会消耗掉大量宝贵的GPU内存从而不必要地限制了图的计算量。我们正积极致力于内存管理的改进，以更好地处理此问题。选项包括更成熟的启发式算法来确定节点执行顺序，重新计算tensor而不是将他们保留在内存中，以及将需要长期存在的tensor从GPU内存中换到相对更丰富的CPU内存中。
4.2 Partial Execution
Client经常只希望计算整个graph的一个subgraph，Session.Run方法对该需求提供了支持，它能够允许client插入任意数据作为任意边的输入，也能接受任意边上的tensor作为输出数据。
Node都有名字，node的输出通过索引来表明它是第几个输出，比如”bar:0”代表”bar”节点的第一个输出。
Session.Run中的两个参数用于描述所需计算的子图部分。1.{name:port -> “fed”tensor值}表示将”fed”tensor值作为哪个边的输入（可以是一组）；2.[name:port, ]表示需要把哪个operation的第几个输出返回给client（可以是一组）。
基于两个参数，图会被重新转换。对输入，图中每个name节点的第port个边被断开，重新连接到一个feed节点（如图6中b与c的连接边被替换成feed与c的连接边）；对输出同理，替换成连接到一个fetch节点的边（如图6中fetch）。一旦在图中有了feed/fetch节点，就可以从fetch开始反向依赖分析出图中需要执行的部分子图。图6中就分析出d，e节点不必执行。
4.3 device约束
Client能够通过device约束来控制每个node能在哪个或哪些或满足哪类条件的device上执行。比如“该node只能在type是GPU的device上执行”，“该node必须跟v3 node在同一个device上执行”等。通过约束，就可以在兼顾node执行效率的同时对某些问题进行人为约束，比如约束一个device上的内存消耗上限。
相应地修改3.2.1节的算法如下：首先对各node分析出它能运行的devices集合，然后筛选出满足约束的devices集合。
4.4 Control Flow
很多时候，使用控制流能更简洁高效地描述机器学习算法。正如Arvind描述的数据流机器学习方法一样，TensorFlow引入了少量基础的控制流运算符，从而使得TensorFlow能处理循环数据流图。具体来说，Switch和Merge运算符让我们可以根据bool tensor的值跳过整个subgraph的执行；Enter，Leave和NextIteration运算符让我们可以表达iteration。使用这些控制流运算符，可以轻松地将高级语法（如if和while）转换成dataflow graph。
TensorFlow runtime实现了tags和frames的概念，类似于MIT的TaggedToken机器。循环的每次迭代都有一个唯一tag，它的执行状态用一个frame表示。一次迭代的输入只要已经就绪，该iteration就可以进入执行了，因此可以并行执行iterations。
TensorFlow使用分布式协调机制来执行带控制流的graph。通常，一个循环中可以包含分配给不同devices执行的node，因此管理循环状态问题转化为分布式终止检测问题。TensorFlow的解决方案是graph改写。在graph分区时自动为每个分区加上控制nodes，这些nodes可以安排每次迭代的开始和终止，并确定循环是否终止。对于每次迭代，拥有循环终止谓词的device都会向每个参与循环体运算的device发送一个短控制消息。
当一个graph包含控制流operations时，梯度下降子图就必须考虑它们。比如包含if时，必须知道正向时选择了哪个分支，然后对该分支构造梯度下降子图；包含while时，必须知道迭代执行了几次。它们都需要依赖计算的中间变量。TensorFlow采用的方法是改写graph，记录下反向所需的值。具体本文不展开。
4.5 Input Operations
直接从worker的存储系统中读入Tensor到worker的内存中。相比于feed方式（从client的存储系统发送到worker的存储系统，再读入worker的内存中）少一次网络传输。
4.6 Queues
Queues是TensorFlow中一项有用的功能。它允许graph的不同部分异步执行，并通过enqueue/dequeue交互数据。Enqueue的激活条件是queue中有可用空间；Dequeue的激活条件是其所需条件均已就绪。Queue的一种用法是在前序计算还在继续时提前从磁盘中预取数据。Queue还可以用于分组，比如将一个负责的梯度计算分解成一些列子梯度计算的聚合；再比如把不同长度的输入进行流水线级拆解和聚合，从而能够更均匀地分配各任务。
除了FIFO queue外，TensorFlow还实现一个shuffling queue，用于把queue中元素随机打乱顺序。
4.7 Containers
TensorFlow通过Containers机制来管理longer-lived mutable state。作为longer-lived mutable state的代表，Variable就保存在一个Container中。默认Container是全周期永存的。也允许创建其他有名称的containers。通过containers可以在完全不相连的graphs之间（甚至不同Sessions之间）共享state。
5. 优化
本章介绍TensorFlow实现中的一些优化手段。
5.1 Common Subexpression Elimination
Graph上合并公共operations。
5.2 优化通信和内存使用
Scheduling operations的顺序可以提升性能（不仅是执行时间，甚至主要是内存和通信代价）。具体来说，scheduling以减少两个连续operations（中间结果需要保存在内存中）的时间间隔，从而减少了内存峰值，这对于内存稀缺的GPU尤其有意义。此外，scheduling跨devices的数据通信能减少对网络资源的争抢。其实scheduling可以有很多优化子目标，这里重点介绍一个尤其有意义的。Recv节点在正常情况下会很保守地在很早的时候启动，然而这没有必要，通过使用运筹学中常见的as-soon-as-possible/as-late-as-possible（ASAP/ALAP）计算，能够分析出graph中相关的关键路径，以估计合适启动recv节点。然后通过在graph中加入控制边，来达到延迟启动的效果。
5.3 Asynchronous Kernels
Synchronous kernels在其计算函数执行完之后就结束了，而TensorFlow还支持了non-blocking kernels（即Asynchronous kernels）。Asynchronous kernels使用了不一样的接口，能够给其计算函数传递一个Continuation，该Continuation在计算函数结执行完后调用。对于有许多活动线程的环境，这一点对其内存等资源的使用时一种优化，避免了等待事件（比如IO）时持续占用执行线程。需要实现Asynchronous kernels的operations包括：recv，Enqueue，Dequeue等，避免了它们在同步模式下对资源的占用。
5.4 使用高性能库打造Kernels
比如矩阵乘使用BLAS和cuBLAS，卷积使用cuda-convnet和cuDNN库。
TensorFlow大量使用了Eigen库，并且扩展了Eigen以支持任意维tensor。
5.5 精度换性能
对于对精度不敏感的机器学习算法，比如典型的training神经网络。我们可以牺牲精度以换取通信效率。比如float32转成float16再通信（float32->float16->float32），相当于网络带宽提升了一倍。
6. 现状和经验
目前已开源，有很多示例和文档。已支持C++和Python前端。
在把一部分用DistBelief写的算法移植到TensorFlow的过程中，我们总结了一些经验，下面给大家分享下。
具体来说，我们讨论inception（目前最先进的卷积神经网络）。它负责吧224x224像素的图片分类到1000个类别（比如；猎豹，垃圾车等）。这个网络中有13.6 million个学习参数，在TensorFlow中将转化成一张包含36,000个operations的graph，仅单图推理就需要2 billon个mula operations。
构图后我们发现assembling和debugging这36,000个operations很有挑战。验证正确性非常困难，因为这个网络具有随机性，而且只有在几个小时后才知道对不对。在这种情况下，我们发现下面这些策略对于把模型转移到TensorFlow中是实用的。
写一个展示模型参数数量的工具。该工具能发现复杂网络中的细微错误，比如由于auto-dup引入的错误实例化operation和variable的错误。
从小case入手逐步做大。我们第一个移植的网络是一个应用于CIFAR-10数据集的小网络，调试这个网络可以阐明单个operation（如max-pooling）中的微妙情况。
先保证不学习时网络的正确性。把学习率设置成0能让网络行为相对静态化，有助于发现一些问题。
先在单机调通在去distributed环境。
防止数值错误。数学库在处理non-finite浮点数时行为经常不一致，最好在调试时就实时发现这些问题。
通过运行子网络，来了解数值误差范围。在TensorFlow和另一个系统上分别运行同样的子网络，看它们的误差是否在一定范围内。
上面这些被实践证明有效。最终我们的TensorFlow版本比DistBelief版本快了6倍。
7. 通用编程习语
TensorFlow的基础数据流图模型可以多种方式用于机器学习应用程序。我们关注的领域之一是加速计算密集型网络在大型数据集上的训练。本节介绍为此目的而开发的几种技术，并说明如何使用TensorFlow实现这些技术。
本节假设模型使用随机梯度下降法（SGD）对中等大小的mini-batches（100-1000个示例）训练模型。
Data Parallel Training
加速SGD的一种简单技术是把单个operation（图中model）的运算在mini-batch内实现并行。比如说，mini-batch包含1000个元素时，我们可以把operation复制10份，每个operation处理100个元素的梯度计算，最后把这些梯度综合起来再同步地更新参数，这样从外部看起来跟单个网络串行处理这一个mini-batch的行为是一样的。单client驱动整个training循环，如图7。
也可以异步地实现data parallel。同样把operation复制多份，但跟同步式不同，异步式的每个operation都独自更新parameters，而不是收集好全部operation的输出才更新。每个operation对应一个client，如图7。
Model Parallel Training
Operation级别并行，即每个operation都确定在一个device上执行，整图的各operation可以分别在不同device上执行。如图8，展示了TensorFlow在三个devices上并行的图示。
Concurrent Steps for Model Computation Pipelining
上面介绍的都是devices间并行，这里说的是单device内流水并行，从而充分利用device的资源。跟data parallel有些类似，只是这里不需要把operation复制多份分发到多个devices，而是把单device上的运算分解成steps从而达到流水的效果（filling in the gaps）。如图9[图中model和update复制多份单是都在单个device上，update和model实现流水线]。
8. 性能
9. 工具
9.1 Tensor Board：graph可视化，summary统计信息可视化
9.2 性能可视化工具
10. Future work
继续开发新网络，在此过程中发现问题并完善TensorFlow。
模块化子图重用，用户可以把一部分子图发布，别的用户可以直接用这一部分，而且这个子图可以是任意前端写的。
TensorFlow性能优化有几个具体方向。JIT compiler利用运行时信息（如shape）把子图编译成更高效的版本，该编译器可以理解图的语义，并能通过loop-fusion, blocking和tiling等技术提升性能；此外可以对特殊shape的场景实现特例化版本。
Node placement和node scheduling算法改进。
11. Related Work
很多其他系统跟TensorFlow在诸多方面有可比较性。Theano[7]，Torch[13]，Caffe[26]，Chainer和Computational Network Toolkit是一些主要用于training神经网络的系统。跟distributed TensorFlow不同，这些系统都只能把计算映射到单机上执行。下面说一些相同之处，跟Theano和Chainer类似，TensorFlow也支持符号差分，从而可以更轻松地编写基于梯度的神经网络网络优化算法。跟Caffe类似，TensorFlow也有一个用C++编写的core，从而简化了训练好的模型在各种生产环境（包括移动设备等内存和计算资源有限的环境）中部署的问题。
TensorFlow跟它的前身DistBelief和后来具有类似设计的系统（如Project Adam和Parameter Server ）有一些相同的设计特征，都允许把计算分布到多机器上的多devices中，而且允许用户通过相对high-level的方式描述机器学习模型。TensorFlow跟DistBelief和Project Adam的不同之处在于，TensorFlow更加flexible，也更易于表达各种机器学习模型和优化算法。通过把有状态的参数节点表示成Variables，并把Variable的更新抽象成graph中的operation，从而简化了整个设计。相比之下DistBelief，Project Adam和Parameter Server都使用一个独立的子系统来管理参数，专用于通信和更新参数值，显得有些复杂。
用于描述图像处理流水线的Halide使用了跟TensorFlow数据流图类似的中间表示。不同之处在于，Halide有其操作语义的高级知识，并使用这些知识生成高度优化的代码段（比如多个operations合并，考虑并行性和数据局部性）。但是Halide只能在单机运行，不支持分布式模式。将来，TensorFlow也希望使用使用类似的operations间动态编译框架。
也有些系统已经支持分布式模式以在集群上执行数据流图。Dryad和Flume演示了如何把复杂的workflow转化成数据流图。CIEL和Naiad则引入了对基于数据依赖的控制流的支持：CIEL把iteration描述成一个动态展开的DAG，Naiad则使用一个有环静态图来支持低延迟iteration。Spark对重复访问相同数据的计算专门进行了优化。使用的是“弹性分布式数据集”（RDDs）技术，它是前期计算的soft-state缓存。Dandelion能在包括GPU在内的一系列异构设备上执行数据流图。TensorFlow从这些系统中都有借鉴，最终使用了一个混合的数据流模型。其数据流调度器（选择下一个执行的node）使用了类似于Dryad，Flume，CIEL和Spark的算法；其distributed架构最接近Naiad，只用一张经过优化的图，并把图缓存在各devices上以最大程度地减少协调开销。类似于Spark和Naiad，当集群有足够的RAM时效果最佳。TensorFlow用一种混合的方式来支持iteration：多个子图副本同时执行，它们之间共享一些Variables。副本间可以通过Variables异步地共享数据，也可以使用同步的方式（如队列）。TensorFlow还graph内iteration的处理使用了两种混合的方式（CIEL和Naiad的混合体）：为了简单起见，每个node仅在其所有输入都已就绪时才激活（类似CIEL）；为了效率起见，graph表示为静态的周期性数据流（类似Naiad）。
12. 总结
文中介绍了TensorFlow，它是一个flexible数据流编程接口，以及对该接口在单机和distributed场景下的实现。已开源。
