无法在Tensorflow中使用新的tf.Operation在Python shell中工作

Question

我试图在张量流中添加一个新的整数矩阵乘法OP，我无法将其成功注册为tf操作，因此可以在python中将其称为tf.intmatmul。

我做的步骤： 1）在math_ops.cc文件中添加了新的REGISTER_OP - IntMatMul。

2）在核心/内核路径中添加了一个新的内核实现 - int_matmul_op.cc和相应的头文件 - int_matmul_op.h

3）我在core / kernels / BUILD文件中添加了OP的依赖项。这将为此OP添加内核链接。

4）在Python Wrapper文件中添加了此OP的定义（作为'intmatmul'），即python / ops / math_ops.py - 此文件调用gen_math_ops.int_mat_mul

5）使用Bazel从源代码重新构建并使用pip包重新安装Tensorflow。

但是当我尝试将此OP用作tf.intmatmul时，我收到一条错误消息，指出模块未定义。 我现在不确定我在这里缺少什么。是否有任何缺失的链接？ 我是否还需要在核心/ BUILD文件中添加任何OP链接？

Answer 1

更新：

所以事实证明这比预期的要复杂得多。这些是必须考虑的事情：

• 显然，为了将函数公开为公共API（即tf.级别），其名称必须在模块的开头列在其前面带有@@的文档字符串中。请查看示例math_ops.py。
• 内核定义必须绝对正确才能将操作反映为操作，即使您仍然可以从内部模块访问它（例如，执行from tensorflow.python.ops import math_ops）。

-

如the docs所示，必须在C ++中使用CamelCase标识符注册的操作名称是＆＃34;翻译＆＃34;在Python中使用snake_case。请改为tf.int_mat_mul。

作为旁注，该教程提供了实现自定义操作的其他指导，无需从源代码重新编译TensorFlow，而是从自定义库加载。