I have implemented a kernel for my custom Op, and put it into /tensorflow/core/user_ops as custom_op.cc. Inside the Op I do all the registering stuff, like REGISTER_OP and REGISTER_KERNEL_BUILDER.
Then I implemented gradient for this Op in Python, and I put it in the same folder as custom_op_grad.py. I did all the registering here as well (@ops.RegisterGradient).
I have created the BUILD file, with the following content:
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",
],
)
After that, I rebuild 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
When I try to use my Op after all this, by calling tf.user_ops.custom_op it tells me that module doesn't have it.
Maybe there are some additional steps I have to do? Or I am doing something wrong with the BUILD file?
Ok, I found the solution. I just removed the BUILD file, and my custom Op was successfully built and was importable in Python using tensorflow.user_ops.custom_op().
To use the gradient I had to put it's code directly inside the tensorflow/python/user_ops/user_ops.py. Not the most elegant solution, but working for now.
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments