0

/tensorflow/core/user_opsカスタム Op 用のカーネルを実装し、それをasに入れましたcustom_op.cc。Op 内では、 や などのすべての登録作業を行いREGISTER_OPますREGISTER_KERNEL_BUILDER。

次に、この Op のグラデーションを Python で実装し、 と同じフォルダーに配置しました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ますか?

4

1 に答える 1

0

わかりました、私は解決策を見つけました。ファイルを削除したところBUILD、カスタム Op が正常にビルドされ、tensorflow.user_ops.custom_op().

グラデーションを使用するには、そのコードを .xml 内に直接配置する必要がありましたtensorflow/python/user_ops/user_ops.py。最もエレガントなソリューションではありませんが、今のところ機能しています。

于 2016-05-30T18:12:03.073 に答える