TensorFlow装饰器@tf.function异常行为求解：计数器不递增问题

首先明确@tf.function的核心运行逻辑：被装饰的函数会在首次被调用时将Python代码转换为TensorFlow静态计算图（该阶段称为追踪阶段），后续调用只要输入签名没有变化，都会直接复用已构建的静态图，不再执行原始的Python代码逻辑。同时，对于Python原生类型的可变变量，tf.function在追踪阶段会将其值作为常量固化到图中，不会在后续图执行阶段同步Python侧的变量修改。

第一种写法（step被装饰，循环在Python侧）的运行逻辑 #

# 对应第一种写法的核心结构
@tf.function 
def step(self):
    self.count.increment()

def start(self):
    for i in range(10):
        self.step()

运行流程：

1. 第一次调用step()时触发图追踪：
   • 执行Python代码self.count.increment()，Python原生的count值从0变为1，同时生成打印1的图节点
   • 由于count是Python原生int，不是tf.Variable，TensorFlow会将当前的count值1作为常量固化到静态图中，最终生成的静态图逻辑只有「打印常量1」，没有动态递增的逻辑
2. 后续9次调用step()：
   • 输入签名没有变化，直接复用已构建的静态图，不会再执行Python侧的increment逻辑，所以每次都输出1，计数器不会递增。

第二种写法（start被装饰，循环在图侧）的运行逻辑 #

# 对应第二种写法的核心结构
@tf.function 
def start(self):
    for i in range(10):
        self.count.increment()

运行流程：

1. 第一次调用start()时触发图追踪：
   • 函数内部的for i in range(10)是Python固定次数循环，AutoGraph会在追踪阶段完整展开这个循环，依次执行10次increment调用
   • 追踪阶段依次执行：count从0变1打印1、变2打印2……直到变10打印10，所有打印操作和递增逻辑都被作为不同的节点固化到静态图中
2. 静态图构建完成后执行，会依次输出1到10的递增序列，符合预期。

补充说明 #

如果要让第一种写法也正常工作，只需要将Count类中的计数器改为tf.Variable类型即可：

class Count:
    def __init__(self):
        self.count = tf.Variable(0) # 改为TensorFlow内置可变状态类型
    def increment(self):
        self.count.assign_add(1) # 使用图内操作修改变量
        tf.print(self.count)

tf.Variable是TensorFlow图内置的状态类型，tf.function不会将其值固化为常量，会在每次图执行时正确处理它的更新逻辑，第一种写法下也能得到递增的输出。

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