TensorFlow基础剖析

一．概述

TensorFlow 是一个使用数据流图 (Dataflow Graph) 表达数值计算的开源软件库。它使 用节点表示抽象的数学计算，并使用 OP 表达计算的逻辑；而边表示节点间传递的数据流， 并使用 Tensor 表达数据的表示。数据流图是一种有向无环图 (DAG)，当图中的 OP 按 照特定的拓扑排序依次被执行时，Tensor 在图中流动形成数据流，TensorFlow 因此而得名。

在分布式运行时，数据流图的被分裂为多个子图，并被有效地部署到集群中的多个机 器上并发执行。在一个机器内，注册的子图被二次分裂为更小的子图，它们被部署在本地 设备集上并发执行。TensorFlow 支持多种异构设备的分布式计算，包括 CPU, GPU, ASIC。 TensorFlow 跨平台的卓越表现，使得它能够灵活地部署在各种计算平台上，包括台式机、 服务器、移动终端。

TensorFlow 最初由 Google Brain 的研究员和工程师们开发出来，用于开展机器学习和 深度神经网络的研究，包括语言识别、计算机视觉、自然语言理解、机器人、信息检索。但 是，TensorFlow 系统架构的通用性和灵活性，使其广泛地用于其他科学领域的数值计算。

1.1 前世今生

Google Brain 项目始于 2011 年，用于研究超大规模的深度神经网络。在项目早期阶段， Google Brain 构建了第一代分布式深度学习框架 DistBelief，并在 Google 内部的产品中得 到了大量的应用。

基于 DistBelief 的经验，Google Brain 对深度学习训练和推理的需求，及其深度学习框 架的系统行为和属性，有了更全面更深刻地理解，并于 2015.11 重磅推出第二代分布式深度 学习框架 TensorFlow。TensorFlow 作为 DistBelief 的后继者，革命性地对既有系统架构做 了全新设计和实现，TensorFlow 一经发布，便在深度学习领域一鸣惊人，在社区中形成了 巨大的影响力。

DistBelief

DistBelief 是一个用于训练大规模神经网络的的分布式系统，是 Google 第一代分布式 机器学习框架。自 2011 年以来，在 Google 内部大量使用 DistBelief 训练大规模的神经网 络，并广泛地用于机器学习和深度学习领域的研究和应用，包括非监督学习、语言表示、图 像分类、目标检测、视频分类、语音识别、序列预测、行人检测、强化学习等。

编程模型

DistBelief 的编程模型是基于层的 DAG 图。层可以看做是一种组合多个运算操作符的 复合运算符，它完成特定的计算任务。例如，全连接层完成 f(WTx + b) 的复合计算，包括 输入与权重的矩阵乘法，随后再与偏置相加，最后在线性加权值的基础上应用激活函数，实 施非线性变换。

架构

DistBelief 使用参数服务器 (Parameter Server, 常称为 PS) 的系统架构，训练作业包括 两个分离的进程：无状态的 Worker 进程，用于模型的训练；有状态的 PS 进程，用于维护 模型的参数。如图??（第??页）所示，在分布式训练过程中，各个模型副本异步地从 PS 上 拉取训练参数 w，当完成一步迭代运算后，推送参数的梯度 ∆w 到 PS 上去，并完成参数的 更新。

图 1-1 DistBelief: Parameter Server 架构

缺陷

但是，对于深度学习领域的高级用户，DistBelief 的编程模型，及其基于 PS 的系统架 构，缺乏足够的灵活性和可扩展性。

1. 优化算法：添加新的优化算法，必须修改 PS 的实现；get(), put() 的抽象方法， 对某些优化算法并不高效。

2. 训练算法：支持非前馈的神经网络面临巨大的挑战性；例如，包含循环的 RNN， 交替训练的对抗网络，及其损失函数由分离的代理完成的增强学习模型。

3. 加速设备：DistBelief 设计之初仅支持多核 CPU，并不支持多卡的 GPU，遗留的 系统架构对支持新的计算设备缺乏良好的可扩展性。

TensorFlow

DistBelief 遗留的架构和设计，不再满足深度学习与日俱增的需求变化，Google 毅然放 弃了既有的 DistBelief 实现，并决定在其基础上做全新的系统架构设计，TensorFlow 应运 而生，开创了深度学习领域的新纪元。

编程模型

TensorFlow 使用数据流图表达计算过程和共享状态，使用节点表示抽象计算，使用边 表示数据流。如图??（第??页）所示，展示了 MNIST 手写识别应用的数据流图。在该模型 中，前向子图使用了 2 层全连接网络，分别为 ReLU 层和 Softmax 层。随后，使用 SGD 的 优化算法，构建了与前向子图对应的反向子图，用于计算训练参数的梯度。最后，根据参数 更新法则，构造训练参数的更新子图，完成训练参数的迭代更新。

图 1-2 TensorFlow: 数据流图

设计原则

TensorFlow 的系统架构遵循了一些基本的设计原则，用于指导 TensorFlow 的系统实 现。

1.       延迟计算：图的构造与执行分离，并推迟计算图的执行过程；

2. 原子 OP：OP 是最小的抽象计算单元，支持构造复杂的网络模型；

3. 抽象设备：支持 CPU, GPU, ASIC 多种异构计算设备类型；

4. 抽象任务：基于任务的 PS，对新的优化算法和网络模型具有良好的可扩展性。

优势

相对于其他机器学习框架，TensorFlow 具有如下方面的优势。

1. 高性能：TensorFlow 升级至 1.0 版本性能提升显著，单机多卡 (8 卡 GPU) 环境 中，Inception v3 的训练实现了 7.3 倍的加速比；在分布式多机多卡 (64 卡 GPU) 环境 中，Inception v3 的训练实现了 58 倍的加速比；

2. 跨平台：支持多 CPU/GPU/ASIC 多种异构设备的运算；支持台式机，服务器， 移动设备等多种计算平台；支持 Windows，Linux，MacOS 等多种操作系统；

3. 分布式：支持本地和分布式的模型训练和推理；

4. 多语言：支持 Python, C++, Java, Go 等多种程序设计语言；

5. 通用性：支持各种复杂的网络模型的设计和实现，包括非前馈型神经网络；

6. 可扩展：支持 OP 扩展，Kernel 扩展，Device 扩展，通信协议的扩展；

7. 可视化：使用 TensorBoard 可视化整个训练过程，极大地降低了 TensorFlow 的 调试过程；

8. 自动微分：TensorFlow 自动构造反向的计算子图，完成训练参数的梯度计算；

9. 工作流：TensorFlow 与 TensorFlow Serving 无缝集成，支持模型的训练、导入、 导出、发布一站式的工作流，并自动实现模型的热更新和版本管理。

1.2 社区发展 TensorFlow 是目前炙手可热的深度学习框架。自开源以来，TensorFlow 社区相当活 跃。在 Github 上收获了来自众多的非 Google 员工的数万次代码提交，并且每周拥有近百 个 Issue 被提交。在 Stack Overflow 上也拥有上万个关于 TensorFlow 的问题被提问和回答。 在各种技术大会上，TensorFlow 也是一颗闪亮的明星，得到众多开发者的青睐。

开源 2015.11，Google Research 发布文章：TensorFlow: Google's latest machine learning system, open sourced for everyone，正式宣布新一代机器学习系统 TensorFlow 开源。随 后，TensorFlow在 Github上代码仓库短时间内获得了大量的 Star和 Fork。如图??（第??页） 所示，TensorFlow 的社区活跃度已远远超过其他竞争对手，成为目前最为流行的深度学习 框架。

毫无疑问，TensorFlow 的开源对学术界和工业界产生了巨大的影响，它极大地降低 了深度学习在各个行业中应用的难度。众多的学者、工程师、企业、组织纷纷地投入到了 TensorFlow 社区，并一起完善和改进 TensorFlow，推动其不断地向前演进和发展。

里程碑 TensorFlow 自 2015.11 开源以来，平均一个多月发布一个版本。如图??（第??页）所 示，展示了 TensorFlow 几个重要特性的发布时间。

工业应用 TensorFlow 自开源发展两年以来，在生产环境中被大量应用使用。在医疗方面，帮助 医生预测皮肤癌发生的概率；在音乐、绘画领域，帮助人类更好地理解艺术；在移动端，多 款移动设备搭载 TensorFlow 训练的机器学习模型，用于翻译工作。TensorFlow 在 Google 内部应用的增长也十分迅速，多个重量级产品都有相关应用，包括：Google Search, Google Gmail, Google Translate, Google Maps 等，如图??（第??页）所示。

二．编程环境

2.1 代码结构

克隆源码 首先，从 Github 上克隆 TensorFlow 的源代码。

$ git clone git@github.com:tensorflow/tensorflow.git

然后，切换到最新的稳定分支上。例如，r1.4 分支。

$ cd tensorflow $ git checkout r1.4

源码结构 运行如下命令，打印出 TensorFlow 源码的组织结构。

$ tree -d -L 1 ./tensorflow

其中，本书将重点关注 core, python 组件，部分涉及 c, cc, stream_executor 组件。

示例代码 2-1 TensorFlow 源码结构

截止当前最新发布的 1.4 版本，TensorFlow 代码库拥有大约 100 万代码。其中，包括 53 万行 C/C++ 代码，37 万行 Python 代码，而且代码规模在不断膨胀之中。其中，Python 提供的 API 是最完善的；相比之下，其他编程语言的 API 尚未成熟，甚至处于起步阶段。

Core

内核的源码结构如下所示，主要包括平台，实用函数库，基础框架，Protobuf 定义，本
地运行时，分布式运行时，图操作，OP 定义，以及 Kernel 实现等组成，这是本书重点剖
析的组件之一，将重点挖掘基础框架中隐藏的领域模型，追踪整个运行时环境的生命周期
和图操作的详细过程，并揭示常见 OP 的 Kernel 实现原理和运行机制。

其中，core 主要由 C++ 实现，大约拥有 26 万行代码。

Python

Python 定义和实现了 TensorFlow 的编程模型，并对外开放
API 供程序员使用，其源
码结构如下所示，这也是本书重点剖析的部分。

其中，该组件由 Python 实现，大约有 18 万行代码。

Contrib

contrib 是第三方贡献的编程库，它也是
TensorFlow 标准化之前的实验性编程接口，
犹如 Boost 社区与 C++ 标准之间的关系。当 contrib 的接口成熟后，便会被 TensorFlow 标准化，并从 contrib 中搬迁至 core, python 中，并正式对外发布。

由于 TensorFlow 社区相当活跃，contrib 的变更相当频繁，截止 1.4 版本，大约有 23 万行代码，主要由 Python 设计和实现的编程接口，部分运行时由 C++ 实现。

StreamExecutor

StreamExecutor 是 Google 另一个开源组件库，它提供了主机端 (host-side) 的编程模 型和运行时环境，实现了 CUDA 和 OpenCL 的统一封装。使得在主机端的代码中，可以将 Kernel 函数无缝地部署在 CUDA 或 OpenCL 的计算设备上执行。 目前，StreamExecutor 被大量应用于 Google 内部 GPGPU 应用程序的运行时。其中， TensorFlow 运行时也包含了一个 StreamExecutor 的快照版本，用于封装 CUDA 和 OpenCL 的运行时。本书将简单介绍 CUDA 的编程模型和线程模型，并详细介绍 StreamExecutor 的系统架构与工作原理，揭示 Kernel 函数的实现模式和习惯用法。

Compiler
众所周知，灵活性是 TensorFlow 最重要的设计目标和核心优势，因此 TensorFlow 的 系统架构具有良好的可扩展性。TensorFlow 可用于定义任意图结构，并使用异构的计算 设备有效地执行。但是，熊掌与鱼翅不可兼得，当低级 OP 组合为计算子图时，并期望在 GPU 上有效执行时，运行时将启动更多的 Kernel 的运算。 因此，TensorFlow 分解和组合 OP 的方法，在运行时并不能保证以最有效的方式运行。 此时，XLA 技术孕育而生，它使用 JIT 编译技术来分析运行时的计算图，它将多个 OP 融 合在一起，并生成更高效的本地机器代码，提升计算图的执行效率。

XLA 技术目前处于初级的研发阶段，是 TensorFlow 社区较为活跃研究方向，截止目前 代码规模大约为 12.5 万行，主要使用 C++ 实现。

2.2 工程构建

在开始之前，尝试 TensorFlow 源码的构建过程，了解 TensorFlow 的基本构建方式、工 具，及其依赖的组件库、第三方工具包，对于理解 TensorFlow 工作原理具有很大的帮助。 但是，因篇幅受限，本章仅以 Mac OS 系统为例，讲述 TensorFlow 的源码编译、安装、及 其验证过程。其他操作系统，请查阅 TensorFlow 发布的官方文档。
环境准备

TensorFlow 的前端是一个支持多语言的编程接口。因此，编译 TensorFlow 源代码之 前，需要事先安装相关的编译器、解释器、及其运行时环境。例如，使用 Python 作为编 程接口，需要事先安装 Python 解释器。其次，在构建系统之前，也需要事先安装 GCC 或 Clang 等 C++ 编译器，用于编译后端系统实现。因为 TensorFlow 使用 C++11 语法实现， 所以要保证安装 C++ 编译器要支持 C++11。另外，TensorFlow 使用 Bazel 的构建工具， 可以将其视为更高抽象的 Make 工具。不幸的是，Bazel 使用 Java8 实现，其依赖于 JDK。 因此在安装 Bazel 之前，还得需要事先安装 1.8 及以上版本的 JDK。
安装 JDK
推荐从 Oracle 官网上下载 1.8 版本的 JDK，然后创建相关的环境变量，并将其添加到 ～/.bashrc 的配置文件中。

$ echo 'export JAVA_HOME=$(/usr/libexec/java_home)' >> ~/.bashrc

$ echo 'export PATH="$JAVA_HOME/bin:$PATH"' >> ~/.bashrc

安装 Bazel
在 Mac OS 上，可以使用 brew 安装 Bazel。
$ brew install bazel
如果系统未安装 brew，可以执行如下命令先安装 brew。当然，安装 brew 需要事先安 装 Ruby 解释器，在此不再冗述。
$ ruby -e "$(curl -fsSL \ https://raw.githubusercontent.com/Homebrew/install/master/install)"

安装 Swig

TensorFlow 使用 Swig 构建多语言的编程环境，自动生成相关编程语言的包装器。因 此，在构建之前需要事先安装 Swig 的工具包。

$ brew install swig

安装 Python 依赖包

使用 pip 安装 TensorFlow 所依赖的 Python 包。

$ sudo pip install six numpy wheel autograd

如果系统未安装 pip，则可以使用 brew 预先安装 pip。

$ brew install pip

安装 CUDA 工具包

当系统安装了 CUDA 计算兼容性大于或等于 3.0 的 GPU 卡时，则需要安装 CUDA 工 具包，及其 cuDNN，实现 TensorFlow 运行时的 GPU 加速。推荐从 NVIDIA 官网上下载

2.2 工程构建 17

CUDA Toolkit 8 及以上版本，并安装到系统中，配置相关环境变量。

$ echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.bashrc $ echo 'export LD_LIBRARY_PATH=$CUDA_HOME/lib:$LD_LIBRARY_PATH' >> ~/.bashrc

然后，再下载 cuDNN 5.1 及以上版本，并将其解压至 CUDA_HOME 相关系统目录中。

$ sudo tar -xvf cudnn-8.0-macos-x64-v5.1.tgz -C /usr/local

配置 至此，编译环境一切就绪，执行./configure 配置 TensorFlow 的编译环境了。当系统 支持 GPU，则需要配置相关的 CUDA/cuDNN 编译环境。

$ ./configure

构建 当配置成功后，使用 Bazel 启动 TensorFlow 的编译。在编译启动之前，会尝试从代码 仓库中下载相关依赖库的源代码，包括 gRPC, Protobuf, Eigen 等，并自动完成编译。

$ bazel build --config=opt //tensorflow/tools/pip_package:build_pip_package

当支持 GPU 计算时，添加--config=cuda 编译选项。

$ bazel build -c opt --config=cuda \ //tensorflow/tools/pip_package:build_pip_package

编译成功后，便可以构建 TensorFlow 的 Wheel 包。

$ bazel-bin/tensorflow/tools/pip_package/build_pip_package \ /tmp/tensorflow_pkg

安装 当 Wheel 包构建成功后，使用 pip 安装 TensorFlow 到系统中。

创建 Eclipse 工程
创建一个 Eclipse C++ 工程，如图??（第??页）所示。确定唯一的项目名称，手动地指 定 TensorFlow 源代码的根目录，并选择 Makefile 的空工程。然后，按照 Properties > C/C ++ General > Paths and Symbols > Includes 配置头文件的搜索目录。

2.3 代码生成 在构建 TensorFlow 系统时，Bazel 或 CMake 会自动生成部分源代码。理解代码生成器 的输出结果，可以加深理解系统的行为模式。
2.4 技术栈 如图??（第??页）所示，按照系统的层次结构展示了 TensorFlow 的技术栈，构成了 TensorFlow 生态系统的核心。