TensorFlowのカスタムOpをPythonでインポートするには？

Question

自分のカスタムOp用にカーネルを実装し、/tensorflow/core/user_opsにcustom_op.ccというように入れました。オペラの中では、REGISTER_OPとREGISTER_KERNEL_BUILDERのようなものをすべて登録します。

次に、このOpのための勾配をPythonで実装しました。これをcustom_op_grad.pyと同じフォルダに入れました。私もここで登録しました（@ops.RegisterGradient）。

私は次の内容で、ビルドファイルを作成している

：私は呼び出すことで、すべてのこの後の私のオペアンプを使用しようとすると

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 

：その後

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を再構築tf.user_ops.custom_opモジュールはそれを持っていないと私に伝えます。

多分、追加する必要があるいくつかの手順がありますか？または、BUILDファイルで何か問題が起きていますか？

Answer 1

私は解決策を見つけました。 BUILDファイルを削除したところ、私のカスタムOpは正常に構築され、tensorflow.user_ops.custom_op()を使用してPythonでインポート可能でした。

グラデーションを使用するには、コードを直接tensorflow/python/user_ops/user_ops.pyの中に入れなければなりませんでした。最も洗練されたソリューションではなく、今のところ働いています。