TensorFlow 增加自定义运算符 ——转自 慢慢学TensorFlow 微信公众号

TensorFlow 支持多种常用运算符（Operations，简称 Op），如矩阵乘 matmul、卷积 conv2d 等。

由这些运算符可以构建复杂的计算图（Graph）。核函数（Kernel）是运算符的特定实现，可以运行在特定类型设备上（如 CPU 或 GPU ）。

TensorFlow 通过注册机制来确定支持的运算符集合和核函数，便于进一步扩展。

今天给一个利用 C++ 编写 TensorFlow Op 的例子，读者可以很快编写自定义的运算符实现自己想要的计算。

首先在一台安装了 TensorFlow 的机器上（笔者用的是 Ubuntu 14.04 系统 + GTX 1080）编写如下 C++ 代码：

上述代码首先包含编译 TensorFlow Op 必要的头文件（op.h 和 op_kernel.h），之后利用 REGISTER_OP 宏注册 TensorFlow Op，名字为 "ZeroOut"，输入参数为 int32 类型的 Tensor，输出也是 int32 类型 Tensor。

该 Op 的计算过程由相应的 Kernel 函数完成，利用 C++ 派生类的特性，从 OpKernel 基类派生了 ZeroOutOp 类，该类中的 Compute 函数最为关键，从输入 Tensor 获取输入数据，并为输出 Tensor 分配空间。利用一个 for 循环实现置零运算。注意到输出 Tensor 的第一个元素重新赋值为输入 Tensor 的第一个元素值。

保存为 zero_out.cpp，然后使用以下命令编译：

注意，这里 -I 后面的路径是我机器上安装的 TensorFlow 的包含路径，读者应自行验证是否兼容。

小贴士：依次运行 python，import tensorflow，help(tensorflow) 就能看到 TensorFlow 实际安装的路径。

编译成功后，当前目录下多了一个 zero_out.so 文件，如下图所示：

这个动态链接库可以动态载入到 TensorFlow 运行环境。我们用 Python 编写一个测试例程，代码如下：

在当前目录保存为 test.py，之后运行 python test.py，结果如图：

可以看到，我们的测试例程中输入 Python 数组 [1, 2, 3, 4, 5]，经过 zero_out 运算符之后，输出为 [1, 0, 0, 0, 0]，可见实现了预期功能。

读者不妨对代码做些修改，实现 Tensor 元素的倒序。

掌握了今天的内容，读者可以很方便地编写自己特殊的运算符，有关运算符的更多内容，请参考：https://www.tensorflow.org/how_tos/adding_an_op/

扫描下面二维码关注此公众号！