系列文章:

从零开始掌握 tensorflow 算子开发系列文章

作为算子开发教程的第一篇,我们首先简单介绍 Tensorflow 的架构以及的算子 (Operator) 在 Tensorflow 的计算图扮演的角色。在简单了解算子之后,我们将实践介绍如何从 Tensorflow 源码编译构建属于自己的 Tensorflow Python 包。

Tensorflow 架构介绍

Tensorflow 是一个多语言的项目,Tensorflow 的底层功能主要由 C 与 C++ 实现,并在此基础上派生出了 Python api、Java api、C++ api。因为 Python 相较于 C++ 这种静态语言,能支持交互式编程写起来比较舒服,基本上都是 python 来写训练模型的脚本,而在线上部署模型的时候 C++、Java 使用的频率更多。等这个系列博客结束之后,有时间我们会继续介绍 Tensorflow 部署的相关内容。

用 Tensorflow 训练模型的时候,在模型跑起来之前会看到这段日志:

1
[INFO] Graph was finalized

Tensorflow 定义了一个图,Tensorflow 实际运行的时候会先将 python 代码定义的网络结构解析为一个有向无环的计算图,通过这个计算图再调度计算资源运行模型。熟悉 Tensorflow 的读者通过 TensorBoard 可以浏览计算图的全貌,如下图所示:

这个计算图中的每一个节点,除了输入和输出节点外,每个中间的节点都代表对张量 (Tensor) 的一个操作。从 checkpoint 目录下的 graph.pbtxt 文件中,我们可以找到每一个节点的结构,例如一个矩阵乘法的节点:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
node {
   name: "dnn/dense/MatMul"
   op: "MatMul"
   input: "dnn/input_layer/Reshape_143"
   input: "dense/kernel/read"
   attr {
     key: "T"
     value {
       type: DT_FLOAT
     }
   }
   省略 ...
 }

我们可以看到每个节点有几个域:name、op、input、attr,其中 name、input、attr 都很容易理解,分别是节点的名字、输入 tensor 以及节点的额外属性,op 则是我们这个系列博客的主题——算子(Operator),张量操作的具体实现。代码块中的 MatMul 就是矩阵乘法算子。如果我们在 Python 代码中使用了 tf.matmul 函数,Tensorflow 就会在计算图中生成一个 MatMul 节点,在加载模型的时候,会对将计算图进行编译,此时根据节点的 op 域从运行时的上下文中调用 MatMul 算子对应的内核(kernel),例如 cpu 环境下调用的是 MatMul 算子的 cpu 内核,gpu 环境下调用的是 gpu 内核。相当于算子类似 C++ 的抽象类,定义了张量操作的接口,kernel 是抽象类的实现。通过动态代理根据运行时的上下文采用不同的实现。下一小节介绍 Tensorflow 源码结构的时候会进一步反映算子和内核的差别。

Tensorflow 源码结构

在简单了解了 Tensorflow 的整体架构之后,我们来看看 Tensorflow 的源码结构,Tensorflow 使用了 Bazel 来组织项目。Bazel 是由 Google 开源的一个自动化构建系统,与之类似的有 Java 的 Maven。Bazel 通过 starlark(一种 Python 方言) 脚本来组织项目结构,约定项目根路径必须包含一个 WORKSPACE 文件作为根路径的标识。因此我们在 Tensorflow 源码的根路径上可以看到这个文件,里面从多个 .bzl 文件中加载执行了函数:

1
2
3
# 定义项目名为 org_tensorflow
workspace(name = "org_tensorflow")
# 从当前依赖的 tensorflow 包的 work