Workspace of LionHeart

Focus on Spark/TensorFlow and other BigData and ML platforms

Tensorflow架构介绍

导航[-]

Tensorflow的架构为大规模分布式训练和预测设计，单页为实验新的机器学习模型和进行系统层面上的优化提供了足够多的灵活性。

本文介绍Tensorflow的系统架构设计，展示Tensorflow如何将灵活性与扩展性结合。

本文内容主要从Tensorflow官方文档中翻译整理而得。

阅读本文之前需要了解Tenflow变成中像计算流图、算子和session的概念。如果不了解可以参考如下文档：

概述

Tensorflow的运行时是一个跨平台的库，下图展示了整体结构。用户级别的代码与核心的运行时的代码通过一套C API分隔开。

本文主要介绍如下几层：

  • Client:
    • 将计算定义为数据流图的形式。
    • 通过一个session初始化图的执行。
  • Distributed Master
    • 根据Session.run()传递的参数将图裁剪为子图。
    • 将子图拆解为多个部分用于分发到不通的进程和设备上。
    • 分发子图分片到Worker Services
    • 初始化Woker Service待执行的子图分片计算。
  • Worker Services (每个task一个)
    • 调度子图算子的执行，这些算子会调用对应于可用硬件（CPU\GPU等）的kernal实现。
    • 发送算子计算结果以及接收来自其他Worker Service的计算结果。
  • Kernel Implementations
    • 负责实现每个图算子的执行。

下图说明了这些组件是如何交互的。

解释如下：

  • “/job:worker/task:0” 和 “/job:ps/task:0” 都是worker services上的task。
  • “PS”代表parameter server，是负责存储和更新模型参数的task。其他的task会将训练过程中优化的参数更新发送给parameter server。
  • 这样的职能划分并非必须的，但在分布式训练场景中很常见。

注意

  • Distributed Master和Worker Service只在支持分布式版本的tensorflow中存在。
  • 单进程版本的Tensorflow包含一个特殊的Session实现完成Distributed Master做的全部事情但只在本地进程中进行设备间的通信。

组件介绍

下面以一个示例图的处理过程介绍TensorFlow核心层的更多实现细节。

Client

用户编写客户端tensorflow程序来构建计算流图。

客户端程序既可以直接调用各种算子操作组合而成，也可以调用更方便的API库如Estimator等直接构建神经网络和其他高层抽象实现。

Tensorflow支持多种客户端语言，主流的是Python和C++，因为Google的内部使用者对这两种语言更熟悉一些。

大多数训练库仍为Python实现，但C++更高效。

Client会创建一个session，将图的定义以tf.GraphDef这种protobuffer的形式发送给Distributed Master节点。

当client对流图中的一个或多个节点进行求值时，会触发调用让Distributed Master进行计算的初始化操作。

在下图中，Client构建了一个图将权值w附加到一个特征向量x并增加了偏移b，然后在变量中存储计算结果。

tf.Session相关代码：

Distributed Master

Distributed Master负责以下工作：

  • 对图进行裁剪来获取满足Client对节点求值必须计算的子图。
  • 将图进行拆分得到每个参与计算的设备负责的图分片。
  • 缓存这些拆分的图分片，以备后续步骤中被重新使用。

master节点能够看到每步计算的整体情况，因此可以对计算流程进行一些标准的优化，例如公共子表达式求值、常量折叠等。

随后master会负责协调一系列的task执行来计算优化后的子图。

下图展示了示例graph的一种可能的分发方案。Distributed Master将模型参数进行分组从而在Parameter Server中将他们存储在一起。

在图的边被切分的分片中，Distributed Master会插入send和receive节点来在不同的task间传递信息。

随后Distributed Master将图的分片分发给不同节点上的task。

相关代码：

Worker Service

每个task的Worker Service负责如下工作：

  • 处理Distributed Master发来的请求。
  • 调度执行构成本地局部子图的算子的kernels实现。
  • 协调task之间的直接通信。

我们优化了Worker Service，支持以较低的开销运行大规模的流图。

我们目前的实现能够支持每秒数以万计的子图执行，这通常能够支持大量副本细粒度地进行训练。

Worker Service将算子的kernel实现分发给设备并尽可能并行化，例如使用多CPU核或者GPU。

我们针对每对源设备和目的设备的组合情况，定制了相应的Send和Recv算子。包括：

  • 本地CPU和GPU设备的数据交换使用cudaMemcpyAsync() API来配合计算和数据转移工作。
  • 两个本地GPU使用端到端的DMA技术来避免通过本机CPU进行代价昂贵的数据拷贝操作。

对于task之间的数据交换，tensorflow支持多种协议，包括：

  • 基于TCP协议的gRPC。
  • 基于融合以太网的RDMA。

另外我们也初步支持了NVIDIA为多GPU通信提供的NCCL库。

相关代码：

Kernel

Tensorflow运行时包含了超过200种标准的算子实现，包括数学相关、数组相关、控制流相关以及状态管理相关的算子。

每一种算子在实现上都针对不同的设备种类进行了优化。

其中，许多算子是使用Eigen::Tensor实现的，它基于C++模板生成了高效的能在多核CPU以及GPU上执行的代码。

另一方面，我们也使用想cuDNN等第三方库尽可能地让kernel的实现更加高效。

如果现有算子无法满足需求，用户可以注册自定义的算子实现到kernel中。

相关代码：