从ModelCheckpoint恢复模型后GradientTape无法计算梯度

你遇到的核心问题是：保存并加载完整Keras模型后，Grad-CAM计算中tape.gradient(logits, last_conv_layer_output)返回None，但同会话训练后直接运行或仅加载权重时正常。

原因在于：

• 完整模型保存加载后，模型层与张量的关联是序列化重建的，你手动通过model.layers[-1].kernel和model.layers[-1].bias计算logits的方式，与grad_model输出的last_conv_layer_output之间没有建立有效的梯度追踪关系——这些权重变量是静态加载的，和动态前向传播的张量流脱节。
• 而训练时模型是动态构建的，权重变量与张量的梯度关联未被破坏；仅加载权重时，模型结构是重新定义的，张量流从头构建，梯度追踪自然正常。

方案1：修改Grad-CAM实现，保留模型原生梯度关联 #

不要手动计算logits，而是利用模型层的原生调用逻辑确保梯度追踪有效：

def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
    # 构建梯度模型：输入→[最后卷积层输出, 最后一层的输入]
    grad_model = tf.keras.models.Model(
        model.inputs,
        [model.get_layer(last_conv_layer_name).output, model.layers[-1].input]
    )

    with tf.GradientTape() as tape:
        tape.watch(img_array)
        last_conv_layer_output, last_layer_input = grad_model(img_array)
        # 直接通过模型层的运算逻辑计算logits，确保梯度追踪
        logits = tf.matmul(last_layer_input, model.layers[-1].kernel) + model.layers[-1].bias
        # 显式将运算绑定到当前计算图
        logits = tf.identity(logits)

    grads = tape.gradient(logits, last_conv_layer_output)
    # 后续热力图计算逻辑...

更稳妥的方式是提前拆分模型最后一层的激活与线性运算：

# 修改原模型定义
CNN_model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3)),
    tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation="relu", name='conv1'),
    tf.keras.layers.Conv2D(filters=64, kernel_size=3, activation="relu", name='conv2'),
    tf.keras.layers.MaxPooling2D(pool_size=2, name='maxpool1'),
    tf.keras.layers.Conv2D(filters=128, kernel_size=3, activation="relu", name='conv3'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu', name='dense1'),
    tf.keras.layers.Dense(1, activation=None, name='logits'),  # 单独输出logits（无激活）
    tf.keras.layers.Activation('sigmoid', name='prediction')
], name='CNN_model')

然后在Grad-CAM中直接获取logits：

def make_gradcam_heatmap(img_array, model, last_conv_layer_name):
    grad_model = tf.keras.models.Model(
        model.inputs,
        [model.get_layer(last_conv_layer_name).output, model.get_layer('logits').output]
    )

    with tf.GradientTape() as tape:
        tape.watch(img_array)
        last_conv_layer_output, logits = grad_model(img_array)

    grads = tape.gradient(logits, last_conv_layer_output)
    # 后续处理...

方案2：仅保存/加载模型权重（已验证可行） #

如果不想修改模型结构或Grad-CAM代码，可采用这种方式：

1. 训练时仅保存权重：

callbacks = [
    tf.keras.callbacks.ModelCheckpoint("best_CNN_weights.h5",
                                       save_best_only=True,
                                       save_weights_only=True),
]

2. 加载时先重建模型结构，再加载权重：

# 重新定义与训练时完全一致的模型结构
CNN_model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3)),
    tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation="relu", name='conv1'),
    tf.keras.layers.Conv2D(filters=64, kernel_size=3, activation="relu", name='conv2'),
    tf.keras.layers.MaxPooling2D(pool_size=2, name='maxpool1'),
    tf.keras.layers.Conv2D(filters=128, kernel_size=3, activation="relu", name='conv3'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu', name='dense1'),
    tf.keras.layers.Dense(1, activation='sigmoid', name='prediction')
], name='CNN_model')

# 加载预训练权重
CNN_model.load_weights("best_CNN_weights.h5")

这种方式下模型的张量流是重新构建的，手动计算logits的梯度追踪会正常工作。

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