TensorFlow:大规模机器学习的异构分布式系统
TensorFlow从名称上看就是两个部分——张量tensor和流flow。非常形象的组合。众所周知，矩阵已经成为机器学习中的基础单元，若干的针对矩阵的计算优化使得现如今的机器学习成为可能。而一些矩阵的方法也是一些重要的机器学习算法的基础。张量就是矩阵概念的推广，其表示更多维度的矩阵。而计算流是一种抽象过程，在如今的深度学习领域，这种一层层地计算可以很形象地看做是张量在计算模型上的流动。而这里的流可以看做是更加一般的计算过程，可以在不同的层级间跨越式流动。
本文作者均来自Google Research:Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mane, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viegas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng
摘要
TensorFlow[1]是一个表达机器学习算法的接口，并且是执行算法的实现框架。使用TensorFlow表示的计算可以在众多异构的系统上方便地移植，从移动设备如手机或者平板电脑到成千的GPU计算集群上都可以执行。该系统灵活，可以被用来表示很多的算法包括，深度神经网络的训练和推算算法，也已经被用作科研和应用机器学习系统在若干的计算机科学领域或者其他领域中，例如语言识别、计算机视觉、机器人、信息检索、自然语言理解、地理信息抽取和计算药物发现。该论文描述了TensorFlow的接口和我们在Google构建的结构实现。TensorFlow API和参考实现都已经作为开源项目按照Apache2.0协议在2015年11月发布，可以在这里查看。
1 引言
Google大脑项目开始于2011年，目的是探索在科研和Google的产品中超大规模深度神经网络的使用。作为这个项目的早期工作，我们构建了DistBelif——第一代的可扩展分布式训练和推断系统[14]，这个系统工作得很不错。我们和其他Google的同事使用DistBelief进行了广泛的研究包括非监督学习[31]、语言表示[35,52]、图像分类模型和目标检测[16,48]，视频分类[27]、语音识别[56,21,20]、序列预测[47]、Go 的移动选择[34]、行人检测[2]、强化学习[38] 等等。另外，和 Google Brain 团队合作中，超过 50 个 Google 内部的团队和其他 Alphabet 公司也已经部署了使用 DistBelief 的深度神经网络在众多产品中，包括 Google Search[11]、广告产品、语音识别系统[50,6,46]、Google Photos[43]、Google Maps 和 街景[19]、Google 翻译[18]、Youtube 和很多其他的产品。
基于我们使用 DistBelief 的经验和对于期望用来训练和使用神经网络的系统特性和需求更加完备地理解，我们构建了 TensorFlow——第二代大规模机器学习模型的实现和部署的系统。TensorFlow 使用通过类似数据流模型的计算，将这些计算映射到不同的硬件平台例如使用包含一个或者多个 GPU 显卡的装有 Android 和 iOS 的单个机器上进行推断，到运行在数百台包含数千个 GPU 的大规模系统训练和推断。拥有一个单一的系统可以扩展分布到众多的平台上可以大大简化真实场景中机器学习系统的使用，正如我们在用分离的系统进行大规模训练和小规模的部署，会产生巨大的维护代价和较差的抽象效果。TensorFlow 的计算被表示为含状态的数据流图（在第二节详细讲解），我们聚焦在让这个系统足够灵活能够快速地实验研究中产生的新模型，并同时充分地提升产品级训练的性能和部署机器学习模型健壮性。为扩展神经网络训练搞更大的部署环境，TensorFlow 允许 client 简单地表达不同类型的并行通过复制和并行执行一个核心模型数据流图，依赖不同计算设备合作更新一个共享的参数或者其他的状态。 对计算描述的微妙变动可以使用较低的代价来达到和尝试很多不同的并行的方法。一些 TensorFlow 的用途借助参数更新的一致性来实现灵活性，我们可以在一些更大的部署环境中轻易表达和利用这些同步上的松弛。对比 DistBelief，TensorFlow 的编程模型更加灵活，性能也更好，支持在大规模的异构硬件平台上训练和使用很多的模型。
DistBelief 的内部用户已经切换成 TensorFlow 了。这些客户依赖 TensorFlow 来研究和产品，执行诸如在移动电话计算机视觉模型的推断到使用数百台机器进行千亿级样本的千亿级参数的深度神经网络的训练[11,47,48,18,53,41]。尽管这些应用集中在机器学习和深度神经网络上，我们希望 TensorFlow 的抽象可以用在其他的领域中，例如其他的机器学习算法或者可能其他类型的数值计算。我们按照 Apache 2.0 协议在 2015 年 11 月开源了 TensorFlow API，可以在 www.tensorflow.org 查看。
本文下面的部分更加细致地描述了 TensorFlow。第二节介绍编程模型和 TensorFlow 接口的基本概念，第三节介绍单机和分布式的实现 。第四节给出了基本编程模型的扩展，第五节介绍了一些基本实现的优化方法。第六节给出了一些使用 TensorFlow 的实验结果，第七节描述了一些使用 TensorFlow 编程的 idiom，第九节则是一些在 TensorFlow 核心外围的工具。第十节和第十一节分别讨论了未来和相关的工作，最后一节给出了总结性想法。
2 编程模型和基本概念
TensorFlow 的计算由一个有向图描述，这个图中由一个节点集合组成。该图表达了数据流计算，作出了一些类型的节点保持和更新持久状态和让分支及循环控制结构类似于 Naiad 的行为方式的扩展。客户一般都会使用 TensorFlow 支持的前端语言（C++或者Python）构建一个计算图。在图 1 中展示了一段样例使用 Python 构建并执行了一个 TensorFlow 的计算图，结构计算图在图 2 中展示。
在一幅 TensorFlow 图中，每个节点（node）有一个或者多个输入和零个或者多个输出，表示一种操作（operation）的实例化。流过图中正常的边（输出到输入）的值都是张量（tensor），任意维度的数组其中基础元素类型是指定的或者在图的构造过程中自动推断出来的。特别的边，我们称之为控制依赖（control dependencies），同样也存在在图中：这类边上没有数据流过，但是他们表示源节点必须在目标节点的控制依赖开始执行前完成运行。因为我们的模型包括可变状态，控制依赖可以被直接用来强制在保证关系的发生。（Since our model includes mutable state, control dependencies can be used directly by clients to enforce happens before relationships.）我们的实现同样会插入控制依赖来确保独立操作之间的顺序，比如说作为控制内存使用最高峰值的方式。
操作和核(Kernel)
一个操作有一个名字。它表示一个抽象的计算（比如说，“矩阵相乘”或者“相加”）。一个操作可以有属性（attribute），所有的属性必须提供或者在图构造的过程中推断出以实例化一个节点来执行操作。属性通常的使用方式是让操作在不同的张量元素类型上多态（例如，两个 float 类型的张量和两个 int32 类型的张量）。核（kernel）是一种操作的特别实现，可以运行在一个特定类型的设备上（如 CPU 或者 GPU）。TensorFlow 的 binary 定义了可以通过注册（registration）机制实现的操作和核的集合上，这个集合可以通过连接额外的操作/核的定义/注册。表 1 展示了内置于 TensorFlow 核心库的一些操作类型。
会话(session)
客户端通过创建会话（session）和 TensorFlow 系统进行交互。为了创建一个计算图，会话接口支持外部（external）方法来提升当前由包含额外节点和边的会话的图（当会话创建时初始的图是空的）。另一个由会话接口提供的主要的操作就是 Run，以需要计算的输出名称和替换某些输出节点的张量的操作集合作为其参数输入。通过控制 Run 的参数，TensorFlow 的实现可以计算所有节点的必须执行传递闭包来计算需要的输出，然后安排执行合适节点来保证他们的依赖关系（在3.1小节详细讲解）。大多数 TensorFlow 的使用都是针对一个图启动一个会话，然后执行整个图或者通过 Run 调用来执行分离的子图数千或者数百万次。
变量(variable)
在大多数计算中，图都是执行多次的。大多数的张量在一次执行后不会存活。然而，变量（variable）是一种特别的操作可以返回一个在图执行若干次过程中存活的持久化的可变张量的句柄。这个句柄可以传递给一系列特定的操作，例如 Assign 和 AssignAdd（等同于 +=）就可以改变其引用的张量了。对应 TensorFlow 在机器学习中的应用，模型的参数典型地就存放在变量引用的张量中，并作为模型训练图的 Run 的一部分进行更新。
3 实现
TensorFlow 系统的主要部分就是客户端，它使用了会话接口来和 master 及一个或者多个的 worker processes 进行通信，每个 worker process 负责对一个或者多个计算设备（CPU 核或者 GPU card）的任意访问和在这些设备上进行图节点的计算按照 master 的要求执行。我们有本地和分布式实现的 TensorFlow 接口。本地实现通常是客户端、master 和 worker 都是在同一台机器上在一个单一的操作系统进程（可能包括多个设备，比如说装了多个 GPU card的设备）上运行。分布式实现采用了本地实现的很多的代码，但是扩展了对客户端、master 和 worker 可以在不同的机器的不同的进程上运行的场景支持。在我们的分布式环境中，这些不同的任务对应于 cluster 调度系统分配在 job 中的容器中[51]。这两种不同的模式在图 3 中进行的展示。本节剩下的部分讨论了在两种实现中遇到的问题，3.3 节讨论了针对分布式实现的一些问题。
设备
设备是 TensorFlow 的计算核心。每个 worker 负责一个或者多个设备，每个设备有一个设备类型和一个名字。设备名字由识别设备类型的部分，在 worker 中的设备索引，以及在分布式设定中，worker 的 job和任务（或者 localhost 当设备是和进程在同一机器时）的标志构成。一些例子如/job:localhost/device:cpu:0 或者 /job:worker/task:17/device:gpu:3。我们已实现了 CPU 和 GPU 的设备接口而其他的设备类型也有了通过注册机制完成的设备实现方式。每个设备对象负责管理分配和解除分配设备内存，对在 TensorFlow 实现中的更高层请求任意 kernel 的执行调度管理。
张量
实现中的张量是一种有类型的、多维度数组。我们支持若干张量元素类型，包含大小为从 8 bit 到 64 bit 的带符号和无符号整型，IEEE 浮点数和双精度类型、复数类型和字符串类型（任意长的字节数组）。合适大小的后台存储通过一个分配器进行管理，该分配器由张量所处的设备确定。张量的后端存储缓存是引用计数的并在没有引用存在时解除分配。
3.1 单设备执行
首先考虑最简单的执行场景：单一的worker进程运行在单一的设备上。图上的节点按照代表节点之间的顺序执行。特别地，我们会在每个节点上保持一个计数来记录这个节点上还没有执行的依赖。一旦这个计数变为 0，节点就可以被调度使用，并会加入到待续的队列中。待续队列按照某个非指定的顺序处理，指派节点执行的kernel 到设备对象上。当一个节点完成执行，所有依赖这个完成的节点的节点的计数都会增加
3.2 多设备执行
一旦系统有了多个设备，有两个主要的复杂情形出现：确定图中每个节点的计算所处的设备，然后管理由上一步确定的置放决定产生的设备间的所需的数据通信。后续部分讨论这两个问题。
3.2.1 节点的置放
给定计算图，TensorFlow 实现的主要责任之一就是将计算映射到可用的设备集合上。这个算法的简单版本下面给出。参见第 4.3 节有关该算法支持的扩展。
该置放算法的输入是一个代价模型，包括对每个图节点的输入和输出张亮的规模的估计，和对每个节点在给于其输入张量时的计算时间的。这个代价模型或者是基于关联不同操作类型的启发式规则的静态估计，或者基于实际的为更早的图的执行而做的置放决定集合衡量。
置放算法首先运行模拟的图的执行过程。模拟按照下面描述进行，对每个节点使用贪心策略选择一个设备。节点到设备的置放过程也是用作真实执行的置放。
置放算法从计算图的源点开始，在系统中的每个设备上模拟相应的活动。对每个在遍历中抵达的节点，可选 available 设备的集合会被考虑到（设备可能会由于其没能提供实现了特定操作的kernel而不可选）。对那些拥有多个可选设备的节点，置放算法使用一种贪心策略来检查在每个可能谁被上置放节点需要完成的时间的效果完成决策。这种启发式规则考虑了根据代价模型在那种设备上估计的和衡量的执行时间，还有任何用来从其他设备传输输入到该节点的通信的代价。其中节点的操作完成最快的设备会被选作该操作的设备，置放决策然后会继续针对图中其他的节点进行处理，包含那些已经做好模拟执行的下游节点。第 4.3 节描述了一些扩展，让用户可以提供提示和部分限制来指导置放算法。这个算法现在还在持续开发的过程中。
3.2.2 交叉设备通信
一旦置放节点被计算，图会被分割成一系列的子图，每个子图在分布在1台设备上。任何从x到y的交叉设备的边会被移除，然后被一个从x到新的发送节点的边代替