如何修复TensorFlow Keras回调中批量结束时验证集预测权重未更新问题

问题根源 #

在on_batch_end中调用model.predict()时，由于Keras默认使用tf.function进行图模式训练，predict方法的独立前向传播流程可能无法同步获取刚更新的批量权重，导致使用初始权重计算预测结果。

修复步骤 #

• 直接调用模型实例获取预测：替换predict方法为模型的__call__接口，确保使用当前内存中的最新权重。
• 优化指标计算逻辑：使用Keras/TensorFlow内置的损失函数，避免手动张量转换的冗余操作。
• 优化数据传递方式：将验证数据传入回调初始化方法，避免依赖全局变量，提升代码健壮性。

修正后的代码 #

class LossHistory(keras.callbacks.Callback):
    def __init__(self, validation_data):
        super().__init__()
        self.validation_data = validation_data
        self.mse_val_per_epoch = []  # 按epoch存储每个batch的验证MSE

    def on_epoch_begin(self, epoch, logs=None):
        self.mse_val_per_epoch.append([])  # 为当前epoch初始化存储列表

    def on_batch_end(self, batch, logs=None):
        x_val, y_val = self.validation_data
        # 使用当前模型最新权重做预测，training=False切换为预测模式
        ypred = self.model(x_val, training=False)
        # 计算并存储MSE
        mse = tf.reduce_mean(tf.square(y_val - ypred)).numpy()
        self.mse_val_per_epoch[-1].append(mse)

# 初始化回调时传入验证数据
mycallback = LossHistory(validation_data=(validation_data[0], validation_data[1]))
# 启动训练
mymodel.fit(training_data, training_target, callbacks=[mycallback])

关键说明 #

1. self.model(x_val, training=False)：直接调用模型实例会跳过predict方法的额外封装，确保使用当前批量更新后的权重；training=False会让模型进入预测模式，关闭Dropout、BatchNormalization等训练专属行为，和predict的行为保持一致。
2. 存储结构优化：用mse_val_per_epoch按epoch存储每个batch的MSE，方便后续分析不同epoch内的batch级验证指标变化。
3. 性能提示：如果验证数据集较大，每个batch都做全量验证会显著增加训练耗时，建议使用小批量验证样本，或调整验证频率（比如每N个batch验证一次）。

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