全面讲解TensorFlow架构与设计

　　TensorFlow，以下简称TF，是Google去年发布的机器学习平台，发布以后由于其速度快，扩展性好，推广速度还是蛮快的。江湖上流传着Google的大战略，Android占领了移动端，TF占领神经网络提供AI服务，未来的趋势恰好是语音图像以及AI的时代，而GoogleIO上发布的Gbot似乎正是这一交叉领域的初步尝试。我们今天探讨一下TensorFlow架构与设计，有需要的小伙伴，参考一下。

     TensorFlow是什么?

　　TensorFlow基于数据流图，用于大规模分布式数值计算的开源框架。节点表示某种抽象的计算，边表示节点之间相互联系的张量。

　　计算图实例

　　TensorFlow支持各种异构的平台，支持多CPU/GPU，服务器，移动设备，具有良好的跨平台的特性;TensorFlow架构灵活，能够支持各种网络模型，具有良好的通用性;此外，TensorFlow架构具有良好的可扩展性，对OP的扩展支持，Kernel特化方面表现出众。

　　TensorFlow最初由Google大脑的研究员和工程师开发出来，用于机器学习和神经网络方面的研究，于2015.10宣布开源，在众多深度学习框架中脱颖而出，在Github上获得了最多的Star量。

　　本文将阐述TensorFlow的系统架构，帮助读者加深理解TensorFlow的工作机理。

　　本文假设读者已经了解TensorFlow的基本编程模型，包括计算图, OP, Tensor, Session等基本概念。

　　系统概述

　　TensorFlow的系统结构以C API为界，将整个系统分为「前端」和「后端」两个子系统：

　　前端系统：提供编程模型，负责构造计算图;

　　后端系统：提供运行时环境，负责执行计算图。

　　TensorFlow系统架构

　　如上图所示，重点关注系统中如下4个基本组件，它们是系统分布式运行机制的核心。

　　Client

　　Client是前端系统的主要组成部分，它是一个支持多语言的编程环境。它提供基于计算图的编程模型，方便用户构造各种复杂的计算图，实现各种形式的模型设计。

　　Client通过Session为桥梁，连接TensorFlow后端的「运行时」，并启动计算图的执行过程。

　　Distributed Master

　　在分布式的运行时环境中，Distributed Master根据Session.run的Fetching参数，从计算图中反向遍历，找到所依赖的「最小子图」。

　　然后，Distributed Master负责将该「子图」再次分裂为多个「子图片段」，以便在不同的进程和设备上运行这些「子图片段」。

　　最后，Distributed Master将这些「子图片段」派发给Work Service;随后Work Service启动「子图片段」的执行过程。

　　Worker Service

　　对于每以个任务，TensorFlow都将启动一个Worker Service。Worker Service将按照计算图中节点之间的依赖关系，根据当前的可用的硬件环境(GPU/CPU)，调用OP的Kernel实现完成OP的运算(一种典型的多态实现技术)。

　　另外，Worker Service还要负责将OP运算的结果发送到其他的Work Service;或者接受来自其他Worker Service发送给它的OP运算的结果。

　　Kernel Implements

　　Kernel是OP在某种硬件设备的特定实现，它负责执行OP的运算。

　　组件交互

　　组件交互

　　如上图所示，假设存在两个任务：

　　/job:ps/task:0: 负责模型参数的存储和更新

　　/job:worker/task:0: 负责模型的训练或推理

　　接下来，我们将进一步抽丝剥茧，逐渐挖掘出TensorFlow计算图的运行机制。

　　客户端

　　Client基于TensorFlow的编程接口，构造计算图。目前，TensorFlow主流支持Python和C++的编程接口，并对其他编程语言接口的支持日益完善。

　　此时，TensorFlow并未执行任何计算。直至建立Session会话，并以Session为桥梁，建立Client与后端运行时的通道，将Protobuf格式的GraphDef发送至Distributed Master。

　　也就是说，当Client对OP结果进行求值时，将触发Distributed Master的计算图的执行过程。

　　如下图所示，Client构建了一个简单计算图。它首先将w与x进行矩阵相乘，再与截距b按位相加，最后更新至s。

　　构造计算图

　　Distributed Master

　　在分布式的运行时环境中，Distributed Master根据Session.run的Fetching参数，从计算图中反向遍历，找到所依赖的最小子图。

　　然后Distributed Master负责将该子图再次分裂为多个「子图片段」，以便在不同的进程和设备上运行这些「子图片段」。

　　最后，Distributed Master将这些图片段派发给Work Service。随后Work Service启动「本地子图」的执行过程。

　　Distributed Master将会缓存「子图片段」，以便后续执行过程重复使用这些「子图片段」，避免重复计算。

　　执行图计算

　　如上图所示，Distributed Master开始执行计算子图。在执行之前，Distributed Master会实施一系列优化技术，例如「公共表达式消除」，「常量折叠」等。随后，Distributed Master负责任务集的协同，执行优化后的计算子图。

　　子图片段

　　子图片段

　　如上图所示，存在一种合理的「子图片段」划分算法。Distributed Master将模型参数相关的OP进行分组，并放置在PS任务上。其他OP则划分为另外一组，放置在Worker任务上执行。

　　SEND/RECV节点

　　插入SEND/RECV节点

　　如上图所示，如果计算图的边被任务节点分割，Distributed Master将负责将该边进行分裂，在两个分布式任务之间插入SEND和RECV节点，实现数据的传递。

　　随后，Distributed Master将「子图片段」派发给相应的任务中执行，在Worker Service成为「本地子图」，它负责执行该子图的上的OP。

　　Worker Service

　　对于每个任务，都将存在相应的Worker Service，它主要负责如下3个方面的职责：

　　处理来自Master的请求;

　　调度OP的Kernel实现，执行本地子图;

　　协同任务之间的数据通信。

　　执行本地子图

　　Worker Service派发OP到本地设备，执行Kernel的特定。它将尽最大可能地利用多CPU/GPU的处理能力，并发地执行Kernel实现。

　　另外，TensorFlow根据设备类型，对于设备间的SEND/RECV节点进行特化实现：

　　使用cudaMemcpyAsync的API实现本地CPU与GPU设备的数据传输;

　　对于本地的GPU之间则使用端到端的DMA，避免了跨host CPU昂贵的拷贝过程。

　　对于任务之间的数据传递，TensorFlow支持多协议，主要包括：

　　gRPC over TCP

　　RDMA over Converged Ethernet

　　Kernel Implements

　　TensorFlow的运行时包含200多个标准的OP，包括数值计算，多维数组操作，控制流，状态管理等。每一个OP根据设备类型都会存在一个优化了的Kernel实现。在运行时，运行时根据本地设备的类型，为OP选择特定的Kernel实现，完成该OP的计算。

　　TensorFlow Core

　　其中，大多数Kernel基于Eigen::Tensor实现。Eigen::Tensor是一个使用C++模板技术，为多核CPU/GPU生成高效的并发代码。但是，TensorFlow也可以灵活地直接使用cuDNN实现更高效的Kernel。

　　决定设备

　　使用一个cost model算法来进行预估时间，计算后使用贪心算法来分配设备。在决定设备的时候，也可以预先设置一些约束，比如，某个op只能在GPU上执行等。

　　预估时间有两种方法：

　　使用启发式的算法，通过把输入和输出的类型以及tensor的大小输入进去，得到时间的预估

　　使用模拟的方法，对图的计算进行一个模拟，得到各个计算在其可用的设备上的时间。

　　寻找合适设备是TF系统区分与之前很多系统的地方，之前的系统比如Parameter Server，是参数分离出来，运算在一起，同时使用数据切分来达到分布式。而TF是把每个op都映射到某个机器上，意味着每个op可能在不同的机器上，这是对系统的进一步剖离，因而可以达到更高的可扩展性。

　　跨设备通信

　　当两个需要通信的op在不同的机器上时，就需要跨设备通信，当它们需要通信时，TF会在它们之间的联系中添加Send和Recv节点，通过Send和Recv之间进行通信来达到op之间通信的效果。如下所示：

　　为了优化网络通信，TF会将相同的数据传送合并，如a->b和a->c的传送合并，这一点可以通过Send和Recv很方便的实现。而通过实现Send和Recv，将master节点的通信调度任务解放出来，master就只需要向图中的各个节点发出运行命令就够了，增加了系统的可扩展性。

　　Send和Recv通过TCP或RDMA来传输数据

　　错误处理

　　在分布式系统中，常见的错误来自于两个方面：

　　Send/Recv的网络传输

　　master和worker process之间的心跳同步

　　当错误发生的时候，TF会将整个图的计算停止，并从上一次保存的状态重新执行。为了保存状态，每个Variable的节点都去连接一个Save的节点。这些save节点会每隔一段时间或每隔几次迭代运行一次。

　　此外，TensorFlow实现了矢量化技术，使得在移动设备，及其满足高吞吐量，以数据为中心的应用需求，实现更高效的推理。

　　如果对于复合OP的子计算过程很难表示，或执行效率低下，TensorFlow甚至支持更高效的Kernle实现的注册，其扩展性表现相当优越。

　　技术栈

　　最后，按照TensorFlow的软件层次，通过一张表格罗列TensorFlow的技术栈，以便更清晰地对上述内容做一个简单回顾。

　　TensorFlow技术栈

     结束语：看完文章的小伙伴，估计都了解“TensorFlow架构与设计”差不多了吧！如果还想了解更多关于这方面的知识，可以登陆课课家浏览，这里面有全面的知识内容还有视频教学，随时欢迎你的到来哟~