设置layer.trainable后TensorFlow model.trainable_variables未更新问题

问题原因 #

当你设置model.trainable = False时，Keras会执行两个关键操作：

1. 递归将所有子层的trainable属性设为False；
2. 将模型自身的_trainable标志设为False，此时模型会直接返回空的trainable_variables集合，忽略子层后续修改的trainable属性。

这就是你单独修改目标层trainable = True后，model.trainable_variables仍为空的核心原因。

两种解决方案 #

方案一：直接操作层的trainable属性（推荐） #

不要修改模型级别的trainable，直接遍历所有层设置为不可训练，再将目标层设为可训练。这种方式模型的_trainable保持True，会自动根据子层状态收集可训练变量：

import tensorflow as tf 
import random

# 加载预训练模型
model = tf.keras.applications.mobilenet_v2.MobileNetV2(
    input_shape=None,
    alpha=1.0,
    weights='imagenet',
    classifier_activation=None)

# 逐个设置所有层为不可训练
for layer in model.layers:
    layer.trainable = False

# 随机选择5个可训练层
selected_layers = [layer.name for layer in model.layers]
while len(selected_layers) > 5:
    rand_index = random.randint(0, len(selected_layers) - 1)
    del selected_layers[rand_index]

# 设置目标层为可训练
for layer in model.layers:
    if layer.name in selected_layers:
        layer.trainable = True

print(model.trainable_variables)  # 正常输出选中层的参数

方案二：重置模型的编译缓存 #

如果你已经设置了model.trainable = False，修改完目标层的trainable后，需要触发模型重新计算可训练变量集合。可以调用model.compile()（即使不需要实际训练，随便指定一个优化器即可），或者使用内部APImodel._reset_compile_cache()：

# 原代码中设置完目标层trainable后添加：
model.compile(optimizer='sgd')  # 仅触发缓存重置，无需实际使用优化器
# 或者使用内部方法：model._reset_compile_cache()

print(model.trainable_variables)  # 此时会输出选中层的参数

验证梯度计算 #

修改后可以正常手动计算梯度：

x = tf.random.normal((1, 224, 224, 3))
y = tf.constant([0])

with tf.GradientTape() as tape:
    pred = model(x)
    loss = tf.keras.losses.sparse_categorical_crossentropy(y, pred, from_logits=True)

dloss_dparams = tape.gradient(loss, model.trainable_variables)
# 检查梯度是否存在（应为全True）
print([g is not None for g in dloss_dparams])

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