Tensorflow代码解析

推荐会员: zhouyf 所属分类: 技术教程,行业精选 发布时间: 2017-03-22 18:44
摘要

2015年11月9日,Google发布深度学习框架TensorFlow并宣布开源,并迅速得到广泛关注,在图形分类、音频处理、推荐系统和自然语言处理等场景下都被大面积推广。TensorFlow系统更新快速,官方文档教程齐全,上手快速且简单易用,支持Python和C++接口。本文依据对Tensorflow(简称TF)白皮书[1]、TF Github[2]和TF官方教程[3]的理解,从系统和代码实现角度讲解TF的内部实现原理。以Tensorflow r0.8.0为基础,本文由浅入深的阐述Tensor和Flow的概念。先介绍了TensorFlow的核心概念和基本概述,然后剖析了OpKernels模块、Graph模块、Session模块。

1. TF系统架构

1.1 TF依赖视图

TF的依赖视图如图1所示[4],描述了TF的上下游关系链。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图 1 TensorFlow依赖视图

TF托管在github平台,有google groups和contributors共同维护。

TF提供了丰富的深度学习相关的API,支持Python和C/C++接口。

TF提供了可视化分析工具Tensorboard,方便分析和调整模型。

TF支持Linux平台,Windows平台,Mac平台,甚至手机移动设备等各种平台。

1.2 TF系统架构

图2是TF的系统架构,从底向上分为设备管理和通信层、数据操作层、图计算层、API接口层、应用层。其中设备管理和通信层、数据操作层、图计算层是TF的核心层。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图2 TF系统架构

底层设备通信层负责网络通信和设备管理。设备管理可以实现TF设备异构的特性,支持CPU、GPU、Mobile等不同设备。网络通信依赖gRPC通信协议实现不同设备间的数据传输和更新。

第二层是Tensor的OpKernels实现。这些OpKernels以Tensor为处理对象,依赖网络通信和设备内存分配,实现了各种Tensor操作或计算。Opkernels不仅包含MatMul等计算操作,还包含Queue等非计算操作,这些将在第5章Kernels模块详细介绍。

第三层是图计算层(Graph),包含本地计算流图和分布式计算流图的实现。Graph模块包含Graph的创建、编译、优化和执行等部分,Graph中每个节点都是OpKernels类型表示。关于图计算将在第6章Graph模块详细介绍。

第四层是API接口层。Tensor C API是对TF功能模块的接口封装,便于其他语言平台调用。

第四层以上是应用层。不同编程语言在应用层通过API接口层调用TF核心功能实现相关实验和应用。

1.3 TF代码目录组织

图3是TF的代码结构视图,下面将简单介绍TF的目录组织结构。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图3 TF代码目录组织结构

Tensorflow/core目录包含了TF核心模块代码。

  • public: API接口头文件目录,用于外部接口调用的API定义,主要是session.h 和tensor_c_api.h。
  • client: API接口实现文件目录。
  • platform: OS系统相关接口文件,如file system, env等。
  • protobuf: 均为.proto文件,用于数据传输时的结构序列化.
  • common_runtime: 公共运行库,包含session, executor, threadpool, rendezvous, memory管理, 设备分配算法等。
  • distributed_runtime: 分布式执行模块,如rpc session, rpc master, rpc worker, graph manager。
  • framework: 包含基础功能模块,如log, memory, tensor
  • graph: 计算流图相关操作,如construct, partition, optimize, execute等
  • kernels: 核心Op,如matmul, conv2d, argmax, batch_norm等
  • lib: 公共基础库,如gif、gtl(google模板库)、hash、histogram等。
  • ops: 基本ops运算,ops梯度运算,io相关的ops,控制流和数据流操作
  • Tensorflow/stream_executor目录是并行计算框架,由google stream executor团队开发。
  •  Tensorflow/contrib目录是contributor开发目录。
  • Tensroflow/python目录是python API客户端脚本。
  • Tensorflow/tensorboard目录是可视化分析工具,不仅可以模型可视化,还可以监控模型参数变化。
  • third_party目录是TF第三方依赖库。
  • eigen3: eigen矩阵运算库,TF基础ops调用
  • gpus: 封装了cuda/cudnn编程库

2. TF核心概念

TF的核心是围绕Graph展开的,简而言之,就是Tensor沿着Graph传递闭包完成Flow的过程。所以在介绍Graph之前需要讲述一下符号编程、计算流图、梯度计算、控制流的概念。

2.1 Tensor

在数学上,Matrix表示二维线性映射,Tensor表示多维线性映射,Tensor是对Matrix的泛化,可以表示1-dim、2-dim、N-dim的高维空间。图4对比了矩阵乘法(Matrix Product)和张量积(Tensor Contract),可以看出Tensor的泛化能力,其中张量积运算在TF的MatMul和Conv2D运算中都有用到。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图4 Tensor contract

Tensor在高维空间数学运算比Matrix计算复杂,计算量也非常大,加速张量并行运算是TF优先考虑的问题,如add, contract, slice, reshape, reduce, shuffle等运算。

TF中Tensor的维数描述为阶,数值是0阶,向量是1阶,矩阵是2阶,以此类推,可以表示n阶高维数据。

TF中Tensor支持的数据类型有很多,如tf.float16, tf.float32, tf.float64, tf.uint8, tf.int8, tf.int16, tf.int32, tf.int64, tf.string, tf.bool, tf.complex64等,所有Tensor运算都使用泛化的数据类型表示。

TF的Tensor定义和运算主要是调用Eigen矩阵计算库完成的。TF中Tensor的UML定义如图4。其中TensorBuffer指针指向Eigen::Tensor类型。其中,Eigen::Tensor[5][6]不属于Eigen官方维护的程序,由贡献者提供文档和维护,所以Tensor定义在Eigen unsupported模块中。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图5 Tensor数据结构定义

图5中,Tensor主要包含两个变量m_data和m_dimension,m_data保存了Tensor的数据块,T是泛化的数据类型,m_dimensions保存了Tensor的维度信息。

Eigen:Tensor的成员变量很简单,却支持非常多的基本运算,再借助Eigen的加速机制实现快速计算,参考章节3.2。Eigen::Tensor主要包含了

  • 一元运算(Unary),如sqrt、square、exp、abs等。
  • 二元运算(Binary),如add,sub,mul,div等
  • 选择运算(Selection),即if / else条件运算
  • 归纳运算(Reduce),如reduce_sum, reduce_mean等
  • 几何运算(Geometry),如reshape,slice,shuffle,chip,reverse,pad,concatenate,extract_patches,extract_image_patches等
  • 张量积(Contract)和卷积运算(Convolve)是重点运算,后续会详细讲解。

2.2 符号编程

编程模式通常分为命令式编程(imperative style programs)和符号式编程(symbolic style programs)。

命令式编程容易理解和调试,命令语句基本没有优化,按原有逻辑执行。符号式编程涉及较多的嵌入和优化,不容易理解和调试,但运行速度有同比提升。

这两种编程模式在实际中都有应用,Torch是典型的命令式风格,caffe、theano、mxnet和Tensorflow都使用了符号式编程。其中caffe、mxnet采用了两种编程模式混合的方法,而Tensorflow是完全采用了符号式编程,Theano和Tensorflow的编程模式更相近。

命令式编程是常见的编程模式,编程语言如python/C++都采用命令式编程。命令式编程明确输入变量,并根据程序逻辑逐步运算,这种模式非常在调试程序时进行单步跟踪,分析中间变量。举例来说,设A=10, B=10,计算逻辑:

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

第一步计算得出C=100,第二步计算得出D=101,输出结果D=101。

符号式编程将计算过程抽象为计算图,计算流图可以方便的描述计算过程,所有输入节点、运算节点、输出节点均符号化处理。计算图通过建立输入节点到输出节点的传递闭包,从输入节点出发,沿着传递闭包完成数值计算和数据流动,直到达到输出节点。这个过程经过计算图优化,以数据(计算)流方式完成,节省内存空间使用,计算速度快,但不适合程序调试,通常不用于编程语言中。举上面的例子,先根据计算逻辑编写符号式程序并生成计算图

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

其中A和B是输入符号变量,C和D是运算符号变量,compile函数生成计算图F,如图6所示。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图6 符号编程的正向计算图

最后得到A=10, B=10时变量D的值,这里D可以复用C的内存空间,省去了中间变量的空间存储。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图 6是TF中的计算流图,C=F(Relu(Add(MatMul(W, x), b))),其中每个节点都是符号化表示的。通过session创建graph,在调用session.run执行计算。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图7 TF符号计算图

和目前的符号语言比起来,TF最大的特点是强化了数据流图,引入了mutation的概念。这一点是TF和包括Theano在内的符号编程框架最大的不同。所谓mutation,就是可以在计算的过程更改一个变量的值,而这个变量在计算的过程中会被带入到下一轮迭代里面去。

Mutation是机器学习优化算法几乎必须要引入的东西(虽然也可以通过immutable replacement来代替,但是会有效率的问题)。 Theano的做法是引入了update statement来处理mutation。TF选择了纯符号计算的路线,并且直接把更新引入了数据流图中去。从目前的白皮书看还会支持条件和循环。这样就几乎让TF本身成为一门独立的语言。不过这一点会导致最后的API设计和使用需要特别小心,把mutation 引入到数据流图中会带来一些新的问题,比如如何处理写与写之间的依赖。[7]

2.3 梯度计算

梯度计算主要应用在误差反向传播和数据更新,是深度学习平台要解决的核心问题。梯度计算涉及每个计算节点,每个自定义的前向计算图都包含一个隐式的反向计算图。从数据流向上看,正向计算图是数据从输入节点到输出节点的流向过程,反向计算图是数据从输出节点到输入节点的流向过程。

图8是2.2节中图6对应的反向计算图。图中,由于C=A*B,则dA=B*dC, dB=A*dC。在反向计算图中,输入节点dD,输出节点dA和dB,计算表达式为dA=B*dC=B*dD, dB=A*dC=A*dD。每一个正向计算节点对应一个隐式梯度计算节点。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图8 符号编程的反向计算图

反向计算限制了符号编程中内存空间复用的优势,因为在正向计算中的计算数据在反向计算中也可能要用到。从这一点上讲,粗粒度的计算节点比细粒度的计算节点更有优势,而TF大部分为细粒度操作,虽然灵活性很强,但细粒度操作涉及到更多的优化方案,在工程实现上开销较大,不及粗粒度简单直接。在神经网络模型中,TF将逐步侧重粗粒度运算。

2.4 控制流

TF的计算图如同数据流一样,数据流向表示计算过程,如图9。数据流图可以很好的表达计算过程,为了扩展TF的表达能力,TF中引入控制流。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

图9 Graph的数据流

在编程语言中,if…else…是最常见的逻辑控制,在TF的数据流中也可以通过这种方式控制数据流向。接口函数如下,pred为判别表达式,fn1和fn2为运算表达式。当pred为true是,执行fn1操作;当pred为false时,执行fn2操作。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

TF还可以协调多个数据流,在存在依赖节点的场景下非常有用,例如节点B要读取模型参数θ更新后的值,而节点A负责更新参数θ,则节点B必须等节点A完成后才能执行,否则读取的参数θ为更新前的数值,这时需要一个运算控制器。接口函数如下,tf.control_dependencies函数可以控制多个数据流执行完成后才能执行接下来的操作,通常与tf.group函数结合使用。

从系统和代码实现角度解析TensorFlow的内部实现原理 | 深度

TF支持的控制算子有Switch、Merge、Enter、Leave和NextIteration等。

TF不仅支持逻辑控制,还支持循环控制。TF使用和MIT Token-Tagged machine相似的表示系统,将循环的每次迭代标记为一个tag,迭代的执行状态标记为一个frame,但迭代所需的数据准备好的时候,就可以开始计算,从而多个迭代可以同时执行。

3. TF 代码分析初步
3.1 TF总体概述
为了对TF有整体描述,本章节将选取TF白皮书[1]中的示例展开说明,如图 3 1所示是一个简单线性模型的TF正向计算图和反向计算图。图中x是输入,W是参数权值,b是偏差值,MatMul和Add是计算操作,dMatMul和dAdd是梯度计算操作,C是正向计算的目标函数,1是反向计算的初始值,dC/dW和dC/dx是模型参数的梯度函数。
图 3 1 tensorflow计算流图示例
以图 3 1为例实现的TF代码见图 3 2。首先声明参数变量W、b和输入变量x,构建线性模型y=W*x+b,目标函数loss采用误差平方和最小化方法,优化函数optimizer采用随机梯度下降方法。然后初始化全局参数变量,利用session与master交互实现图计算。

图 3 2 TF线性模型示例的实现代码

图 3 2中summary可以记录graph元信息和tensor数据信息,再利用tensorboard分析模型结构和训练参数。
图 3 3是上述代码在Tensorboard中记录下的Tensor跟踪图。Tensorboard可以显示scaler和histogram两种形式。跟踪变量走势可更方便的分析模型和调整参数。

图 3 3 Tensorboard显示的TF线性模型参数跟踪
图 3 4是图 3 1示例在Tensorboard中显示的graph图。左侧子图描述的正向计算图和反向计算图,正向计算的输出被用于反向计算的输入,其中MatMul对应MatMul_grad,Add对应Add_grad等。右上侧子图指明了目标函数最小化训练过程中要更新的模型参数W、b,右下侧子图是参数节点W、b展开后的结果。

图 3 4 Tensorboard显示的TF线性模型graph
图 3 4中,参数W是命名空间(Namespace)类型,展开后的W主要由Assign和Read两个OpNode组成,分别负责W的赋值和读取任务。
命名空间gradients是隐含的反向计算图,定义了反向计算的计算逻辑。从图 3 1可以看出,更新参数W需要先计算dMatMul,即图 3 4中的MatMul_grad操作,而Update_W节点负责更新W操作。为了进一步了解UpdateW的逻辑,图 3 5对MatMul_grad和update_W进行了展开分析。

图 3 5 MatMul_grad计算逻辑
图 3 5中,子图(a)描述了MatMul_grad计算逻辑,子图(b)描述了MatMul_grad输入输出,子图(c)描述了update_W的计算逻辑。首先明确MatMul矩阵运算法则,假设 z=MatMul(x, y),则有dx = MatMul(dz, y),dy = MatMul(x, dz),由此可以推出dW=MatMul(dAdd, x)。在子图(a)中左下侧的节点b就是输入节点x,dAdd由Add_grad计算输出。update_W的计算逻辑由最优化函数指定,而其中的minimize/update_W/ApplyGradientDescent变量决定,即子图(b)中的输出变量Outputs。
另外,在MatMul_grad/tuple命名空间中还隐式声明了control dependencies控制依赖操作,这在章节2.4控制流中相关说明。
3.2 Eigen介绍
在Tensoflow中核心数据结构和运算主要依赖于Eigen和Stream Executor库,其中Eigen支持CPU和GPU加速计算,Stream Executor主要用于GPU环境加速计算。下面简单讲述Eigen库的相关特性,有助于进一步理解Tensorflow。
3.2.1 Eigen简述
Eigen是高效易用的C++开源库,有效支持线性代数,矩阵和矢量运算,数值分析及其相关的算法。不依赖于任何其他依赖包,安装使用都很简便[8]。具有如下特性:
Ø  支持整数、浮点数、复数,使用模板编程,可以为特殊的数据结构提供矩阵操作。比如在用ceres-solver进行做优化问题(比如bundle adjustment)的时候,有时候需要用模板编程写一个目标函数,ceres可以将模板自动替换为内部的一个可以自动求微分的特殊的double类型。而如果要在这个模板函数中进行矩阵计算,使用Eigen就会非常方便。
Ø  支持逐元素、分块、和整体的矩阵操作。
Ø  内含大量矩阵分解算法包括LU,LDLt,QR、SVD等等。
Ø  支持使用Intel MKL加速
Ø  部分功能支持多线程
Ø  稀疏矩阵支持良好,到今年新出的Eigen3.2,已经自带了SparseLU、SparseQR、共轭梯度(ConjugateGradient solver)、bi conjugate gradient stabilized solver等解稀疏矩阵的功能。同时提供SPQR、UmfPack等外部稀疏矩阵库的接口。
Ø  支持常用几何运算,包括旋转矩阵、四元数、矩阵变换、AngleAxis(欧拉角与Rodrigues变换)等等。
Ø  更新活跃,用户众多(Google、WilliowGarage也在用),使用Eigen的比较著名的开源项目有ROS(机器人操作系统)、PCL(点云处理库)、Google Ceres(优化算法)。OpenCV自带到Eigen的接口。
Eigen库包含 Eigen模块和unsupported模块,其中Eigen模块为official module,unsupported模块为开源贡献者开发的。
Eigen unsupported 模块中定义了数据类型Tensor及相关函数,包括Tensor的存储格式,Tensor的符号表示,Tensor的编译加速,Tensor的一元运算、二元运算、高维度泛化矩阵运算,Tensor的表达式计算。本章后续所述Tensor均为Eigen::Tensor
Eigen运算性能评估如图 3 6所示[9],eigen3的整体性能比eigen2有很大提升,与GOTO2、INTEL_MKL基本持平。

图 3 6矩阵运算常用库比较
3.2.2 Eigen 存储顺序
Eigen中的Tensor支持两种存储方式:
Ø  Row-major表示矩阵存储时按照row-by-row的方式。
Ø  Col-major表示矩阵存储时按照column-by-column的方式。
Eigen默认采用Col-major格式存储的(虽然也支持Row-major,但不推荐),具体采用什么存储方式取决于算法本身是行遍历还是列遍历为主。例如:A=[[a11, a12, a13], [a21, a22, a23]]的存储序列见图 3 7。

图 3 7 Row-major和Column-major存储顺序
3.2.3 Eigen 惰性求值
在编程语言理论中,存在及早求值(Eager Evaluation) 和惰性求值(Lazy Evaluation)
Ø  及早求值:大多数编程语言所拥有的普通计算方式
Ø  惰性求值:也认为是“延迟求值”,可以提高计算性能,最重要的好处是它可以构造一个无限的数据类型。
关于惰性求值,举例如下:
Vec3 = vec1 + vec2;
及早求值形式需要临时变量vec_temp存储运算结果,再赋值给vec3,计算效率和空间效率都不高:
Vec_temp = vec1 + vec2;
Vec3 = vec_temp
而惰性求值不需要临时变量保存中间结果,提高了计算性能:
Vec_symbol_3 = (vec_symbol_1 + vec_symbol_2);
Vec3 = vec_symbol_3.eval(vec1, vec2)
由于Eigen默认采用惰性计算,如果要求表达式的值可以使用Tensor::eval()函数。Tensor::eval()函数也是session.run()的底层运算。例如:
Tensor<float, 3> result = ((t1 + t2).eval() * 0.2f).exp();
3.2.4 Eigen 编译加速
编译加速可以充分发挥计算机的并行计算能力,提高程序运行速度。
举例如下:
普通的循环相加运算时间复杂度是O(n):

如果指令集支持128bit并行计算,则时间复杂度可缩短为O(n/4):

Eigen编译时使用了SSE2加速。假设处理float32类型,指令集支持128bit并行计算,则一次可以计算4个float32类型,速度提升4倍。
3.2.5 Eigen::half
Tensorflow支持的浮点数类型有float16, float32, float64,其中float16本质上是Eigen::half类型,即半精度浮点数[10]。关于半精度浮点数,英伟达2002年首次提出使用半精度浮点数达到降低数据传输和存储成本的目的。
在分布式计算中,如果对数据精度要求不那么高,可以将传输数据转换为float16类型,这样可以大大缩短设备间的数据传输时间。在GPU运算中,float16还可以减少一般的内存占用。
在Tensorflow的分布式传输中,默认会将float32转换为float16类型。Tensorflow的转换方式不同于nvidia的标准,采用直接截断尾数的方式转化为半精度浮点数,以减少转换时间。
图 3 8是双精度浮点数(float64)存储格式。

图 3 8 双精度浮点数
图 3 9是单精度浮点数(float32)存储格式。

图 3 9 单精度浮点数
图 3 10是半精度浮点数(float16)存储格式。

图 3 10 半精度浮点数
浮点数存储格式分成3部分,符号位,指数和尾数。不同精度是指数位和尾数位的长度不一样。
3.3 设备内存管理
TF设备内存管理模块利用BFC算法(best-fit with coalescing)实现。BFC算法是Doung Lea’s malloc(dlmalloc)的一个非常简单的版本。它具有内存分配、释放、碎片管理等基本功能[11]。
BFC将内存分成一系列内存块,每个内存块由一个chunk数据结构管理。从chunk结构中可以获取到内存块的使用状态、大小、数据的基址、前驱和后继chunk等信息。整个内存可以通过一个chunk的双链表结构来表示。

图 3 11内存分块结构图
用户申请一个内存块(malloc)。根据建立的chunk双链表找到一个合适的内存块(后面会说明什么是合适的内存块),如果该内存块的大小是用户申请大小的两倍以上,那么将该内存块切分成两块,这就是split操作。返回其中一块给用户,并将该内存块标识为占用。Spilt操作会新增一个chunk,所以需要修改chunk双链表以维持前驱和后继关系。
用户释放一个内存块(free)。先将该块标记为空闲。然后根据chunk数据结构中的信息找到其前驱和后继内存块。如果前驱和后继块中有空闲的块,那么将刚释放的块和空闲的块合并成一个更大的chunk(这就是merge操作,合并当前块和其前后的空闲块)。再修改双链表结构以维持前驱后继关系。这就做到了内存碎片的回收。
BFC的核心思想是:将内存分块管理,按块进行空间分配和释放;通过split操作将大内存块分解成小内存块;通过merge操作合并小的内存块,做到内存碎片回收。
但是还留下许多疑问。比如说申请内存空间时,什么样的块算合适的内存块?如何快速管理这种块?
BFC算法采取的是被动分块的策略。最开始整个内存是一个chunk,随着用户申请空间的次数增加,最开始的大chunk会被不断的split开来,从而产生越来越多的小chunk。当chunk数量很大时,为了寻找一个合适的内存块而遍历双链表无疑是一笔巨大的开销。为了实现对空闲块的高效管理,BFC算法设计了bin这个抽象数据结构。
Bin数据结构中,每个bin都有一个size属性,一个bin是一个拥有chunk size >= bin size的空闲chunk的集合。集合中的chunk按照chunk size的升序组织成单链表。BFC算法维护了一个bin的集合:bins。它由多个bin以及从属于每个bin的chunks组成。内存中所有的空闲chunk都由bins管理。

图 3 12 bins集合的结构图
图 3 12中每一列表示一个bin,列首方格中的数字表示bin的size。bin size的大小都是256的2^n的倍。每个bin下面挂载了一系列的空闲chunk,每个chunk的chunk size都大于等于所属的bin的bin size,按照chunk size的升序挂载成单链表。BFC算法针对bins这个集合设计了三个操作:search、insert、delete。
Search 操作:给定一个chunk size,从bins中找到大于等于该chunk size的最小的那个空闲chunk。Search操作具体流程如下。如果bin以数组的形式组织,那么可以从index = chunk size /256 >>2的那个bin开始查找。较好的情况是开始查找的那个bin的chunk链表非空,那么直接返回链表头即可。这种情况时间复杂度是常数级的。最坏的情况是遍历bins数组中所有的bin。对于一般大小的内存来说,bins数组元素非常少,比如4G空间只需要23个bin就足够了(256 * 2 ^ 23 > 4G),因此也很快能返回结果。总体来说search操作是非常高效的。对于固定大小内存来说,查找时间是常数量级的。
Insert 操作:将一个空闲的chunk插入到一个bin所挂载的chunk链表中,同时需要维持chunk链表的升序关系。具体流程是直接将chunk插入到index = chunk size /256 >>2的那个bin中即可。
Delete操作:将一个空闲的chunk从bins中移除。
TF中内存分配算法实现文件core/common_runtime/bfc_allocator.cc,GPU内存分配算法实现文件core/common_runtime/gpu/gpu_bfc_allocator.cc。
3.4 TF开发工具介绍
TF系统开发使用了bazel工具实现工程代码自动化管理,使用了protobuf实现了跨设备数据传输,使用了swig库实现python接口封装。本章将从这三方面介绍TF开发工具的使用。
3.4.1 Swig封装
Tensorflow核心框架使用C++编写,API接口文件定义在tensorflow/core/public目录下,主要文件是tensor_c_api.h文件,C++语言直接调用这些头文件即可。
Python通过Swig工具封装TF库包间接调用,接口定义文件tensorflow/python/ tensorflow.i。其中swig全称为Simplified Wrapper and Interface Generator,是封装C/C++并与其它各种高级编程语言进行嵌入联接的开发工具,对swig感兴趣的请参考相关文档。
在tensorflow.i文件中包含了若干个.i文件,每个文件是对应模块的封装,其中tf_session.i文件中包含了tensor_c_api.h,实现client向session发送请求创建和运行graph的功能。
3.4.2 Bazel编译和调试
Bazel是Google开源的自动化构建工具,类似于Make和CMake工具。Bazel的目标是构建“快速并可靠的代码”,并且能“随着公司的成长持续调整其软件开发实践”。
TF中几乎所有代码编译生成都是依赖Bazel完成的,了解Bazel有助于进一步学习TF代码,尤其是编译测试用例进行gdb调试。
Bazel假定每个目录为[package]单元,目录里面包含了源文件和一个描述文件BUILD,描述文件中指定了如何将源文件转换成构建的输出。
以图 3 13为例,左子图为工程中不同模块间的依赖关系,右子图是对应模块依赖关系的BUILD描述文件。
图 3 13中name属性来命名规则,srcs属性为模块相关源文件列表,deps属性来描述规则之间的依赖关系。”//search: google_search_page”中”search”是包名,”google_search_page”为规则名,其中冒号用来分隔包名和规则名;如果某条规则所依赖的规则在其他目录下,就用”//”开头,如果在同一目录下,可以忽略包名而用冒号开头。
图 3 13中cc_binary表示编译目标是生成可执行文件,cc_library表示编译目标是生成库文件。如果要生成google_search_page规则可运行

如果要生成可调试的二进制文件,可运行

图 3 13 Bazel BUILD文件示例
TF中首次运行bazel时会自动下载很多依赖包,如果有的包下载失败,打开tensorflow/workspace.bzl查看是哪个包下载失败,更改对应依赖包的new_http_archive中的url地址,也可以把new_http_archive设置为本地目录new_local_repository。
TF中测试用例跟相应代码文件放在一起,如MatMul操作的core/kernels/matmul_op.cc文件对应的测试用例文件为core/kernels/matmul_op_test.cc文件。运行这个测试用例需要查找这个测试用例对应的BUILD文件和对应的命令规则,如matmul_op_test.cc文件对应的BUILD文件为core/kernels/BUILD文件,如下

其中tf_cuda_cc_test函数是TF中自定义的编译函数,函数定义在/tensorflow/ tensorflow.bzl文件中,它会把matmul_op_test.cc放进编译文件中。要生成matmul_op_test可执行文件可运行如下脚本:

3.4.3 Protobuf序列化
Protocol Buffers 是一种轻便高效的结构化数据存储格式,可以用于结构化数据串行化,或者说序列化。它很适合做数据存储或 RPC 数据交换格式。可用于通讯协议、数据存储等领域的语言无关、平台无关、可扩展的序列化结构数据格式。
Protobuf对象描述文件为.proto类型,编译后生成.pb.h和.pb.cc文件。
Protobuf主要包含读写两个函数:Writer(序列化)函数SerializeToOstream() 和  Reader(反序列化)函数 ParseFromIstream()。
Tensorflow在core/probobuf目录中定义了若干与分布式环境相关的.proto文件,同时在core/framework目录下定义了与基本数据类型和结构的.proto文件,在core/util目录中也定义部分.proto文件,感觉太随意了。
在分布式环境中,不仅需要传输数据序列化,还需要数据传输协议。Protobuf在序列化处理后,由gRPC完成数据传输。gRPC数据传输架构图见图 3 14。

图 3 14 gRPC数据传输架构
gRPC服务包含客户端和服务端。gRPC客户端调用stub 对象将请求用 protobuf 方式序列化成字节流,用于线上传输,到 server 端后调用真正的实现对象处理。gRPC的服务端通过observer观察处理返回和关闭通道。
TF使用gRPC完成不同设备间的数据传输,比如超参数、梯度值、graph结构。
5. TF – Graph模块
TF把神经网络模型表达成一张拓扑结构的Graph,Graph中的一个节点表示一种计算算子。Graph从输入到输出的Tensor数据流动完成了一个运算过程,这是对类似概率图、神经网络等连接式算法很好的表达,同时也是对Tensor + Flow的直观解释。
5.1 Graph视图
Tensorflow采用符号化编程,形式化为Graph计算图。Graph包含节点(Node)、边(Edge)、NameScope、子图(SubGraph),图 5 1是Graph的拓扑描述。
Ø  节点分为计算节点(Compute Node)、起始点(Source Node)、终止点(Sink Node)。起始点入度为0,终止点出度为0。
Ø  NameScope为节点创建层次化的名称,图 3 4中的NameSpace类型节点就是其中一种体现。
Ø  边分为普通边和依赖边(Dependecy Edge)。依赖边表示对指定的计算节点有依赖性,必须等待指定的节点计算完成才能开始依赖边的计算。

图 5 1 Graph的拓扑描述
图 5 2是Graph的UML视图模型,左侧GraphDef类为protobuf中定义的graph结构,可将graph结构序列化和反序列化处理,用于模型保存、模型加载、分布式数据传输。右侧Graph类为/core/graph模块中定义的graph结构,完成graph相关操作,如构建(construct),剪枝(pruning)、划分(partitioning)、优化(optimize)、运行(execute)等。GraphDef类和Graph类可以相关转换,如图中中间部分描述,函数Graph::ToGraphDef()将Graph转换为GraphDef,函数ConvertGraphDefToGraph将GraphDef转换为Graph,借助这种转换就能实现Graph结构的网络传输。

图 5 2 Graph的UML视图
Graph-UML图中还定义了Node和Edge。Node定义函数操作和属性信息,Edge连接源节点和目标节点。类NodeDef中定义了Op、Input、Device、Attr信息,其中Device可能是CPU、GPU设备,甚至是ARM架构的设备,说明Node是与设备绑定的。类FunctionDefLibrary主要是为了描述各种Op的运算,包括Op的正向计算和梯度计算。FunctionDef的定义描述见图 5 3。

图 5 3 FunctionDef的定义
图 5 4是FunctionDef举例,对MatMulGrad的梯度描述,其中包含函数参数定义、函数返回值定义、模板数据类型定义、节点计算逻辑。

图 5 4 FunctionDef举例:MatMulGrad
5.2 Graph构建
有向图(DAG)由节点和有向边组成。本章节主要讲述TF如何利用<Nodes, Edges>组合成完整的graph的。假设有如下计算表达式:t1=MatMul(input, W1)。

图 5 5 Graph简单示例
图 5 5中图计算表达式包含3个节点,2条边,描述为字符串形式如下。

5.3 Graph局部执行
Graph的局部执行特性允许使用者从任意一个节点输入(feed),并指定目标输出节点(fetch)。图 5 6是TF白皮书中描述Graph局部执行的图。[15]

图 5 6 Graph局部执行

5.4 Graph设备分配
TF具有高度设备兼容性,支持X86和Arm架构,支持CPU、GPU运算,可运行于Linux、MacOS、Android和IOS系统。而且,TF的设备无关性特征在多设备分布式运行上也非常有用。
Graph中每个节点都分配有设备编号,表示该节点在相应设备上完成计算操作。用户既可以手动指定节点设备,也可以利用TF自动分配算法完成节点设备分配。设备自动算法需要权衡数据传输代价和计算设备的平衡,尽可能充分利用计算设备,减少数据传输代价,从而提高计算性能。
Graph设备分配用于管理多设备分布式运行时,哪些节点运行在哪个设备上。TF设备分配算法有两种实现算法,第一种是简单布放算法(Simple Placer),第二种基于代价模型(Cost Model)评估。简单布放算法按照指定规则布放,比较简单粗放,是早期版本的TF使用的模型,并逐步被代价模型方法代替。
5.4.1 Simple Placer算法
TF实现的Simple Placer设备分配算法使用union-find方法和启发式方法将部分不相交且待分配设备的Op节点集合合并,并分配到合适的设备上。
Union-find(联合-查找)算法是并查集数据结构一种应用。并查集是一种树型的数据结构,其保持着用于处理一些不相交集合(Disjoint Sets)的合并及查询问题。Union-find定义了两种基本操作:Union和Find。
Ø  Find:确定元素属于哪一个子集。它可以被用来确定两个元素是否属于同一子集。
Ø  Union:将两个子集合并成同一个集合。即将一个集合的根节点的父指针指向另一个集合的根节点。
启发式算法(Heuristic Algorithm)定义了节点分配的基本规则。Simple Placer算法默认将起始点和终止点分配给CPU,其他节点中GPU的分配优先级高于CPU,且默认分配给GPU:0。启发式规则适用于以下两种场景:
Ø  对于符合GeneratorNode条件(0-indegree, 1-outdegree, not ref-type)的节点,让node与target_node所在device一致,参见图 5 7。

TF中Simple Placer的实现定义在文件core/common_runtime/simple_placer.cc。文件中主要定义了两个类:ColocationGraph和SimplePlacer。ColocationGraph类利用Union-find算法将节点子集合合并成一个节点集合,参考成员函数ColocationGraph:: ColocateNodes实现。SimplePlacer类实现节点分配过程,下面将主要介绍SimplePlacer:: Run()函数的实现过程。

5.4.2 代价模型
TF使用代价模型(Cost Model)会在计算流图生成的时候模拟每个device上的负载,并利用启发式策略估计device上的完成时间,最终找出预估时间较低的graph设备分配方案。[1]
Cost model预估时间的方法有两种:
Ø  使用启发式的算法,通过把输入和输出的类型以及tensor的大小输入进去,得到时间的预估
Ø  使用模拟的方法,对图的计算进行一个模拟,得到各个计算在其可用的设备上的时间。
启发式策略会根据如下数据调整device的分配:节点任务执行的总时间;单个节点任务执行的累计时间;单个节点输出数据的尺寸。

图 5 9代价模型UML视图
TF中代价模型的实现定义在文件core/graph/costmodel.cc和core/common_runtime/ costmodel_manager.cc,其UML视图参见图 5 9。
Cost model manager从graph创建cost model,再评估计算时间,如下。

Ø  Function Inlining (函数内联)
函数内联处理可减少方法调用的成本。在TF中包含以下几种方法:
Ø  RemoveListArrayConverter(g):” Rewrites _ListToArray and _ArrayToList to a set of Identity nodes”.
Ø  RemoveDeadNodes(g):删除DeatNode。DeatNode的特征是”not statefull, not _Arg,  not reachable from _Retval”.
Ø  RemoveIdentityNodes(g):删除Identity节点。如n2=Identity(n1) + Identity(n1);  优化后:  n2=n1 + n1;
Ø  FixupSourceAndSinkEdges(g):固定source和sink的边
Ø  ExpandInlineFunctions(runtime, g):展开内联函数的嵌套调用
其中_ListToArray、_ArrayToList、_Arg、_Retval均在core/ops/function_ops.cc中定义。
Graph优化相关测试文件在common_runtime/function_test.cc,调试方法:

7. Reference
[1]. Abadi M, Agarwal A, Barham P, et al. Tensorflow: Large-scale machine learning on heterogeneous distributed systems[J]. arXiv preprint arXiv:1603.04467, 2016.
[2]. Tensorflow in Github. https://github.com/tensorflow/tensorflow.
[3]. Tensorflow official website. https://www.tensorflow.org/
[4]. TensorFlow™ – Open Source Library for Machine Learning Applications. https://delftswa. gitbooks.io/desosa2016/content/tensorflow/chapter.html
[5]. Eigen unsupported documentation. https://eigen.tuxfamily.org/dox-devel/unsupported/ index.html
[6]. Tensor official documentation. https://bitbucket.org/eigen/eigen/src/default/unsupported/ Eigen/CXX11/src/Tensor/README.md?fileviewer=file-view-default
[7]. 如何评价Tensorflow和其它深度学习系统. http://weibo.com/p/10016039076107377 75666.
其中评估时间的函数EstimateComputationCosts是对graph中每个node依次评估,节点计算时间评估函数如下。
[8]. Eigen official documentation. https://eigen.tuxfamily.org/dox-devel/index.html
[9]. Eigen Benchmark. http://eigen.tuxfamily.org/index.php?title=Benchmark
[10]. Half-precision floating-point format. https://en.wikipedia.org/wiki/Half-precision_floating -point_format
[11]. Tensorflow设备内存分配算法解析. http://weibo.com/p/1001603980563068394770
[12]. Tensorflow: Adding a New Op. https://www.tensorflow.org/how_tos/adding_an_op/#adding -a-new-op
[13]. Tensorflow: Reading data. https://www.tensorflow.org/how_tos/reading_data/
[14]. Tensorflow: Custom Data Readers. https://www.tensorflow.org/how_tos/new_data_formats
[15]. Goldsborough P. A Tour of TensorFlow[J]. 2016.
[16]. Click C. Global code motion/global value numbering[C]//ACM SIGPLAN Notices. ACM, 1995, 30(6): 246-257.
[17]. Kevin Robinson. A tour through the TensorFlow codebase.

作者:姚健,毕业于中科院计算所网络数据实验室,毕业后就职于360天眼实验室,主要从事深度学习和增强学习相关研究工作。目前就职于腾讯MIG事业部,从事神经机器翻译工作。联系方式: yao_62995@163.com

来源:https://zhuanlan.zhihu.com/p/25932160

分享&收藏
关键词:

版权声明:本站原创和会员推荐转载文章,仅供学习交流使用,不会用于任何商业用途,转载本站文章请注明来源、原文链接和作者,否则产生的任何版权纠纷与本站无关,如果有文章侵犯到原作者的权益,请您与我们联系删除或者进行授权,联系邮箱:service@datagold.com.cn。

发表评论

电子邮件地址不会被公开。 必填项已用*标注