PyTorch中CNN训练与验证损失曲线绘制及代码正确性检查

一、代码问题排查与修正 #

你的代码整体训练逻辑框架没问题，但存在两处关键错误，同时有一些优化点：

1. 核心错误修正 #

• 验证损失列表初始化位置错误：你把validationStep_loss = []放在了enumerate(val_dataset)的循环内部，这会导致每个batch都重置列表，最终只保留最后一个batch的损失，计算出的epoch验证损失完全错误。需要将列表初始化移到验证循环外面。
• 验证阶段未禁用梯度计算：在model.eval()模式下，必须配合torch.no_grad()上下文管理器，避免不必要的梯度计算，节省显存和计算资源。

2. 额外优化建议 #

• 避免直接用torch.tensor(data['target'])转换，如果data['target']已经是张量，直接用.float()更安全；同时建议将数据移到和模型相同的设备（CPU/GPU），避免设备不匹配报错。
• (i+1) % 1 == 0的打印条件会每个step都输出日志，建议改成(i+1) % 10 == 0这类间隔，减少冗余输出。

修正后的完整代码：

import torch
import numpy as np

# 自定义RMSE损失函数
def RMSELoss(predicted, target):
    return torch.sqrt(torch.mean((predicted - target)**2))

criterion = RMSELoss
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
epochs = 300

# 假设train_dataset和val_dataset是DataLoader实例
n_total_steps = len(train_dataset)

trainingEpoch_loss = []
validationEpoch_loss = []

# 自动适配CPU/GPU设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

for epoch in range(epochs):
    step_loss = []
    model.train()
    for i, data in enumerate(train_dataset):
        # 数据转float并移到指定设备
        feature = data['data'].float().to(device)
        target = data['target'].float().to(device)
        
        optimizer.zero_grad()
        outputs = model(feature)
        training_loss = criterion(outputs, target)
        training_loss.backward()
        optimizer.step()
        
        step_loss.append(training_loss.item())
        # 每10个step打印一次日志
        if (i+1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {training_loss.item():.4f}')
    # 计算当前epoch的平均训练损失
    trainingEpoch_loss.append(np.mean(step_loss))
 
    model.eval()
    validationStep_loss = []  # 移到验证循环外部
    with torch.no_grad():  # 禁用梯度计算
        for i, data in enumerate(val_dataset):
            feature = data['data'].float().to(device)
            target = data['target'].float().to(device)
            
            outputs = model(feature)
            validation_loss = criterion(outputs, target)
            validationStep_loss.append(validation_loss.item())
    # 计算当前epoch的平均验证损失
    validationEpoch_loss.append(np.mean(validationStep_loss))

二、绘制训练与验证损失曲线 #

我们用Python最常用的matplotlib库来绘制曲线，步骤如下：

1. 导入依赖库 #

import matplotlib.pyplot as plt

2. 绘制曲线代码 #

plt.figure(figsize=(10, 6))
# 绘制训练损失曲线
plt.plot(range(1, epochs+1), trainingEpoch_loss, label='Training Loss', color='blue')
# 绘制验证损失曲线
plt.plot(range(1, epochs+1), validationEpoch_loss, label='Validation Loss', color='red')

# 添加图表标注与样式
plt.title('Training vs Validation Loss (RMSE)')
plt.xlabel('Epochs')
plt.ylabel('RMSE Loss')
plt.legend()
plt.grid(True)

# 显示图表或保存为图片
plt.show()
# 若需要保存图片，取消下方注释
# plt.savefig('loss_curve.png', dpi=300)

曲线解读 #

• 若训练损失和验证损失持续下降并趋于平稳，说明模型训练正常，拟合效果良好。
• 若训练损失持续下降，但验证损失先降后升，说明模型出现过拟合，可考虑添加Dropout、L2正则化，或提前停止训练。
• 若两者均无下降趋势，需检查学习率设置、模型结构合理性或数据质量。

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