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

Tensorflow 架构介绍

TensorFlow Layers

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

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

[INFO] Graph was finalized

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

graphs_computation.png

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

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 文件中加载执行了函数:

# 定义项目名为 org_tensorflow
workspace(name = "org_tensorflow")
# 从当前依赖的 tensorflow 包的 workspace3.bzl 文件中导入 tf_workspace3 函数
load("@//tensorflow:workspace3.bzl", "tf_workspace3") 
# 调用tf_workspace3 函数
tf_workspace3()
...

在 Bazel 中同时也约定了项目的可独立编译的模块——包 (Package) 要包含一个 BUILD 文件。因此我们在 tensorflow 的根路径以及各个一级目录中都能看到一个 BUILD 文件。在 BUILD 文件中会定义多个规则 (rule),这些规则组织起来形成编译链,来完成对整个项目的编译、打包、测试等一系列操作。XXX.bzl 文件相当于是 Bazel 脚本,方便规则复用以及简化 BUILD 文件结构,在下一篇博客中我们将详细介绍 Bazel 用法以及如何组织项目。

除了 Bazel 相关的文件之外我们在根目录下可以看到 Tensorflow 目录,这个目录下包含了 Tensorflow 的源码。其中我们主要关注 core 目录。因为这个目录是 Tensorflow 的底层代码。其他诸如go、cc 之类的语言名目录包含的是各语言上层 API 源码,lite、security 等目录提供了一些额外的功能。

在 core 目录下我们重点关注 ops 目录以及 kernels 目录。ops 目录声明了各种 TensorFlow 的算子。但是内核并不在这个目录下,而是在 kernels 目录下。以 BatchMatMul 算子为例,算子的定义在 ops/math_ops.cc 文件中,但是内核的实现在 kernels.matmul_op_real.cc 文件中,通过 REGISTER_BATCH_MATMUL_TOUT_CPU 宏以及 BatchMatMulV2Op 模板类来实现。在后续的博客中我们将解析 Tensorflow 算子的源码,带读者了解不同算子的实现。

Tensorflow 构建实践

从源代码构建」 这篇官方文档其实已经写得很全面。我们简单实践一下 mac 上的 Tensorflow 的构建。构建的第一步是安装 bazel,mac 使用 homebrew 安装即可:

/bin/bash -c "$(curl -fsSL \
https://raw.githubusercontent.com/Homebrew/install/master/install.sh)"
brew install bazel
bazel --version #打印 bazel 版本

接着进入到 Tensorflow 目录中,修改 Tensorflow 的 .bazelversion 文件设置 bazel 版本为你安装的版本。然后执行./configure命令运行构建前的配置选择,如果要编译 GPU 版本的话只能在 linux 系统下,而且需要先安装 cuda 以及 cudnn

最后执行bazel build //tensorflow/tools/pip_package:build_pip_package构建 Python 包。如果构建成功可以在 bazel-bin/tensorflow/tools/pip_package 目录下找到 build_pip_package 文件,然后执行以下命令:

./bazel-bin/tensorflow/tools/pip_package/build_pip_package --nightly_flag [输出目录]

在输出目录中会有一个 tensorflow 的 whl 包,使用 pip install 安装即可。


标题:从零开始掌握 tensorflow 算子开发 (1):tensorflow 架构以及编译构建
作者:KIRI
地址:https://kirivir.github.io/articles/2021/10/20/1634668182447.html