TensorFlow自定义算子：如何以传引用方式传递可变对象参数

我刚踩过类似的坑，核心问题在于TensorFlow的普通tf.Tensor是不可变对象——你传递它给自定义算子时，拿到的是只读的内存视图，修改操作只会作用在临时副本上，原张量根本不会有变化。要让修改在算子执行后生效，得从算子实现和Python调用两个层面配合调整：

第一步：理解问题本质 #

Python里的tf.Tensor和字符串、整数一样属于不可变类型，函数参数传递的是"只读视图"，算子无法修改其底层内存。而tf.Variable是专门的可变存储容器，它的底层张量是可写的，这才是我们需要传递给算子的对象。

第二步：自定义算子的C++实现（声明可变输入） #

在算子的Compute函数中，必须显式获取可变的输入张量指针，而不是只读的。这里给你一个CPU算子的最小可运行示例：

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

using namespace tensorflow;

// 注册算子，定义输入输出类型和形状
REGISTER_OP("ModifyInput")
    .Input("input: float32")
    .Output("output: float32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0)); // 输出和输入形状保持一致
      return Status::OK();
    });

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

  void Compute(OpKernelContext* context) override {
    // 重点：获取可变的输入张量，第三个参数true允许修改输入
    Tensor* input_tensor = nullptr;
    OP_REQUIRES_OK(context, context->mutable_input(0, &input_tensor, true));
    
    // 获取扁平的张量视图，方便直接修改元素
    auto input_flat = input_tensor->flat<float>();
    
    // 这里写你的修改逻辑，比如把第一个元素改成99.0
    input_flat(0) = 99.0f;

    // 可选：直接输出修改后的输入张量（如果需要同时拿到输出的话）
    context->set_output(0, *input_tensor);
  }
};

// 注册CPU版本的算子
REGISTER_KERNEL_BUILDER(Name("ModifyInput").Device(DEVICE_CPU), ModifyInputOp);

第三步：Python端调用（必须用tf.Variable） #

把你之前的普通张量替换成tf.Variable，这样算子才能修改它的底层存储。示例代码如下：

import tensorflow as tf
import numpy as np
from PIL import Image

# 加载编译好的自定义算子（假设已编译为modify_input_op.so）
tf.load_op_library('./modify_input_op.so')

# 读取图像并转为tf.Variable（关键步骤！）
img = Image.open("your_image_path.jpg")
img_np = np.array(img)[:, :, 0:3].flatten().astype('float32')
# 用tf.Variable包装，这是唯一能被算子修改的可变容器
img_var = tf.Variable(img_np)

print("修改前第一个元素值:", img_var.numpy()[0])

# 调用自定义算子
output = tf.raw_ops.ModifyInput(input=img_var)

# 验证修改是否生效
print("修改后原Variable的第一个元素值:", img_var.numpy()[0])
print("输出张量的第一个元素值:", output.numpy()[0])

额外注意事项 #

• 如果是GPU算子，逻辑完全一致，只需把DEVICE_CPU改成DEVICE_GPU，处理GPU内存时用CUDA相关API即可。
• 绝对不要尝试修改普通tf.Tensor——TensorFlow的设计就是张量不可变，强行操作会触发内存访问错误，或者根本不生效。
• 如果你的场景不适合用tf.Variable，也可以用tf.Tensor配合tf.stop_gradient，但tf.Variable是最直观的解决方案。

内容的提问来源于stack exchange，提问作者Lodrik