深度学习利器： TensorFlow 系统架构及高性能程序设计

2015 年 11 月 9 日谷歌开源了人工智能平台 TensorFlow，同时成为 2015 年最受关注的开源项目之一。经历了从 v0.1 到 v0.12 的 12 个版本迭代后，谷歌于 2017 年 2 月 15 日发布了 TensorFlow 1.0 版本，并同时在美国加州山景城举办了首届 TensorFlow Dev Summit 会议。

TensorFlow 1.0 及 Dev Summit（2017）回顾

和以往版本相比，TensorFlow 1.0 的特性改进主要体现在以下几个方面：

• 速度更快：TensorFlow 1.0 版本采用了 XLA 的编译技术，改进了 TensorFlow 的运行性能及内存利用。从 Benchmark 问题的测试结果来看，对单机 Inception v3 模型，实现了在单机 8 GPUs 上 7.3 倍的运算加速；对分布式 Inception v3 模型，实现了在多机 64 GPUs 上 58 倍的运算加速。
• 更加灵活：该版本除了支持 tf.layers，tf.metrics 及 tf.losses 模型的 High-Level API 外，实现了对 keras（high-level neural networks library）API 的全面兼容。
• 更产品化：TensorFlow Python API 在 v1.0 版本中趋于稳定，为产品兼容性打下坚实基础。

在 TensorFlow 1.0 版本发布的当天，谷歌公司还举办了 TensorFlow 2017 DEV Summit。该日程主要包括以下几个方面的主题演讲：

• XLA (TensorFlow, Compiled) 编译技术 ：介绍采用 XLA 技术最小化图计算执行时间和最大化利用计算资源，用于减少数据训练和模型结果推断时间。
• Hands-on TensorBoard 可视化技术：介绍了如何使用 TensorBoard，以及 TensorFlow 图模型、训练数据的可视化等。
• TensorFlow High-Level API：介绍了使用 Layers, Estimators, and Canned Estimators High-Level API 定义训练模型。
• Integrating Keras & TensorFlow: 介绍了如何在 TensorFlow 中使用 Keras API 进行模型定义及训练。
• TensorFlow at DeepMind：介绍了在 DeepMind 中使用 TensorFlow 平台的典型案例，包括 AlphaGo 等应用。
• Skin Cancer Image Classification：介绍了斯坦福医学院使用 TensorFlow 分类皮肤癌照片，用于医学诊断。
• Mobile and Embedded TensorFlow：介绍了如何把 TensorFlow 模型运行在移动终端、嵌入式设备，包括安卓，iOS 等系统。
• Distributed TensorFlow：系统性地介绍了分布式 TensorFlow 的相关技术，以及如何应用于大规模模型训练。
• TensorFlow Ecosystem：讲解了 TensorFlow 的生态系统，包括生成训练数据，分布式运行 TensorFlow 和 serving models 的产品化流程。
• Serving Models in Production with TensorFlow Serving：系统性讲解了如何在生产环境中应用 TensorFlow Serving 模型。
• ML Toolkit：介绍了 TensorFlow 的机器学习库，如线性回归，KMeans 等算法模型的使用。
• Sequence Models and the RNN API：介绍了如何构建高性能的 sequence-to-sequence 模型，以及相关 API。
• Wide & Deep Learning: 介绍了如何结合 Wide 模型和 Deep 模型构建综合训练模型。
• Magenta，Music and Art Generation：使用增强型深度学习模型生成音乐声音和艺术图片。
• Case Study，TensorFlow in Medicine - Retinal Imaging：使用 TensorFlow 机器学习平台对医学视网膜图片进行分类，辅助医学诊断。

TensorFlow 系统架构

TensorFlow 作为分布式机器学习平台，主要架构如下图所示。RPC 和 RDMA 为网络层，主要负责传递神经网络算法参数。CPU 和 GPU 为设备层，主要负责神经网络算法中具体的运算操作。Kernel 为 TensorFlow 中算法操作的具体实现，如卷积操作，激活操作等。Distributed Master 用于构建子图；切割子图为多个分片，不同的子图分片运行在不同的设备上；Master 还负责分发子图分片到 Executor/Work 端。Executor/Work 在设备（CPUs，GPUs，etc.）上，调度执行子图操作；并负责向其它 Worker 发送和接收图操作的运行结果。C API 把 TensorFlow 分割为前端和后端，前端（Python/C++/Java Client）基于 C API 触发 TensorFlow 后端程序运行。Training libraries 和 Inference libs 是模型训练和推导的库函数，为用户开发应用模型使用。

下图为 Client、Master 及 Worker 的内部工作原理。"/job:worker/task:0" 和 “/job:ps/task:0” 表示 worker 中的执行服务。"job:ps"表示参数服务器，用于存储及更新模型参数。"job:worker"用于优化模型参数，并发参数发送到参数服务器上。Distributed Master 和 Worker Service 只存在于分布式 TensorFlow 中。单机版本的 TensorFlow 实现了 Local 的 Session，通过本地进程的内部通讯实现上述功能。

用户编写 TensorFlow 应用程序生成计算图，Client 组件会创建 Session，并通过序列化技术，发送图定义到 Distributed Master 组件。下图中，Client 创建了一个 s+=w*x+b 的图计算模型。

当 Client 触发 Session 运算的时候，Maser 构建将要运行的子图。并根据设备情况，切割子图为多个分片。下面为 Master 构建的运行子图：

接着切割子图，把模型参数分组在参数服务器上，图计算操作分组在运算 Worker 上。下图为一种可行的图切割策略：

Distributed Master 会根据模型参数的分区情况进行切割边，在 Task 间插入发送和接收 Tensor 信息的通信节点，如下图所示：

接着 Distributed Master 通过 RegisterGraph 方法发送子图分片给 Task，如下图所示：

Master 通过 RunGraph 触发子图运算，Worker 会使用 GPU/CPU 运算设备执行 TensorFlow Kernel 运算。在本节点的 CPU 和 GPU 之间，使用 cudaMemcpyAsync 传输数据；在本节点 GPU 和 GPU 之间，使用 peer-to-peer DMA 传输数据，避免通过 CPU 复制数据。TensorFlow 使用 gRPC（TCP）和 RDMA （Converged Ethernet）技术，实现 Worker 间的数据通信及传输，如下图所示：

高性能程序设计

TensorFlow 内核采用 C/C++ 开发，并提供了 C++，Python，Java，Go 语言的 Client API。特别是 Python API，是目前主流的 TensorFlow 模型开发接口。但为什么还需要采用 C++ API 去训练模型呢？本文基于如下两点考虑，首先当我们采用 Python API 去训练模型的时候，需要不断地用 Python API 调用 C/C++ 底层接口，重复的接口调用一定程度上影响了程序的执行性能。更为重要的是，在 GPU 上训练模型的时候需要大量的内存交换；如果采用 C++ API 去训练模型，可提供更好的运算性能及更好地控制 GPU 内存的分配。

下图为 Python API 的运算架构：在模型训练的每次迭代中，程序通过 Python API 读取 Batch Data，然后通过 TensorFlow Session Run 接口，传递数据给 C++，并触发神经网络训练。如下图所示：

下图为 C++ API 的运算架构：在模型训练的每次迭代中，通过 C++ API 读取 Batch Data 后，直接触发模型训练。减少了不同语言间 API 接口的循环调用及数据传输。如下图所示：

为了采用 C++ API 进行模型训练，我们首先需要编写训练模型，这个编写过程可以采用 Python 语言来完成。我们首先采用 Python API 编写训练模型，然后把图模型转换为 Protobuf 的序列化文件。接着通过 C++ API 加载该模型文件，创建 TensorFlow Session，初始化模型变量，以及加载训练数据并执行神经网络训练。程序架构如下图所示：

下面为使用 Python API 定义训练模型的示例：

with tf.Session() as sess:

#定义 Placeholder Tensor 接入训练数据
   x = tf.placeholder(tf.float32, [None, 32], name="x")
   y = tf.placeholder(tf.float32, [None, 8], name="y")

   #定义训练模型
   w1 = tf.Variable(tf.truncated_normal([32, 16], stddev=0.1))
   b1 = tf.Variable(tf.constant(0.0, shape=[16]))
   w2 = tf.Variable(tf.truncated_normal([16, 8], stddev=0.1))
   b2 = tf.Variable(tf.constant(0.0, shape=[8]))
   a = tf.nn.tanh(tf.nn.bias_add(tf.matmul(x, w1), b1))
   y_out = tf.nn.tanh(tf.nn.bias_add(tf.matmul(a, w2), b2), name="y_out")
   cost = tf.reduce_sum(tf.square(y-y_out), name="cost")
   optimizer = tf.train.AdamOptimizer().minimize(cost, name="train")

   #定义变量初始化操作
   init = tf.initialize_variables(tf.all_variables(), name='init_all_vars_op')

   #把图模型转换为 Protobuf 文件
tf.train.write_graph(sess.graph_def, './', 'mlp.pb', as_text=False)

下面为使用 C++ API 加载 Protobuf 图模型，并执行训练的示例：

#include "tensorflow/core/public/session.h"
#include "tensorflow/core/graph/default_device.h"
using namespace tensorflow;

int main(int argc, char* argv[]) {
   //Protobuf 模型文件名
   std::string graph_definition = "mlp.pb";
   //Tensorflow Sesssion
   Session* session;

   // 定义图模型对象
   GraphDef graph_def;
   SessionOptions opts;

   // 存储 Session 会话的运行结果
   std::vector<Tensor> outputs; 

   #加载 Protobuf 模型文件到图模型对象中
   TF_CHECK_OK(ReadBinaryProto(Env::Default(), graph_definition, &graph_def));

   // 默认在 gpu 0 上执行模型的训练操作
   graph::SetDefaultDevice("/gpu:0", &graph_def);

   // 设定 GPU 显存使用参数
   opts.config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.5);
   opts.config.mutable_gpu_options()->set_allow_growth(true);

   // 创建 TensorFlow 会话
   TF_CHECK_OK(NewSession(opts, &session));

   // 加载图对象到会话中
   TF_CHECK_OK(session->Create(graph_def));

   // 执行模型参数初始化操作
   TF_CHECK_OK(session->Run({}, {}, {"init_all_vars_op"}, nullptr));

   // 定义模型输入数据，包括数据类型和维度信息
   Tensor x(DT_FLOAT, TensorShape({100, 32}));
   Tensor y(DT_FLOAT, TensorShape({100, 8}));

   // 把 Tensor 转换为矩阵，并初始化 Tensor 数据
   auto _XTensor = x.matrix<float>();
   auto _YTensor = y.matrix<float>();
   _XTensor.setRandom();
   _YTensor.setRandom();

   for (int i = 0; i < 10; ++i) {
       // 执行模型的训练操作，{{"x", x}, {"y", y}}表示输入数据 Tensor 名称和 Tensor 对象；{"cost"}表示要获取输出值的操作名称；&outputs 表示执行 "cost" 操作后返回的 Tensor 对象
       TF_CHECK_OK(session->Run({{"x", x}, {"y", y}}, {"cost"}, {}, &outputs)); 

       // 获取执行“cost“操作后的运算结果
       float cost = outputs[0].scalar<float>()(0);
       std::cout << "Cost: " << cost << std::endl;

       // 执行 "train" 操作
       TF_CHECK_OK(session->Run({{"x", x}, {"y", y}}, {}, {"train"}, nullptr)); // Train
       outputs.clear();
   }

   // 关闭 Session 及删除 Session 对象
   session->Close();
   delete session;
   return 0;
}

当 C++ 程序写好后，编译时候需要链接的头文件，开源已经帮我们整理好了，存放于目录 /usr/lib/python2.7/site-packages/tensorflow/include 下。编译和运行的时候需要链接 libtensorflow_cc.so，可以按照下面的方式编译该库文件：bazel build -c opt //tensorflow:libtensorflow_cc.so --copt=-m64 --linkopt=-m64 --spawn_strategy=standalone --genrule_strategy=standalone --verbose_failures。具体可参考 TensorFlow 源代码的官方编译文档。

总结

本文首先回顾了 TensorFlow 1.0 主要新特性及 TensorFlow 2017 Dev Summit 的主要议程。到目前为止 TensorFlow 的 GitHub Star 排名为 51000+， Fork 排名已达 24000+，有 15000+ commits。并随着 TensorFlow 新版本的不断发布以及新特性的不断增加，TensorFlow 使用更加灵活，运行速度更快，使用方式更产品化，已成为目前主流的深度学习平台之一。

接着介绍了 TensorFlow 的系统架构，包括 Client，Master，Worker，Kernel 的相关概念及运行方式，是一种适合大规模分布式训练的机器学习平台。从上述系统架构中可以看到，TensorFlow 内核采用 C/C++ 开发，当采用 Python API 去训练模型的时候，需要不断地用 Python 调用 C/C++ 底层接口，重复的接口调用一定程度上影响了程序的执行性能。如果有最求高性能运算的朋友，可以尝试用下本文高性能运算章节推荐的方法。

参考文献

1. http://www.tensorflow.org
2. 深度学习利器：分布式 TensorFlow 及实例分析
3. 深度学习利器：TensorFlow 使用实战

作者介绍

武维（邮箱：3381209@qq.com），博士，现为 IBM Spectrum Computing 研发工程师。主要从事大数据，深度学习，云计算等领域的研发工作。

感谢杜小芳对本文的审校。

给InfoQ 中文站投稿或者参与内容翻译工作，请邮件至 editors@cn.infoq.com 。也欢迎大家通过新浪微博（ @InfoQ ， @丁晓昀），微信（微信号： InfoQChina ）关注我们。