仓库地址：https://github.com/tensorflow/custom-op

将仓库拉到本地，假设自定义算子名字是inno，复制zero_out文件夹重命名为tensorflow_inno，功能是（1）保留input的第一个数字；（2）input其他位置上的元素都变为原来2倍。

经过文件名修改，自定义算子tensorflow_inno的目录树如下：

.
├── BUILD
├── cc
│   ├── kernels（实现）
│   │   └── inno_kernels.cc
│   └── ops（注册）
│       └── inno_ops.cc
├── __init__.py
└── python
    ├── __init__.py（空）
    └── ops
        ├── __init__.py（空）
        ├── inno_ops.py（调用.so动态链接库）
        ├── _inno_ops.so（编译后产生）
        └── inno_ops_test.py（调用inno_ops.py测试）

1. C实现（cc/目录下）

目录中，cc/下的所有程序都调用tensorflow的c接口，cc/ops注册了inno算子，cc/kernels实现了inno算子。

├── cc
│  ├── kernels（实现）
│  │  └── inno_kernels.cc
│  └── ops（注册）
│    └── inno_ops.cc

先在ops目录下的inno_ops.cc中注册算子：

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

using namespace tensorflow;

REGISTER_OP("Inno")
    .Input("inp: int32")
    .Output("out: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
});

这里定义了算子的名称是Inno，输入名称inp，输出名称out，输出的shape（存疑：等于输入的shape？）。

在kernels下的inno_kernels.cc中实现算子（注意命名需要采用CamelCase，驼峰拼写法）：

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

using namespace tensorflow;

class InnoOp : public OpKernel {
 public:
  explicit InnoOp(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<int32>();

    // 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<int32>();

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

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

到此为止，已经定义好了tensorflow的算子。

2. Python Package（python/目录下）

后面我们会编译Python包（.whl），安装后目录和这里一样，因此可以推断这里是在书写Python Package的脚本。其编译过程是：c脚本->.so动态链接库->Python调用->.whl包。

└── python
    ├── __init__.py（空）
    └── ops
        ├── __init__.py（空）
        ├── inno_ops.py（调用.so动态链接库）
        ├── _inno_ops.so（编译后产生）
        └── inno_ops_test.py（调用inno_ops.py测试）

inno_ops.py里面调用了.so动态链接库，定义了一个接口inno_out。Python安装完成后，可以直接调用inno_out。

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from tensorflow.python.framework import load_library
from tensorflow.python.platform import resource_loader

inno_ops = load_library.load_op_library(
    resource_loader.get_path_to_datafile('_inno_ops.so'))
inno_out = inno_ops.Inno

inno_ops_test.py调用了inno_ops.py测试inno_out算子，使用以下语句判断输入和输出是否相等：self.assertAllClose(inno_out([[1, 2], [3, 4]]), np.array([[1, 4], [6, 8]]))。代码如下：

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

from tensorflow.python.platform import test
try:
  from tensorflow_inno.python.ops.inno_ops import inno_out
except ImportError:
  from inno_ops import inno_out


class InnoTest(test.TestCase):

  def testinno(self):
    with self.test_session():
      self.assertAllClose(
          inno_out([[1, 2], [3, 4]]), np.array([[1, 4], [6, 8]]))


if __name__ == '__main__':
  test.main()

3. 其他文件

1. tensorflow_inno目录下，有一个BUILD文件。注意将其中文件名改为自己的。
2. tensorflow_inno上级目录（自定义算子的根目录），Makefile文件中的路径需要修改为自定义的路径。
3. MANIFEST.in文件里面有两行文字，需要改为自定义名字
   recursive-include tensorflow_zero_out *.so
   recursive-include tensorflow_time_two *.so
4. setup.py文件中定义了包的名称、版本号，可以保持默认。

4. 编译、安装

在custom-op文件夹（tensorflow-inno上级目录）中，编译算子得到.so文件和.whl文件。执行下面三行命令即可安装完成。

(zhouyi) inno-suanfa@ubuntu:~/zhuwk/custom-op$ make inno_pip_pkg
(zhouyi) inno-suanfa@ubuntu:~/zhuwk/custom-op$ cd artifacts/
(zhouyi) inno-suanfa@ubuntu:~/zhuwk/custom-op/artifacts$ pip install tensorflow_custom_ops-0.0.1-cp38-cp38-linux_x86_64.whl 

进入Python测试算子功能，发现功能正常。

(zhouyi) inno-suanfa@ubuntu:~/zhuwk/custom-op/artifacts$ python
Python 3.8.13 (default, Mar 28 2022, 11:38:47) 
[GCC 7.5.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow_inno as tfno
>>> import tensorflow as tf
>>> tfno.inno_out(inp=tf.constant([[1,2,3],[4,5,6]]),name="inno_op")
<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[ 1,  4,  6],
       [ 8, 10, 12]], dtype=int32)>

5. 保存为单算子pb模型

代码：

import tensorflow as tf
import tensorflow_inno as tfno
class newop(tf.keras.models.Model):
    def __init__(self):
        super(newop, self).__init__()
    def call(self, x):
        x = tfno.inno_out(inp=x, name="inno_op")
        return x
opmodel = newop()
opmodel.predict(tf.constant([[1,2],[3,4]]))
print(opmodel.summary())
import numpy as np
full_model = tf.function(lambda x: opmodel(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec((2,2), np.int32))
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()
# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                  logdir="./",
                  name="inno_op.pb",
                  as_text=False)

输出：

Model: "newop"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
Total params: 0
Trainable params: 0
Non-trainable params: 0
_________________________________________________________________
None
'./inno_op.pb'

使用netron可视化，节点属性如下