如何向tensorflow添加操作(教程)

2024-09-29 17:15:04 发布

您现在位置:Python中文网/ 问答频道 /正文

我想在tensorflow中添加我自己的操作。 所以我读了https://www.tensorflow.org/extend/adding_an_op#use_the_op_in_python。 然后我试着增加新的操作,但没用。 错误信息:tensorflow.python.framework.错误_impl.NotFound错误:dlopen(零_出去。所以,6):找不到图像

我把这个源代码放到tensorflow/tensorflow/core/user_ops naming zero_输出.cc. 在

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"
using namespace tensorflow;
REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

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>();
"zero_out.cc" 35L, 1286C

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);

我只是用教程写同样的东西。 然后,我用相同的路径创建了构建文件

^{pr2}$

然后,在终端(path~/tensorflow)中,我键入了bazel构建命令。在

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

我试着运行我的代码。在

import tensorflow as tf
zero_out_module = tf.load_op_library('zero_out.so')
with tf.Session(''):
  zero_out_module.zero_out([[1, 2], [3, 4]]).eval()

但只有错误的按摩。在

error massage: tensorflow.python.framework.errors_impl.NotFoundError: dlopen(zero_out.so, 6): image not found

我想知道怎么了。在

谢谢你的阅读。在


Tags: thecoreinputoutputtensorflowcontextframeworkpublic

热门问题