Question

我尝试向tensorflow添加自定义操作，但我无法从python加载它。问题类似于github中关闭的issue，但那里的解决方案并没有解决我的问题。

操作系统：macOS 10.12

已安装的CUDA和cuDNN版本：无

从源安装TensorFlow 0.11.0。

我按照add new op教程添加了zero_out.cc文件：

#include "tensorflow/core/framework/op.h"

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output(0) = input(0);
  }
};

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

和bazel BUILD文件：

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"]
)

然后我跑：

bazel build -c opt //tensorflow/core/user_ops:zero_out.so

输出：

INFO: Waiting for response from Bazel server (pid 28589)...
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: 5.115s, Critical Path: 0.00s

生成的共享库位于bazel-bin中。当我试图像这样加载它时：

tf.load_op_library('/Users/dtong/code/data/tensorflow/bazel-bin/tensorflow/core/user_ops/zero_out.so')

结果：

python(41716,0x7fffb7e123c0) malloc: *** error for object 0x7f9e9cd2de18: pointer being freed was not allocated
*** set a breakpoint in malloc_error_break to debug

Answer 1

好吧，我找到了解决方案。而不是通过bazel构建用户操作，使用g ++。

g++ -v -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC -I $TF_INC -O2 -undefined dynamic_lookup -D_GLIBCXX_USE_CXX11_ABI=0

它会起作用。原因似乎是我的gcc版本太高了（v6.2.0）。

-D_GLIBCXX_USE_CXX11_ABI=0来自official site的概念＆＃39;构建Op库＆＃39;部分。

但我还是找不到使用bazel的解决方案。我在GitHub中发了一个issue。

无法在tensorflow中加载自定义的op共享库

1 个答案: