简介
前言，当查看tf python 代码时， 碰到一个函数实际由C++ 实现时，
pybind11
动手学深度学习框架(2) - python端如何调用c++的代码
pybind11 是一个轻量级的只包含头文件(header-only)的 c++ 库，用于将 c++ 代码暴露给 python 调用（反之亦可，但主要还是前者）。
#include <pybind11/pybind11.h> // pybind11的头文件
// PYBIND11_MODULE 是一个宏，实现一个 Python 扩展模块
PYBIND11_MODULE(pydemo, m){ // 定义Python模块pydemo，之后在 Python 脚本里必须用这个名字才能 import。m其实是 pybind11::module 的一个实例对象，它只是个普通的变量，起什么名字都可以，但为了写起来方便，一般都用“m”。
m.doc() = "pybind11 demo doc"; // 模块的说明文档
m.def("add", [](int a, int b) -> int { return a + b; }); // def函数，传递一个 Python 函数名和 C++ 的函数、函数对象或者是 lambda 表达式
} // Python模块定义结束
假设这个 C++ 源文件名是“pybind.cpp”，用 g++ 把它编译成在 Python 里调用的模块，生成一个大概这样的文件：pydemo.cpython-35m-x86_64-linux-gnu.so
g++ pybind.cpp \ #编译的源文件
-std=c++11 -shared -fPIC \ #编译成动态库
`python3 -m pybind11 --includes` \ #获得 pybind11 所在的包含路径，让 g++ 能够找得到头文件
-o pydemo`python3-config --extension-suffix` #生成的动态库名字，前面必须是源码里的模块名，而后面那部分则是 Python 要求的后缀名
之后就可以在python 中使用
import pydemo
x = pydemo.add(1,2)
print(x)
进阶：C++ 里的类也能够等价地转换到 Python 里面调用
自定义算子
一个Op可以接收一个或者多个输入Tensor，然后产生零个或者多个输出Tensor，分别利用Input和Output定义。在注册一个Op之后，就需要继承OpKernel，实现他的计算过程Compute函数，在Compute函数中，我们可以通过访问OpKernelContext来获得输入和输出信息。当我们需要申请新的内存地址时，可以通过OpKernelContext去申请TempTensor或者PersistentTensor。一般Tensorflow的Op都采用Eigen来操作Tensor
Adding a Ncw Op对于 TensorFlow，可以自定义 Operation，即如果现有的库没有涵盖你想要的操作, 你可以自己定制一个。为了使定制的 Op 能够兼容原有的库，你必须做以下工作:
- 在一个 C++ 文件中注册新 Op. Op 的注册与实现是相互独立的. 在其注册时描述了 Op 该如何执行. 例如, 注册 Op 时定义了 Op 的名字, 并指定了它的输入和输出.REGISTER_OP("ZeroOut") .Input("to_zero: int32") .Output("zeroed: int32") .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { c->set_output(0, c->input(0)); return Status::OK(); });
- 使用 C++ 实现 Op. 每一个实现称之为一个 “kernel”, 可以存在多个 kernel, 以适配不同的架构 (CPU, GPU 等)或不同的输入/输出类型.
- bazel（tf编译工具） 会检索所有op 并创建一个 Python wrapper. 这个wrapper是创建 Op 的公开 API. 当注册 Op 时, 会自动生成一个默认 默认的包装器. 既可以直接使用默认包装器, 也可以添加一个新的包装器.
- (可选) 写一个函数计算 Op 的梯度，在Python 中注册.@ops.RegisterGradient("ZeroOut") def?_zero_out_grad(op, grad): xxxxxxxxx
- (可选) 写一个函数, 描述 Op 的输入和输出 shape. 该函数能够允许从 Op 推断 shape.
- 测试 Op, 通常使用 Pyhton。如果你定义了梯度，你可以使用Python的GradientChecker来测试它。
There are two main mechanisms for op and kernel registration:
tf.load_op_library()
Bazel BUILD文件如下，执行?bazel build ${BAZEL_ARGS[@]}?可以得到?tensorflow/core/user_ops/zero:zero_out.so?PS: 类似于执行了 上文中的g++
load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
tf_custom_op_library(
name = "zero_out.so", # target name
srcs = ["zero_out.cc"], # the list of the sources to compile,
)
得到so 文件后，tf.load_op_library 动态加载so作为 module 使用（可以参考python module 动态加载加载）。
import tensorflow as tf
# 返回一个 A python module, containing the (op对应的)Python wrappers for Ops defined in the plugin.
# Python Module，是一个 Python 文件，以 .py 结尾，包含了 Python 对象定义和Python语句
zero_out_module = tf.load_op_library('zero_out.so')
with tf.Session():
print(zero_out_module.zero_out([1,2,3,4,5])).eval() # eval 底层执行 session.run
zero_out_module.zero_out?可能仅用于演示，正规的ops 比如tf.matmul 等实现（对应到 编译tf生成的代码 gen_math_ops.py）会涉及到生成opDef （graphDef 的一部分）等逻辑。tensornet框架 自定义ops 示例
tensornet
/core
/kernels
/sparse_table_ops.cc # kernel实现
/ops
/sparse_table_ops.cc # REGISTER_OP
/BUILD
/tensornet
/core
/gen_sparse_table_ops.py # Bazel tf_gen_op_wrapper_py生成
BUILD 文件内容?tcnsorflow c++ op 生成 python调用接口
// 生成 生成op的lib，即so文件
cc_library(
name = "sparse_table_ops_kernels",
srcs = [
"kernels/sparse_table_ops_dummy.cc",
"ops/sparse_table_ops.cc",
],
hdrs = [
"//core/utility:semaphore",
],
linkstatic = 1,
deps = [
"@org_tensorflow//tensorflow/core:framework",
"@org_tensorflow//tensorflow/core:lib",
"@org_tensorflow//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
// 生成python的接口，即gen_sparse_table_ops.py 文件
tf_gen_op_wrapper_py(
name = "sparse_table_ops",
deps = [":sparse_table_ops_kernels"],
cc_linkopts = ['-lrt']
)
TensorFlow模型准实时更新上线的设计与实现
定制好 op 后，如何替换模型计算图中原生的 op 呢？TensorFlow 在模型保存时，会生成 meta_graph_def 文件，文件内容是采用类似 json 的格式描述计算图的结构关系。当加载此文件时，TensorFlow 会根据文件中描述的结构信息构建出计算图。可以修改模型保存的 meta_graph_def 文件，将其中的 op 替换为我们定制的 op，同时修改每个 node 的 input 和 output 关系，以修改 op 之间的依赖关系。PS： 当然这里说的替换原有的op
horovod
很多机器学习框架都会采用如下套路：shell脚本（可选），python端 和 C++端。
- Shell脚本是启动运行的入口，负责解析参数，确认并且调用训练程序；
- Python是用户的接口，引入了C++库，封装了API，负责运行时和底层C++交互；
- C++实现底层训练逻辑；
深度学习分布式训练框架horovod (2)一从使用者角度切入
引入库的作用是获取到 C++ 的函数，并且用 python 封装一下，这样就可以在 python 世界使用 C++代码了。比如下文，python 的 _allreduce 函数就会把功能转发给 C++，由 MPI_LIB.horovod_allreduce 完成。
def _allreduce(tensor, name=None, op=Sum, prescale_factor=1.0, postscale_factor=1.0,
ignore_name_scope=False):
if name is None and not _executing_eagerly():
name = 'HorovodAllreduce_%s' % _normalize_name(tensor.name)
return MPI_LIB.horovod_allreduce(tensor, name=name, reduce_op=op,
prescale_factor=prescale_factor,
postscale_factor=postscale_factor,
ignore_name_scope=ignore_name_scope)
## 初始化时执行
def _load_library(name):
"""Loads a .so file containing the specified operators.
"""
filename = resource_loader.get_path_to_datafile(name)
library = load_library.load_op_library(filename)
return library
# Check possible symbol not found error from tensorflow version mismatch
try:
MPI_LIB = _load_library('mpi_lib' + get_ext_suffix())
except Exception as e:
check_installed_version('tensorflow', tf.__version__, e)
raise e
else:
check_installed_version('tensorflow', tf.__version__)
|