PyTorch训练时如何存储所有epoch验证损失以绘制完整学习曲线？

1. 初始化损失存储容器 #

在run_models.py的epoch训练循环开始前，声明空列表用于存储每轮的验证指标，示例代码如下：

# 放在 for epoch in range(n_epochs) 代码行的上方
val_loss_history = []
# 若需要记录验证集MSE、训练损失可新增对应列表
val_mse_history = []
train_loss_history = []

2. 每轮epoch结束后追加数据 #

找到原代码中每轮训练结束后计算test_res的代码块，在指标计算完成后，将当前epoch的验证指标追加到对应存储列表中：

# 原代码已有的test_res计算逻辑
test_res = compute_val_metrics(model, val_dataloader)
# 新增追加逻辑，字段名需和test_res返回的键名保持一致
val_loss_history.append(test_res['loss'])
val_mse_history.append(test_res['mse'])
# 若要存训练损失，在训练阶段计算完train_loss后追加即可

3. 训练全流程结束后调用绘图逻辑 #

把你写的绘图代码放到所有epoch循环结束的位置，传入完整的损失列表绘制曲线，不要把绘图逻辑放在epoch循环内部，避免每次仅绘制单数据点：

import matplotlib.pyplot as plt

# 绘制验证损失曲线
plt.plot(range(1, len(val_loss_history)+1), val_loss_history, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Learning Curve')
plt.legend()
# 可选择保存到本地或者直接展示
plt.savefig('val_learning_curve.png')
plt.show()

注意事项 #

• 禁止将存储容器的初始化代码放在epoch循环内部，否则每轮循环都会重置列表，仅能保留最后一轮的数据
• 若原代码配置了早停逻辑，需要确保存储列表会记录到早停触发时的所有epoch数据，不会被提前截断
• 若需要同时对比训练、验证曲线，同步在训练阶段每轮结束后追加训练损失到对应列表，绘图时新增对应曲线即可

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