tensorflow 源代码分析（4）-自定义op

添加自定义op

本文介绍如何在tensorflow中增加一个新得OP以及op对应得kernel，还有如何编译和使用新增加得op。

1.新增加op对应得kernel实现，新的op ZeroOutOp类继承自OpKernel，重新实现Compute接口，在该接口内实现op对应得kernel（tensorflow/core/user_ops/zero_out.cc cpu设备上的实现）：

REGISTER_OP 宏来定义 Op 的接口，定义op得输入输出和属性：

Attr：该op接口内得属性值，例如：preserver_index 就是该op得属性，以后Zeroout会自带该属性

Input：该op得输入

Output：该op得输出

SetShapeFn：控制输入输出前后得shape变化

通过REGISTER_KERNEL_BUILDER向tensorflow系统注册op，该实现对应得是CPU：

2.上面三项完成后即可编译该op，有两种方式编译：1.bazel编译，2.g++编译：

1）bazel编译：

     Bazel编译需要用到一个BUILD文件：

   Tf_custom_op_library 声明编译成一个lib库

   Name：指定lib库的名称

   Srcs：指定源文件

   执行：bazel build --config opt //tensorflow/core/user_ops:zero_out.so开始编译so。

编译成功；

2）g++编译：

TF_CFLAGS=($(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )

TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )

g++ -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

通过g++编译过程，最终会生成so

3）此时的接口是C++实现的，可以通过编译成库的方式供python调用（swig），也可在tensorflow源代码中直接添加对应的python接口，这样可以重新编译tenosrflow，然后就可以直接通过python调用自定的op。

tensorflow/tensorflow/core/BUILD中下面rule把user_ops下面的op生成lib库：

通过

把user_ops_op_lib编译到ops库中

通过这些设置后zero_out.cc中的op会自动在bazel-genfiles/tensorflow/python/ops/ gen_user_ops.py文件中实现zero_out的python包装器。

接下通过gen_user_ops.py实现zero_out的python API：

在tensorflow/tensorflow/python/user_ops/user_ops.py中添加：

可以看到_gen_user_ops是从tensorflow.python.ops.gen_user_ops中引用过来的，通过调用_zero_out重新实现zero_out。

在这里应注意接口名字得变化，在python中使用得是pep8格式而在c++中使用得是驼峰式的命名规则，所以在查找对应的接口时需要注意名字的变换。

tensorflow/tensorflow/python/BUILD

主要是把deps中的ops打包成一个python lib库。

tensorflow/tensorflow/python/ops/hidden_ops.txt中添加：

主要是隐藏zeroout自动生成的代码。（todo：不清楚为什么要隐藏，不隐藏不能正常调用）

重新编译tensorflow源代码，生成*.whl文件，使用pip install *.whl 即可安装自定义版本的tensorflow。

 

3.上面介绍的是如何添加一个op对应的cpu上的实现，因为tensorflow是异构的计算平台，所以该op也需要支持GPU上运行，tensorflow是通过文件名的不同区分在不同设备上的实现的，cpu设备的实现一般以*.cc结尾，gpu设备的实现一般以*.cu.cc结尾。所以新增支持的设备就需要实现新的op kernel的实现。

因为该op实现的比较简单，如果该op需要实现负责的功能，比如需要实现矩阵的运算，就需要调用支持矩阵运算的库，matmul op实现矩阵的运算调用的eigen，eigen是一个高层次的C ++库，有效支持线性代数，矩阵和矢量运算，数值分析及其相关的算法，但是只支持cpu设备，在gpu设备上实现的matmul调用的是cublas的加速算法库。

至于后端该如何选择哪种实现，会根据用户自己选择的在哪个设备上执行该op，自己找到对应的kernel实现。