使用CustomCallback实现TensorFlow训练中动态冻结层遇错求助

问题根源 #

你遇到的TypeError: 'NoneType' object is not callable错误，核心原因有两个：

1. 回调类直接引用全局的cnn模型和optimizer实例，训练过程中模型状态变更会导致引用失效；
2. 复用已训练过的优化器实例重新编译模型，其内部状态（如学习率更新、动量值）与当前训练流程冲突，引发调用异常。

另外，不重新编译模型时层无法冻结的原因是：Keras仅在编译阶段根据trainable属性构建计算图和梯度更新逻辑，运行时修改属性不会生效，必须重新编译。

修复方案 #

1. 在回调类中使用self.model访问当前训练的模型实例，避免全局变量依赖；
2. 每次冻结层后，创建新的优化器实例（或重置优化器状态）再编译，避免旧状态干扰；
3. 优化冻结逻辑，仅在需要修改层可训练性时执行编译操作。

修改后的完整代码 #

加载库

import tensorflow as tf
from tensorflow.keras import datasets, layers, models
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from keras.callbacks import ModelCheckpoint, EarlyStopping
from tensorflow.keras.utils import Sequence
from keras.models import load_model

加载数据集

# 加载数据集
(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()
# 归一化
X_train, X_test = X_train/255.0, X_test/255.0

构建模型

cnn = models.Sequential([
    layers.Conv2D(filters=32, kernel_size=(1,1), padding="same", activation="relu", input_shape=(32,32,3)),
    layers.Conv2D(filters=64, kernel_size=(3,3), padding="same", activation="relu"),
    layers.MaxPool2D(pool_size=(2,2)),
    
    layers.Conv2D(filters=64, kernel_size=(3,3), padding="same", activation="relu"),
    layers.Conv2D(filters=128, kernel_size=(5,5), padding="same", activation="relu"),
    layers.MaxPool2D(pool_size=(2,2)),
    
    layers.Flatten(),
    layers.Dense(64, activation="relu"),
    layers.Dense(128, activation="relu"),
    layers.Dense(64, activation="relu"),
    layers.Dense(10, activation="softmax")  
])

修复后的CustomCallback类

class CustomCallback(tf.keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs=None):
        # 定义各阶段要冻结的层索引
        freeze_plan = {
            5: [0,1,2],
            10: [3,4,5],
            15: [6,7,8]
        }
        
        if epoch in freeze_plan:
            # 冻结指定层
            for idx in freeze_plan[epoch]:
                self.model.layers[idx].trainable = False
            
            # 创建新的优化器实例（避免旧状态干扰）
            new_optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
            # 重新编译模型
            self.model.compile(
                optimizer=new_optimizer,
                loss="sparse_categorical_crossentropy",
                metrics=["accuracy"]
            )
            print(f"Epoch {epoch}: 已冻结层{freeze_plan[epoch]}并重新编译模型")

定义优化器并编译

# 初始优化器
initial_optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
# 编译初始模型
cnn.compile(
    optimizer=initial_optimizer,
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

训练模型

results = cnn.fit(
    X_train, y_train,
    epochs=20,
    validation_data=(X_test, y_test),
    batch_size=1024,
    callbacks=[CustomCallback()]
)

额外说明 #

• 如果希望保留优化器的状态（比如学习率衰减），可以在回调中重置优化器的状态，而不是创建新实例，例如：

  # 重置优化器状态
  optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
  optimizer.build(self.model.trainable_variables)
  

• 每次冻结层后重新编译，Keras会重新构建梯度计算图，仅对trainable=True的层计算梯度，实现动态冻结的效果。

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