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 接口