python - 无法获得新的 tf.Operation 在 Tensorflow 的 Python shell 中工作

Question

我正在尝试在 tensorflow 中添加一个新的整数矩阵乘法 OP，但我无法将其成功注册为 tf 操作，因此无法在 python 中将其称为 tf.intmatmul。

我做的步骤: 1) 在 math_ops.cc 文件中添加了一个新的 REGISTER_OP - IntMatMul。

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

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

4) 在 Python 包装器文件中添加了此 OP 的定义(如“intmatmul”)，即 python/ops/math_ops.py - 此文件调用 gen_math_ops.int_mat_mul

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

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

Answer 1

更新:

所以这比预期的要复杂。这些是必须考虑的事项:

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

--

作为the docs表明，必须在 C++ 中使用 CamelCase 标识符注册的操作名称在 Python 中被“翻译”为 snake_case。请尝试使用 tf.int_mat_mul。

附带说明一下，该教程提供了额外的指导来实现自定义操作，而无需从源代码重新编译 TensorFlow，而是从自定义库加载它。