TensorFlow 2.9.1：tf.function下如何动态修改层trainable属性

这个问题本质是tf.function的静态图缓存机制导致的：第一次执行被@tf.function装饰的call、train_step方法时，TensorFlow会根据执行当下的Python侧状态（包括各层trainable属性值、对应的trainable_weights列表）追踪计算流程，编译成静态图后缓存，后续调用会直接复用缓存图，不会重新读取Python侧更新后的属性值。
训练前就设置好trainable再启动训练，第一次图追踪拿到的就是正确的属性值，所以运行符合预期；训练中途修改trainable时，旧的静态图已经缓存完成，属性修改不会同步到已编译的图里，自然结果异常。移除@tf.function后全程走动态图执行，每次都会实时读取最新的Python属性，所以结果正常。

不需要移除tf.function，只需要在修改trainable属性后，显式让被装饰的方法感知状态变化、重新编译静态图即可，以下两种方案可根据场景选择：

方案1：切换trainable后手动清除图缓存 #

该方案实现简单，适合训练阶段划分明确、仅在阶段切换时修改trainable的场景，重追踪仅在切换时触发一次，不会影响训练循环内的执行性能。
在修改层的trainable属性、重新compile之后，手动清除train_step和call方法已缓存的静态图，强制下次调用时根据最新属性重新编译图即可，修改后的train方法示例：

def train(self, inputs, targets, num_epochs=5000):
    self.cube_layer.trainable = False
    self.compile(optimizer=self.optimizer)
    for epoch in range(num_epochs):
        loss = self.train_step(inputs, targets)
    
    self.cube_layer.trainable = True
    self.compile(optimizer=self.optimizer)
    # 清除已缓存的静态图，强制下次调用时重追踪
    self.train_step = tf.function(self.train_step.python_function)
    self.call = tf.function(self.call.python_function)
    for epoch in range(num_epochs):
        loss = self.train_step(inputs, targets)
        
    print("Loss: " +str(loss))

方案2：用tf.Variable承载可训练状态，自动触发重追踪 #

该方案适合需要频繁切换层可训练状态的场景，不需要手动操作图缓存，状态切换逻辑更内聚。
核心逻辑是用非训练的tf.Variable存储层的可训练标志，tf.function会感知该变量的取值变化，在值变更后首次调用时自动重编译静态图，模型改造示例：

class MyModel(keras.Model):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        
        self.square_layer = keras.layers.Dense(2)
        self.cube_layer = keras.layers.Dense(2)
        # 用tf.Variable存储可训练状态，不参与梯度更新
        self.cube_trainable = tf.Variable(False, trainable=False)
        
        self.optimizer = tf.keras.optimizers.Adam()
    
    @tf.function
    def call(self, X):
        # 显式读取控制标志，加入图追踪依赖
        _ = self.cube_trainable.read_value()
        return tf.stack([self.square_layer(X), self.cube_layer(X)], axis=-1)
    
    @tf.function
    def train_step(self, inputs, targets):
        cube_trainable = self.cube_trainable.read_value()
        with tf.GradientTape() as tape:
            predictions = self(inputs)
            loss = tf.reduce_mean(tf.square(predictions - targets))
        # 根据当前状态动态筛选需要更新的权重
        trainable_weights = self.square_layer.trainable_weights
        if cube_trainable:
            trainable_weights += self.cube_layer.trainable_weights
        grads = tape.gradient(loss, trainable_weights)
        self.optimizer.apply_gradients(zip(grads, trainable_weights))
        return loss

    def set_cube_trainable(self, state: bool):
        self.cube_trainable.assign(state)
        self.cube_layer.trainable = state

使用时只需要调用set_cube_trainable切换状态即可，不需要额外处理图缓存。

• 不要直接在tf.function装饰的方法内部修改Python原生布尔类型的trainable值，这类Python侧的副作用不会被静态图正确捕获。
• 单独调用compile()不会清除自定义tf.function方法的缓存，这也是原代码中重新compile后依然不生效的核心原因。

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