如何在TensorFlow C++中定义跨Session.run()持变量的自定义有状态Op？

Alright, let's tackle this problem. You're right that just calling REGISTER_OP().SetIsStateful() isn't enough — that only flags the op as stateful to TensorFlow's optimization pass (so it doesn't get folded or duplicated unnecessarily), but it doesn't actually create or persist any internal state. To make your op retain values across Session.run() calls, you need to use TensorFlow's resource management system in the op's kernel implementation.

Here's a complete, step-by-step solution:

1. C++端：定义Op与实现状态感知的Kernel #

First, we'll create a C++ resource class to hold our persistent foo value, then define the op and its kernel that uses this resource.

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

// 自定义资源类，用于存储持久化的状态
class CounterResource : public tensorflow::ResourceBase {
 public:
  CounterResource() : foo_(0) {}

  // 状态更新方法：每次调用foo加1
  void Increment() { foo_++; }

  // 获取当前状态值
  int GetValue() const { return foo_; }

  // 调试用的字符串输出
  std::string DebugString() const override {
    return tensorflow::strings::StrCat("CounterResource(foo=", foo_, ")");
  }

 private:
  int foo_; // 我们要持久化的状态变量
};

// 注册自定义Op
REGISTER_OP("IncrementCounter")
    .Output("current_value: int32")
    .SetIsStateful() // 标记为有状态Op，避免被优化器误处理
    .SetDoc(R"doc(
A stateful op that increments an internal counter each time it runs.
current_value: The counter's value after incrementing.
)doc");

// 实现Op的Kernel逻辑
class IncrementCounterOp : public tensorflow::OpKernel {
 public:
  explicit IncrementCounterOp(tensorflow::OpKernelConstruction* context)
      : OpKernel(context) {
    // 从设备的资源管理器中查找或创建我们的CounterResource
    OP_REQUIRES_OK(context,
                   context->resource_manager()->LookupOrCreate<CounterResource>(
                       context->device()->name(), "persistent_counter", &counter_,
                       [](CounterResource** resource) {
                         *resource = new CounterResource;
                         return tensorflow::Status::OK();
                       }));
  }

  void Compute(tensorflow::OpKernelContext* context) override {
    // 更新状态：将foo加1
    counter_->Increment();

    // 分配输出张量并写入当前值
    tensorflow::Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, tensorflow::TensorShape(),
                                                     &output_tensor));
    output_tensor->scalar<int>()() = counter_->GetValue();
  }

 private:
  tensorflow::core::RefPtr<CounterResource> counter_; // 持有资源的引用
};

// 将Kernel注册到CPU设备（如果需要GPU支持，可添加对应的DEVICE_GPU注册）
REGISTER_KERNEL_BUILDER(Name("IncrementCounter").Device(tensorflow::DEVICE_CPU),
                        IncrementCounterOp);

2. 编译生成动态链接库 #

Compile the C++ code into a shared library that TensorFlow can load. For Linux/macOS, use a command like this (replace paths with your local TensorFlow headers and library locations):

g++ -std=c++11 -shared increment_counter_op.cc -o increment_counter_op.so -fPIC -I/usr/local/include/tensorflow -L/usr/local/lib -ltensorflow_framework

For Windows, you'd compile to a .dll instead, using MSVC or MinGW.

3. Python端：加载Op并测试状态保持 #

Now load the compiled library in Python and test that the state persists across run() calls:

import tensorflow as tf

# 加载自定义Op库
custom_op = tf.load_op_library('./increment_counter_op.so')

# 测试状态持久化
with tf.Session() as sess:
    # 第一次调用：foo从0→1，输出1
    print(sess.run(custom_op.increment_counter()))  # 输出: 1
    # 第二次调用：foo从1→2，输出2
    print(sess.run(custom_op.increment_counter()))  # 输出: 2
    # 第三次调用：foo从2→3，输出3
    print(sess.run(custom_op.increment_counter()))  # 输出: 3

为什么单独用SetIsStateful()没效果？ #

SetIsStateful() only tells TensorFlow that the op has side effects or depends on persistent state, so it won't apply optimizations like constant folding or removing redundant executions. But it doesn't create any state storage on its own. The real work happens in the kernel, where we use TensorFlow's ResourceManager to create a persistent resource that lives across run() calls, tied to the device (CPU/GPU) the op runs on.

Alternative: Python-only "stateful op" (simulated) #

If you don't want to write C++ code, you can simulate this behavior using a tf.Variable wrapped in a custom layer. Note this isn't a true op-level state, but it achieves the same user-facing result:

import tensorflow as tf

class IncrementCounter(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        # 定义一个不可训练的变量来存储状态
        self.foo = tf.Variable(0, trainable=False, dtype=tf.int32)
    
    def call(self, inputs=None):
        # 每次调用时递增foo
        self.foo.assign_add(1)
        return self.foo

# 测试
counter = IncrementCounter()
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(counter()))  # 输出: 1
    print(sess.run(counter()))  # 输出: 2
    print(sess.run(counter()))  # 输出: 3

But for a true custom op with internal state (e.g., if you need to implement low-level logic in C++), the first approach is the correct way to go.

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