Tensorflow add a new op, could not import from python

982 views Asked by At

I did as shown in the tensorflow document: https://www.tensorflow.org/how_tos/adding_an_op/#attrs

  1. get the source code from https://github.com/tensorflow/tensorflow
  2. add a zero_out.cc file under tensorflow/core/user_ops/
  3. add a BUILD file under tensorflow/core/user_ops/
  4. run in terminal: $ bazel build -c opt //tensorflow/core/user_ops:zero_out.so

INFO: Found 1 target...
Target //tensorflow/core/user_ops:zero_out.so up-to-date:
bazel-bin/tensorflow/core/user_ops/zero_out.so
INFO: Elapsed time: 0.125s, Critical Path: 0.00s

But when I run the following code in python, error comes:

Python 2.7.6 (default, Jun 22 2015, 17:58:13)
[GCC 4.8.2] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcuda.so.1 locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcurand.so locally
>>> zero_out_module = tf.load_op_library('zero_out.so')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/yc/tfenv/local/lib/python2.7/site-packages/tensorflow/python/framework/load_library.py", line 64, in load_op_library
    None, None, error_msg, error_code)
tensorflow.python.framework.errors_impl.NotFoundError: zero_out.so: cannot open shared object file: No such file or directory
1

There are 1 answers

0
pfm On

Bazel would generate the op library into the bazel-bin folder.

As you used bazel to build ot, you should try to load it from there:

zero_out_module = tf.load_op_library('bazel-bin/tensorflow/user_ops/zero_out.so')