乐趣区

Tensorflow源码解析1 — 内核架构和源码结构

1 主流深度学习框架对比
当今的软件开发基本都是分层化和模块化的,应用层开发会基于框架层。比如开发 Linux Driver 会基于 Linux kernel,开发 Android app 会基于 Android Framework。深度学习也不例外,框架层为上层模型开发提供了强大的多语言接口、稳定的运行时、高效的算子,以及完备的通信层和设备层管理层。因此,各大公司早早的就开始了深度学习框架的研发,以便能占领市场。当前的框架有数十种之多,主流的如下(截止到 2018 年 11 月)

显然 TensorFlow 是独一无二的王者。第二名 Keras,它是对 TensorFlow 或 Theano 接口的二次封装,严格意义上并不是一个独立的深度学习框架。TensorFlow 目前也已经集成了 Keras,使得安装了 TensorFlow 的用户就可以直接使用 Keras 了。
TensorFlow 之所以能够从数十种框架中脱颖而出,主要优点有

出身高贵,是谷歌出品的。但其他很多框架出身也不差,例如 PyTorch 之于 Facebook,MXNET 之于 Amazon
2015 年就开源了,比较早的俘获了一大批开发者。这个确实是 tf 的一大先发优势,但 PyTorch 的前身 Caffe,以及 MXNET 开源时间都不晚,而且 Caffe 流行时间比 tf 早,后来才被赶超的。更有 Theano 这样的绝对老前辈。由此可见,软件开源是多么重要。目前流行的深度学习框架也基本都开源了。
支持的开发语言多,支持 Python Java Go C++ 等多种流行语言。相比某些框架,确实是优势很大。相比 MXNET 则小巫见大巫了。MXNET 早期发展的一个主要方向就是前端多语言的支持,连 MATLAB R Julia 等语言都支持了。
运行效率高。早期的时候,其实 tf 的运行效率比很多框架都要低一些的。
安装容易,用户上手快,文档齐全,社区活跃。这个是 tf 的一个较大优势,特别是社区方面,也就是我们常说的生态优势。互联网头部集中效应十分明显,体现在开源软件上也是一样。这也是我认为最大的一个优势。

总结起来,TensorFlow 虽然每个方面都不是绝对领先的优势,但贵在每个方面都做的不错,因此最终能够一骑绝尘,独领风骚。
学习 Tensorflow 框架内核,可以理解前端接口语言的支持,session 生命周期,graph 的构建、分裂和执行,operation 的注册和运行,模块间数据通信,本地运行和分布式运行模式,以及 CPU GPU TPU 等异构设备的封装支持等。学习这些,对于模型的压缩 加速 优化等都是大有裨益的。
2 TensorFlow 系统架构
TensorFlow 设计十分精巧,基于分层和模块化的设计思想进行开发的。框架如下图

整个框架以 C API 为界,分为前端和后端两大部分。

前端:提供编程模型,多语言的接口支持,比如 Python Java C++ 等。通过 C API 建立前后端的连接,后面详细讲解。

后端:提供运行环境,完成计算图的执行。进一步分为 4 层

运行时:分为分布式运行时和本地运行时,负责计算图的接收,构造,编排等。
计算层:提供各 op 算子的内核实现,例如 conv2d, relu 等
通信层:实现组件间数据通信,基于 GRPC 和 RDMA 两种通信方式
设备层:提供多种异构设备的支持,如 CPU GPU TPU FPGA 等

模型构造和执行流程
TensorFlow 的一大特点是,图的构造和执行相分离。用户添加完算子,构建好整图后,才开始进行训练和执行,也就是图的执行。大体流程如下

图构建:用户在 client 中基于 TensorFlow 的多语言编程接口,添加算子,完成计算图的构造。
图传递:client 开启 session,通过它建立和 master 之间的连接。执行 session.run() 时,将构造好的 graph 序列化为 graphDef 后,以 protobuf 的格式传递给 master。
图剪枝:master 根据 session.run() 传递的 fetches 和 feeds 列表,反向遍历全图 full graph,实施剪枝,得到最小依赖子图
图分裂:master 将最小子图分裂为多个 Graph Partition,并注册到多个 worker 上。一个 worker 对应一个 Graph Partition。
图二次分裂:worker 根据当前可用硬件资源,如 CPU GPU,将 Graph Partition 按照 op 算子设备约束规范(例如 tf.device(‘/cpu:0’),二次分裂到不同设备上。每个计算设备对应一个 Graph Partition。
图运行:对于每一个计算设备,worker 依照 op 在 kernel 中的实现,完成 op 的运算。设备间数据通信可以使用 send/recv 节点,而 worker 间通信,则使用 GRPC 或 RDMA 协议。

3 前端多语言实现 – swig 包装器
TensorFlow 提供了很多种语言的前端接口,使得用户可以通过多种语言来完成模型的训练和推断。其中 Python 支持得最好。这也是 TensorFlow 之所以受欢迎的一大原因。前端多语言是怎么实现的呢?这要归功于 swig 包装器。
swig 是个帮助使用 C 或者 C ++ 编写的软件能与其它各种高级编程语言进行嵌入联接的开发工具。在 TensorFlow 使用 bazel 编译时,swig 会生成两个 wrapper 文件

pywrap_tensorflow_internal.py:对接上层 Python 调用
pywrap_tensorflow_internal.cc:对接底层 C API 调用。

pywrap_tensorflow_internal.py 模块被导入时,会加载_pywrap_tensorflow_internal.so 动态链接库,它里面包含了所有运行时接口的符号。而 pywrap_tensorflow_internal.cc 中,则注册了一个函数符号表,实现 Python 接口和 C 接口的映射。运行时,就可以通过映射表,找到 Python 接口在 C 层的实现了。

4 tensorflow 源码结构
TensorFlow 源码基本也是按照框架分层来组织文件的。如下

其中 core 为 tf 的核心,它的源码结构如下

5 总结
TensorFlow 框架设计精巧,代码量也很大,我们可以从以下部分逐步学习

TensorFlow 内核架构和源码结构。先从全局上对框架进行理解。
前后端连接的桥梁 –Session,重点理解 session 的生命周期,并通过相关源码可以加深理解 Python 前端如何调用底层 C 实现。
TensorFlow 核心对象—Graph。图 graph 是 TensorFlow 最核心的对象,基本都是围绕着它来进行的。graph 的节点为算子 operation,边为数据 tensor。
TensorFlow 图的节点 — Operation。operation 是图 graph 的节点,承载了计算算子。
TensorFlow 图的边 — Tensor。Tensor 是图 graph 的边,承载了计算的数据。
TensorFlow 本地运行时。
TensorFlow 分布式运行时。和本地运行时有一些共用的接口,但区别也很大。
TensorFlow 设备层。主要了解设备层的定义规范,以及实现。
TensorFlow 队列和并行运算。
TensorFlow 断点检查 checkpoint,模型保存 Saver,以及可视化 tensorboard。这三个为 TensorFlow 主要的工具。

本文作者:扬易阅读原文
本文为云栖社区原创内容,未经允许不得转载。

退出移动版