PyTorch中如何高效实现逐batch/迭代的训练与验证损失追踪

首先要把两个容易搞混的目标分清楚，不然很容易写出高开销的废代码：

• 记录一次全量验证遍历过程中每个测试batch的损失值：属于零额外计算开销的操作，根本不需要重复跑验证流程。
• 记录训练迭代步对应的验证损失变化：不需要每跑1个训练batch就把整个测试集过一遍，用低开销的抽样验证就能实现目标。

你现在的验证流程本来就是逐batch算损失的，之前只是算完就把单batch的临时值丢了，只留了累计的总和算epoch平均，完全不需要多跑任何前向传播，只要在算完单batch损失的时候直接存到列表里就行：

1. 每次进入验证环节的时候，先建个空列表，专门存这次遍历里所有batch的损失
2. 逐batch算验证损失的时候，原来的累计逻辑一点不动，多一步把当前batch的损失值追加到列表里
3. 这一轮验证跑完，把这个列表里的内容合并到全局的验证损失记录里就行

对应修改代码示例（完全贴合你现有的实现逻辑）：

# 全局初始化的记录列表
train_batch_losses = []  # 你已经实现的训练损失记录
val_batch_losses = []    # 新增的验证逐batch损失记录

# 原有训练循环部分
for epoch in range(2):
    model.train()
    for batch_idx, (train_images, train_labels) in enumerate(trainloader):
        optimizer.zero_grad()
        outputs = model(train_images)
        loss = criterion(outputs, train_labels)
        loss.backward()
        optimizer.step()
        # 你已实现的训练损失记录逻辑
        train_batch_losses.append(loss.item())

    # 原有验证流程，仅需小幅修改
    model.eval()
    epoch_val_total_loss = 0
    epoch_total_samples = 0
    epoch_val_batch_losses = []  # 新增：存放当前epoch验证遍历的逐batch损失
    with torch.no_grad():
        for val_images, val_labels in testloader:
            val_outputs = model(val_images)
            val_loss = criterion(val_outputs, val_labels)
            # 新增：直接存储当前验证batch的损失值，无任何额外计算
            epoch_val_batch_losses.append(val_loss.item())
            # 原有累计逻辑完全保留
            epoch_val_total_loss += val_loss.item() * val_images.size(0)
            epoch_total_samples += val_images.size(0)
    # 原有epoch级平均损失计算完全保留
    epoch_val_loss = epoch_val_total_loss / epoch_total_samples
    # 新增：把当前epoch的逐batch验证损失合并到全局记录
    val_batch_losses.extend(epoch_val_batch_losses)

上述修改没有增加任何前向传播或反向传播的计算量，仅多了常数级别的列表存储操作，完全不会影响训练速度。

如果你需要的是和训练损失一样的追踪效果——也就是横轴是训练了多少个batch、纵轴是这个时间点模型的验证损失，千万不要每跑一个训练batch就全量验证，那样速度掉一半都不止。用固定抽样的方式就能把开销压到最低：

• 提前从测试集里抽固定数量的batch（比如10~20个，占测试集总量的5%以内）缓存到内存里，连后面每次加载数据的开销都能省掉
• 设个固定的触发间隔，比如每训练100个batch，就拿之前缓存的那几个验证batch跑一次前向传播
• 每次跑验证的时候，逐batch记下来损失值，同时记一下当前是第几个训练步，这样后面画图的时候能和训练损失的x轴对齐

实现的时候要避开几个常见的坑：

• 别为了记逐batch的验证损失，每跑1个训练batch就把整个测试集过一遍：这种写法相当于把训练时间直接翻倍，完全没必要
• 算epoch级的平均验证损失的时候，别直接把存的逐batch损失求算术平均：因为最后一个验证batch的样本数可能比普通batch少，要按每个batch的实际样本数对损失做加权计算，结果才和你原来的逻辑一致
• 验证的时候别忘切回model.eval()模式，也别漏了torch.no_grad()上下文，不然误开梯度计算会占额外显存，还会平白增加计算量

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