在线文档教程

TensorFlow Architecture(TensorFlow架构)

TensorFlow架构

我们为大规模分布式训练和推理设计了TensorFlow,但它也足够灵活,可以支持新机器学习模型和系统级优化的实验。

本文档描述了使这种规模和灵活性相结合的系统架构。它假定您对TensorFlow编程概念(如计算图,操作和会话)基本熟悉。有关这些主题的介绍,请参阅入门指南。对分布式TensorFlow的一些熟悉也会有所帮助。

本文档面向想要以某种方式扩展TensorFlow的开发人员,而这些开发人员并不支持当前的API,希望针对TensorFlow优化的硬件工程师,实现扩展和分发的机器学习系统的实现者,或者希望在Tensorflow领导下进行研究的人员。阅读完它后,您应该了解TensorFlow体系结构,以便读取和修改核心TensorFlow代码。

Overview

TensorFlow运行时是一个跨平台的库。图1说明了其总体架构。AC API将不同语言的用户级别代码与核心运行时分离。

图1

本文档重点介绍以下几个层次:

  • 客户端

  • 将计算定义为数据流图。

  • 使用会话启动图形执行

  • 分布式主站

  • 根据Session.run()的参数定义修剪图中的特定子图。

  • 将子图划分为多个部分,这些部分运行在不同的流程和设备中。

  • 将图片分发给工作服务。

  • 由工作人员服务启动图块执行。

  • 工人服务(每项任务中的一项)

  • 使用适合可用硬件(CPU,GPU等)的内核实现来安排图操作的执行。

  • 向其他工作者服务发送和接收操作结果。

  • 内核实现

  • 执行单个图形操作的计算。

图2说明了这些组件的相互作用。“/ job:worker / task:0”和“/ job:ps / task:0”都是带有工作服务的任务。“PS”代表“参数服务器”:负责存储和更新模型参数的任务。其他任务在更新这些参数时会优化参数。任务之间的这种特殊分工不是必需的,但是分布式培训是很常见的。

图2

请注意,分布式主和辅助服务仅存在于分布式TensorFlow中。TensorFlow的单进程版本包含一个特殊的Session实现,它执行分布式主机的所有操作,但只与本地进程中的设备进行通信。

以下部分将更详细地介绍核心TensorFlow图层并逐步完成示例图形的处理。

客户端

用户编写构建计算图的客户端TensorFlow程序。该程序可以直接编写单个操作,也可以使用Estimators API等便利库来组成神经网络层和其他更高级别的抽象。TensorFlow支持多种客户端语言,并且我们优先考虑Python和C ++,因为我们的内部用户最熟悉这些语言。随着功能变得更加成熟,我们通常将它们移植到C ++,以便用户可以从所有客户端语言访问优化的实现。大多数培训库仍然是Python,但C ++确实支持高效推理。

客户端创建会话,将会将图形定义作为tf.GraphDef协议缓冲区发送到分布式主机。当客户端评估图中的一个或多个节点时,评估会触发分布式主机的调用以启动计算。

在图3中,客户建立了一个图表,将权重(w)应用于特征向量(x),添加偏差项(b)并将结果保存在变量中。

图3

代码

  • tf.Session分布式master:

  • 修剪图形以获得评估客户请求的节点所需的子图,

  • 分割图形以获得每个参与设备的图形片段,并且

  • 缓存这些碎片以便它们可以在随后的步骤中重新使用。

由于主人看到一个步骤的整体计算,它会应用标准优化,如常见的子表达式消除和常量折叠。然后它协调一组任务中优化子图的执行。

图4

图5显示了我们示例图的一个可能分区。分布式主站将模型参数分组,以便将它们放在参数服务器上。

图5

在分区切割图边时,分布式主机会插入发送和接收节点以在分布式任务之间传递信息(图6)。

图6

然后分布式主机将图形部分发送到分布式任务。

图7

代码

工人服务

每项任务中的工作人员服务:

  • 处理来自主人的请求,

  • 为组成本地子图的操作安排内核的执行,以及

  • 调解任务之间的直接沟通。

我们优化了工作服务,以便以较低的开销运行大型图形。我们目前的实现可以每秒执行数以万计的子图,这使大量副本能够进行快速,细致的训练步骤。工作服务将内核分派给本地设备,并尽可能并行运行内核,例如通过使用多个CPU内核或GPU流。

我们专门针对每对源和目标设备类型的Send和Recv操作:

  • 本地CPU和GPU设备之间的传输使用cudaMemcpyAsync()API来重叠计算和数据传输。

  • 两个本地GPU之间的传输使用对等DMA,以避免通过主机CPU进行昂贵的复制。

对于任务间的传输,TensorFlow使用多种协议,包括:

  • gRPC over TCP.

  • 融合以太网上的RDMA。

我们还对NVIDIA的NCCL库进行了多GPU通信的初步支持(请参阅参考资料tf.contrib.nccl)。

图8

代码

内核实现

运行时包含200多种标准操作,包括数学,数组操作,控制流和状态管理操作。这些操作中的每一个都可以为各种设备优化内核实现。许多操作内核都使用Eigen :: Tensor实现,它使用C ++模板为多核CPU和GPU生成高效的并行代码; 然而,我们自由地使用像cuDNN这样的库,可以实现更高效的内核实现。我们还实施了量化,可在移动设备和高吞吐量数据中心应用等环境中实现更快的推理,并使用gemmlowp低精度矩阵库加速量化计算。

如果将子计算表示为操作组合是困难或低效的,则用户可以注册额外的内核,这提供了用C ++编写的高效实现。例如,我们建议您注册您自己的融合内核以执行一些性能关键操作,例如ReLU和Sigmoid激活函数及其相应的渐变。XLA编译器具有自动内核融合的实验实现。

代码

  • OpKernel 接口