2.4 TensorFlow架构
2.4.1 TensorFlow架构概述
TensorFlow运行时是一个跨平台库，图2.10展示了它的一般架构，其中的C API可将不同语言的用户级代码与核心运行时分开。TensorFlow在设计时借鉴了Android的成功架构，即其生态构建使得层次和背景各异的小微开发者能够简单设计应用，并且易于迁移到类似生态。
图2.10 TensorFlow核心架构
TensorFlow架构包括以下核心组件。
- 客户端：将计算定义为数据流图和使用会话启动图执行。
- 分布式主服务：根据Session.run()的参数定义，从图中裁剪出特定子图；将子图分成在不同进程和设备中运行的多个块；将各个图块分发给工作器服务；通过工作器服务开始执行图块。
- 工作器服务（每项任务一个）：使用适合可用硬件（CPU、GPU等）的内核实现来安排图操作的执行，向其他工作器服务发送操作结果以及从它们那里接收操作结果。
- 内核实现：对单个图操作执行计算。
图2.11说明了这些组件的交互。“/job:worker/task:0”和“/job:ps/task:0”都是工作器服务的任务。PS代表参数服务器（负责存储和更新模型参数的任务）。其他任务会在优化这些参数时向这些参数发送更新。任务之间的这种特殊分工不是必需的，但对于分布式训练来说很常见。
图2.11 计算操作交互
注意，分布式主服务和工作器服务仅存在于分布式TensorFlow中。TensorFlow的单进程版本包含一个特殊的会话实现，它可以执行分布式主服务执行的所有操作，但只与本地进程中的设备进行通信。
以下各部分将更详细地介绍TensorFlow核心层，并逐步说明如何处理示例图。
2.4.2 TensorFlow客户端架构
用户负责编写可构建计算图的客户端TensorFlow程序。该程序可以直接组建各项操作，也可以使用Keras API之类的便利库组建神经网络层和其他更高级别的抽象物。TensorFlow支持多种客户端语言，Python和C++是首先支持的，C++是TensorFlow框架的实现语言，而Python在数据处理工程师群体中有着广泛的使用基础，目前针对Java和JavaScript的API库也日益成熟。随着功能的日益成熟，通常会将它们移植到C++中，以便用户可以通过所有客户端语言访问经过优化的实现。大多数训练库层面仍然只支持Python，主要目的是让科研人员更加专注于机器学习算法的构建。
客户端会创建一个会话，以便将图定义作为tf.GraphDef协议缓冲区发送到分布式主服务。当客户端评估图中的一个或多个节点时，会调用分布式主服务来启动计算。
在图2.12中，客户端构建了一个图，将权重（w）应用于特征向量（x），添加偏差项（b）并将结果保存在变量（s）中。
图2.12 客户端构建图的过程
2.4.3 TensorFlow分布式主服务架构
分布式主服务包含图的低阶操作，包括：对图进行裁剪以获得评估客户端请求的节点所需的子图，划分图以获取每台参与设备的图块，以及缓存这些块以便可以在后续步骤中重复使用。
由于主服务可以看到某一步的总体计算，因此它会应用常见的子表达式消除或常量折叠等标准优化。然后，它会在一组任务中协调经过优化的子图的执行，如图2.13所示。
图2.13 主服务执行计算图开始状态
图2.14显示了示例图的可能划分。分布式主服务已对模型参数进行分组，以便将它们一起放在参数服务器上。
图2.14 模型参数分组的过程
如果图边缘在划分时被切割掉，分布式主服务就会插入发送和接收节点，以便在分布式任务之间传递信息（见图2.15）。
图2.15 任务之间传递信息的过程
然后，分布式主服务会将各个图块传送到分布式任务，如图2.16所示。
图2.16 计算图到达分布式任务工作中心
2.4.4 TensorFlow工作器服务架构
在工作器服务架构中，每项任务中的工作器服务会处理来自主服务的请求，为包含本地子图的操作安排内核的执行时间，以及调解任务之间的直接通信。
我们会优化工作器服务，以便以较低的开销运行大型图。当前的实现每秒可以执行数万个子图，这可让大量副本生成快速、精细的训练步。工作器服务会将各个内核分派给本地设备并尽可能并行运行这些内核，例如通过使用多个CPU核或GPU流。
本地CPU和GPU设备之间的传输使用cudaMemcpyAsync() API来重叠计算和数据传输，如图2.17所示。两个本地GPU之间的传输使用对等DMA，以避免通过主机CPU进行代价高昂的复制。对于任务之间的传输，TensorFlow会使用多种协议，包括：基于TCP的gRPC和基于聚合以太网的RDMA。此外，还提供对NVIDIA的NCCL库的初步支持，以进行多GPU通信。
图2.17 计算图在分布式任务中进行运算和同步
2.4.5 TensorFlow内核架构
内核是TensorFlow真正的能力所在，运行时包含200多个标准操作，包括数学、数组、控制流和状态管理操作。其中，每项操作都具有针对各种设备优化的内核实现。许多操作内核都是使用Eigen::Tensor实现的，后者使用C++模板为多核CPU和GPU生成高效的并行代码。不过，我们可以随意使用cuDNN实现更高效内核的库。另外，还实现了量化，能够在移动设备和高吞吐量数据中心应用等环境中实现更快的推断，并使用GEMMLOWP低精度矩阵库加快量化计算。
如果将子计算表示为操作组合很困难或效率低下，那么用户可以注册提供高效实现（用C++编写）的其他内核。例如，建议为一些对性能要求苛刻的操作注册自己的混合内核，例如ReLU和S型激活函数及其对应的梯度。另外，XLA编译器具有自动内核混合的实验性实现。