用cmake工具 自定义tensorflow算子

自定义tensorflow算子有两种方式：

1. 下载tensorflow源码，在源码中添加cpp代码，然后编译安装。

这种方法的优点是自定义的算子一起被打包进tensorflow，形式比较统一。当然缺点也很明显，其他人想调用该算子时，需要重新安装tensorflow

2. 代码独立编译成动态链接库，然后再tensorflow中调用。

这种方式的优点是非常灵活，编译开发的工作量比较小。

还有一种方式是把算子编译成第三方库，这种方式一般比较少用，所以下面就第二种方式讲解一下自定义算子的过程。

1. 环境准备：

首先配置环境，为了避免和本地的环境混淆，建议在docker容器中测试，docker的基本使用在docker 中镜像和容器区别_kangshuangzhu的博客-CSDN博客_docker镜像和容器的区别中有简单介绍。tensorflow已经为我们准备好了专门用于自定义算子的镜像，拉取镜像：

docker pull tensorflow/tensorflow:2.5.0-custom-op-ubuntu16

运行，进入容器：

docker run -it tensorflow/tensorflow:2.5.0-custom-op-ubuntu16 /bin/bash

 该容器已经安装好了Python和tensorflow2.5

2. 算子开发

创建一个新的文件夹 custom-op

mkdir custom-op
cd custom-op

创建三个文件，下面会介绍代码的用途

zero_out_ops.cc

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

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();
    });

 zero_out_kernels.cc

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

CMakeLists.txt 

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(tfcustomop_project)

# Define our library target
add_library(zero_out SHARED zero_out_ops.cc zero_out_kernels.cc)

# 在命令行中运行python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))' 
# 可以得到tensorflow的include路径以及 编译器的链接库版本 分别写进include_directories 和 set 配置中，注意删掉路径最前面的-I
include_directories(/usr/local/lib/python3.6/dist-packages/tensorflow/include)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -D_GLIBCXX_USE_CXX11_ABI=0")


# 在命令行中执行 python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))'
# 可以得到tensorflow动态库路径，把两个路径拼接到一起，得到完成路径填到target_link_libraries
target_link_libraries(zero_out "/usr/local/lib/python3.6/dist-packages/tensorflow/libtensorflow_framework.so.2")

3. 编译动态链接库

在custom-op中创建新文件夹build，然后进入build，执行cmake.. 和make

mkdir build
cmake ..
make

 可以看到生成了了一个动态链接库，zero_out.so

4. 测试算子

python -c 'import tensorflow as tf; aa = tf.load_op_library("./build/libzero_out.so"); print(aa.zeroout([1,1,1,1]))'

可以看到输出结果[1,0,0,0]

代码解释

op和kernel的作用

上面是一个非常简单的，没有用到gpu的算子开发。一般一个算子包括ops和kernel两部分组成，当然，你也可以把他们放在一个代码中实现。

op定义了运算名称，输入名称，输出名称，输入数据类型，输出数据类型

kernel则定义了具体的运算代码。

在应用中，op用于静态图中创建图时生成op，kernel用于运行静态图时创建真实的运算操作。

可以看到在kernel中，运算操作是除了tensor的第一个元素，其他全都置为0

函数名称

需要注意的是，op和kernel中的函数名称要一直，否则会出现未定义op或者该函数没有定义kernel的错误。仔细观察会发现，在op和kernel中定义的函数名称都是ZeroOut， 但是在Python中调用的时候用的是zero_out。经过本人实验，这里确实会把驼峰明明强制转换成下划线命名。此外操作名称必须首字母大写，而且不能和库中已经注册的其它操作重名。

参考文档：【架构分析】TensorFlow 自定义算子实现原理分析_HaoBBNuanMM的博客-CSDN博客_tensorflow 算子