如何在 Python 中导入 TensorFlow 中的自定义操作？答案

【问题标题】：How to make custom Op in TensorFlow importable in Python?如何在 Python 中导入 TensorFlow 中的自定义操作？
【发布时间】：2016-05-29 19:07:18
【问题描述】：

我已经为我的自定义 Op 实现了一个内核，并将其作为 custom_op.cc 放入 /tensorflow/core/user_ops。在 Op 内部，我负责所有注册工作，例如 REGISTER_OP 和 REGISTER_KERNEL_BUILDER。

然后我在 Python 中为这个 Op 实现了渐变，我把它和 custom_op_grad.py 放在了同一个文件夹中。我也在这里完成了所有注册 (@ops.RegisterGradient)。

我已经创建了 BUILD 文件，内容如下：

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
tf_custom_op_library(
        name = "custom_op.so",
        srcs = ["custom_op.cc"],
)

py_library(
        name = "custom_op_grad",
        srcs = ["custom_op_grad.py"],
        srcs_version = "PY2",
        deps = [
        ":custom_op_grad",
        "//tensorflow:tensorflow_py",
        ],
)

在那之后，我重建了 Tensorflow：

pip uninstall tensorflow
bazel clean
bazel build -c opt //tensorflow/tools/pip_package:build_pip_package
cp -r bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/__main__/* bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/
bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
pip install /tmp/tensorflow_pkg/tensorflow-0.8.0-py2-none-any.whl

当我在这一切之后尝试使用我的 Op 时，通过调用 tf.user_ops.custom_op 它告诉我模块没有它。

也许我还需要做一些额外的步骤？或者我对BUILD 文件做错了什么？

【问题讨论】：

标签： python machine-learning neural-network tensorflow bazel

【解决方案1】：

好的，我找到了解决方案。我刚刚删除了BUILD 文件，我的自定义操作已成功构建，并且可以使用tensorflow.user_ops.custom_op() 在Python 中导入。

要使用渐变，我必须将其代码直接放在tensorflow/python/user_ops/user_ops.py 中。不是最优雅的解决方案，但现在可以使用。

【讨论】：