python - 如何在自定义操作中改变 Tensorflow 变量？

Question

我正在尝试修改简单的 Adding a New Op这样它就不会创建一个新的张量作为返回值，但它实际上会改变输入张量并返回它。我知道这是可能的，因为 scatter Op 正在做同样的事情，但是查看 scatter Op 源代码，鉴于我缺乏 C++ 经验，我不知道该怎么做。

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.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();
    });



#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
    Tensor input_tensor = context->mutable_input(0, true);
    auto input = input_tensor.flat<int32>();

    // We always return the input ref.
    context->forward_ref_input_to_ref_output(0, 0);

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

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

如果我编译上面的代码并运行一个简单的 Python 脚本来测试它，我会得到以下错误:

Python(14820,0x700003a7e000) malloc: *** error for object 0x7fd5c45a5a88: pointer being freed was not allocated
*** set a breakpoint in malloc_error_break to debug

我需要更改代码中的哪些内容才能满足我的需求？

Answer 1

我觉得你最好修改一下抓取输入输出的过程。实际上根据你的REGISTER_OP，它不是引用输入，所以

context->mutable_input(0, true)

会是

context->input(0)

此外，设置输出将更改为

context->set_output(0, context->input(0))

我觉得设置输出后就可以了。