图像Captioning模型训练10轮后训练与验证损失上升问题

问题背景 #

训练由BERT（文本编码器）和Xception（视觉编码器）并行构成的图像Captioning模型时，前10轮训练表现正常，但第10轮后训练与验证损失急剧上升，具体代码、参数及训练日志如下：

模型代码 #

num_epochs = 20  # 实际训练建议至少30轮
batch_size = 1

vision_encoder = create_vision_encoder(num_projection_layers=1, projection_dims=256, dropout_rate=0.1)
text_encoder = create_text_encoder(num_projection_layers=1, projection_dims=256, dropout_rate=0.1)
dual_encoder = DualEncoder(text_encoder, vision_encoder, temperature=0.05)
dual_encoder.compile(optimizer=tfa.optimizers.AdamW(learning_rate=0.001, weight_decay=0.001)) # run_eagerly=True)

from tensorflow.keras.callbacks import LearningRateScheduler
import math
def step_decay(epoch):
   initial_lrate = 0.001
   drop = 0.005
   epochs_drop = 10.0
   lrate = initial_lrate * math.pow(drop, math.floor((1+epoch)/epochs_drop))
   return lrate
lrate = LearningRateScheduler(step_decay)
callbacks_list = [lrate]

print(f"GPU数量: {len(tf.config.list_physical_devices('GPU'))}")
print(f"样本数量（图像-标题对）: {train_example_count}")
print(f"Batch Size: {batch_size}")
print(f"每轮步数: {int(np.ceil(train_example_count / batch_size))}")
train_dataset = get_dataset(os.path.join(tfrecords_dir, "train-*.tfrecord"), batch_size)
valid_dataset = get_dataset(os.path.join(tfrecords_dir, "valid-*.tfrecord"), batch_size)
# 创建学习率自适应衰减回调
reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.2, patience=3)
# 创建早停回调
early_stopping = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=5, restore_best_weights=True)
history = dual_encoder.fit(
  train_dataset,
  epochs=num_epochs,
  validation_data=valid_dataset,
  callbacks=[reduce_lr, early_stopping, callbacks_list],
  )
print("训练完成。保存视觉编码器和文本编码器...")
vision_encoder.save("/content/drive/MyDrive/vision_encoder")
text_encoder.save("/content/drive/MyDrive/text_encoder")
print("模型已保存。")

训练日志 #

GPU数量: 1
样本数量（图像-标题对）: 3500
Batch Size: 1
每轮步数: 3500
Epoch 1/20
3500/3500 [==============================] - 217s 62ms/step - loss: 5.1028e-04 - val_loss: 1.9643e-04 - lr: 0.0010
Epoch 2/20
3500/3500 [==============================] - 218s 62ms/step - loss: 8.8274e-05 - val_loss: 3.3228e-05 - lr: 0.0010
Epoch 3/20
3500/3500 [==============================] - 220s 63ms/step - loss: 0.3582 - val_loss: 4.2012e-04 - lr: 0.0010
Epoch 4/20
3500/3500 [==============================] - 216s 62ms/step - loss: 9.6259e-04 - val_loss: 3.7130e-05 - lr: 0.0010
Epoch 5/20
3500/3500 [==============================] - 213s 61ms/step - loss: 1.7488e-05 - val_loss: 6.3365e-06 - lr: 2.0000e-04
Epoch 6/20
3500/3500 [==============================] - 208s 59ms/step - loss: 2.9985e-06 - val_loss: 1.0982e-06 - lr: 0.0010
Epoch 7/20
3500/3500 [==============================] - 207s 59ms/step - loss: 1.0761 - val_loss: 0.0212 - lr: 0.0010
Epoch 8/20
3500/3500 [==============================] - 211s 60ms/step - loss: 0.0062 - val_loss: 4.6654e-05 - lr: 2.0000e-04
Epoch 9/20
3499/3500 [============================>.] - ETA: 0s - loss: 2.2375e-05Epoch 10/20
3500/3500 [==============================] - 210s 60ms/step - loss: 234.2512 - val_loss: 309.9704 - lr: 5.0000e-06
Epoch 11/20
3500/3500 [==============================] - 211s 60ms/step - loss: 310.0370 - val_loss: 309.7400 - lr: 1.0000e-06
训练完成。保存视觉编码器和文本编码器...
WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 124). These functions will not be directly callable after loading.
模型已保存。

问题根因分析 #

1. 学习率调度冲突：同时启用了自定义LearningRateScheduler和ReduceLROnPlateau两种学习率调整策略，两者逻辑互相干扰。从日志可见，第10轮学习率被骤降至5e-06，后续进一步降到1e-06，过低的学习率导致模型参数无法有效更新，甚至出现梯度消失或权重震荡，最终引发损失暴涨。
2. 自定义学习率衰减过于极端：step_decay函数中drop=0.005意味着每10轮学习率变为原来的0.005倍，衰减幅度远超合理范围，直接让学习率趋近于0，模型彻底失去优化能力。
3. Batch Size过小：batch_size=1会导致梯度更新噪声极大，训练稳定性极差，后期极易出现损失波动甚至失控。
4. 早停机制未及时触发：EarlyStopping的patience=5设置无法在损失突然飙升时及时终止训练，导致模型继续在异常状态下迭代。

修复方案 #

1. 移除冲突的学习率调度器：保留ReduceLROnPlateau即可，它能根据验证损失自适应调整学习率，避免手动设置的极端衰减。若坚持使用自定义调度，需修改为更温和的衰减逻辑，例如将drop设为0.5，每10轮学习率减半：

   def step_decay(epoch):
      initial_lrate = 0.001
      drop = 0.5  # 更温和的衰减系数
      epochs_drop = 10.0
      lrate = initial_lrate * math.pow(drop, math.floor((1+epoch)/epochs_drop))
      return lrate
   

2. 增大Batch Size：根据GPU显存情况，将batch_size调整为8、16或32，提升梯度更新的稳定性，降低训练噪声。
3. 优化早停配置：将EarlyStopping的patience值降至3，或增加min_delta参数监控损失变化幅度，更早终止异常训练。
4. 检查损失函数与温度系数：确认DualEncoder的损失函数适配图像-文本匹配任务的对比学习逻辑，当前temperature=0.05过小，可尝试调整至0.1~0.5区间，提升损失计算的数值稳定性。

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