TensorFlow 架构

TensorFlow 旨在用于进行大规模分布式训练和推断,但 TensorFlow 也非常灵活,足以支持利用新的机器学习模型和系统级优化进行实验。

本文档对这种可以同时实现规模性和灵活性的系统架构进行了介绍。本文假定您基本熟悉 TensorFlow 编程概念,例如计算图、操作和会话。如需查看关于这些内容的简介,请参阅本文档。对分布式 TensorFlow 有一定了解也很有帮助。

本文档面向以下人员:希望以当前 API 不支持的某种方式扩展 TensorFlow 的开发者,希望针对 TensorFlow 进行优化的硬件工程师,负责扩展和分发的机器学习系统实现者,或想要深入了解 Tensorflow 的任何人。读完本文档后,您应该对 TensorFlow 架构有足够的了解,可以读懂和修改核心 TensorFlow 代码。

概述

TensorFlow 运行时是一个跨平台库。图 1 展示了它的一般架构。C API 可将不同语言的用户级代码与核心运行时分开。

TensorFlow 层

图 1

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

  • 客户端
    • 将计算定义为数据流图。
    • 使用会话启动图执行。
  • 分布式主服务
    • 根据 Session.run() 的参数定义,从图中裁剪出特定子图。
    • 将子图分成在不同进程和设备中运行的多个块。
    • 将各个图块分发给工作器服务。
    • 通过工作器服务开始执行图块。
  • 工作器服务(每项任务一个)
    • 使用适合可用硬件(CPU、GPU 等)的内核实现来安排图操作的执行。
    • 向其他工作器服务发送操作结果以及从它们那里接收操作结果。
  • 内核实现
    • 对单个图操作执行计算。

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

TensorFlow 架构图

图 2

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

以下各部分更详细地介绍了核心 TensorFlow 层,并逐步说明了如何处理示例图。

客户端

用户负责编写可构建计算图的客户端 TensorFlow 程序。该程序可以直接组建各项操作,也可以使用 Estimators API 之类的便利库组建神经网络层和其他更高级别的抽象物。TensorFlow 支持多种客户端语言,但我们优先考虑 Python 和 C++,因为我们的内部用户最熟悉这些语言。随着功能的日益成熟,我们通常会将它们移植到 C++,以便用户可以通过所有客户端语言访问经过优化的实现。大多数训练库仍然只支持 Python,但 C++ 确实支持高效推断。

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

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

TensorFlow 架构图:客户端

图 3

代码

分布式主服务

分布式主服务会:

  • 对图进行裁剪,以获得评估客户端请求的节点所需的子图,
  • 划分图以获取每台参与设备的图块,以及
  • 缓存这些块,以便它们可以在后续步骤中重复使用。

由于主服务可以看到某一步的总体计算,因此它会应用常见的子表达式消除和常量折叠等标准优化。然后,它会在一组任务中协调经过优化的子图的执行。

TensorFlow 架构图:主服务

图 4

图 5 显示了示例图的可能划分。分布式主服务已对模型参数进行分组,以便将它们一起放在参数服务器上。

划分图

图 5

如果图边缘在划分时被切割掉,分布式主服务会插入发送和接收节点,以便在分布式任务之间传递信息(图 6)。

划分图

图 6

然后,分布式主服务会将各个图块传送到分布式任务。

划分图

图 7

代码

工作器服务

每项任务中的工作器服务会:

  • 处理来自主服务的请求,
  • 为包含本地子图的操作安排内核的执行时间,以及
  • 调解任务之间的直接通信。

我们会优化工作器服务,以便以较低的开销运行大型图。我们当前的实现每秒可以执行数万个子图,这可让大量副本生成快速、精细的训练步。工作器服务会将各个内核分派给本地设备并尽可能并行运行这些内核,例如通过使用多个 CPU 核或 GPU 流。

我们会特殊对待每对源设备和目标设备类型的发送和接收操作:

  • 本地 CPU 和 GPU 设备之间的传输使用 cudaMemcpyAsync() API 来重叠计算和数据传输。
  • 两个本地 GPU 之间的传输使用对等 DMA,以避免通过主机 CPU 进行代价高昂的复制。

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

  • 基于 TCP 的 gRPC。
  • 基于聚合以太网的 RDMA。

此外,我们还提供对 NVIDIA 的 NCCL 库的初步支持,以进行多 GPU 通信(请参阅 tf.contrib.nccl)。

划分图

图 8

代码

内核实现

运行时包含 200 多个标准操作,包括数学、数组操作、控制流和状态管理操作。其中每项操作都可以具有针对各种设备优化的内核实现。许多操作内核都是使用 Eigen::Tensor 实现的,后者使用 C++ 模板为多核 CPU 和 GPU 生成高效的并行代码;不过,我们可以随意使用 cuDNN 之类能实现更高效内核的库。我们还实现了量化,能够在移动设备和高吞吐量数据中心应用等环境中实现更快的推断,并使用 gemmlowp 低精度矩阵库加快量化计算。

如果将子计算表示为操作组合很困难或效率低下,则用户可以注册提供高效实现(用 C++ 编写)的其他内核。例如,我们建议为一些对性能要求苛刻的操作注册自己的混合内核,例如 ReLU 和 S 型激活函数及其对应的梯度。XLA 编译器具有自动内核混合的实验性实现。

代码