跳到主要内容
SUN / 笔记
← 全部文章

PyTorch 训练流程与验证

深度学习1628 字 · 5 分钟
#PyTorch#模型训练

1. 训练流程

flowchart LR
    A[原始数据] --> B[划分 train / val / test]
    B --> C[Dataset 与 transform]
    C --> D[DataLoader]
    D --> E[模型前向]
    E --> F[损失与指标]
    F --> G[反向传播与参数更新]
    G --> H[验证与选模]
    H --> I[独立测试]
    I --> J[保存配置、权重与结果]

三个数据集的职责必须分开：

  • 训练集：用于计算梯度和更新参数。
  • 验证集：用于调参、早停和选择最佳检查点。
  • 测试集：只在方案冻结后做最终估计，不能每轮参与选模。

2. Dataset 与 DataLoader 的边界

Dataset 回答“第 i 个样本是什么”，DataLoader 回答“如何组成批次并送给训练循环”。

python
from torch.utils.data import DataLoader, Dataset
import torch


class ToyDataset(Dataset):
    def __init__(self):
        self.x = torch.arange(100, dtype=torch.float32).unsqueeze(1)
        self.y = (self.x[:, 0] % 2).long()

    def __len__(self):
        return len(self.x)

    def __getitem__(self, index):
        return self.x[index], self.y[index]

常用参数：

参数 作用 实践建议
batch_size 每批样本数 受显存、吞吐和泛化需求共同影响
shuffle 每轮重排索引 训练集通常开启，验证和测试关闭
num_workers 数据加载子进程数 从 0 开始调试，再通过实测提高
pin_memory 把 CPU 张量放到页锁定内存 CUDA 训练时通常值得测试
drop_last 丢弃最后一个不足批次 BatchNorm 或固定形状需求下再开启
collate_fn 自定义组批 可变长序列、图或复杂样本结构
persistent_workers 跨 epoch 保留 worker num_workers > 0 且多轮训练时可测试

Windows 多进程加载

Windows 默认使用 spawn 创建 worker，不是“只能设为 0”。使用正数时应：

python
def main():
    dataset = ToyDataset()
    loader = DataLoader(dataset, batch_size=16, num_workers=2)
    for batch in loader:
        pass


if __name__ == "__main__":
    main()

自定义 Dataset、collate_fn 和 worker_init_fn 应定义在模块顶层，避免使用不能被 pickle 的 lambda。Notebook 或调试阶段先使用 num_workers=0，错误栈更清楚。

3. 数据预处理与数据增强

  • 确定性预处理：解码、缩放、标准化、格式转换。
  • 随机增强：随机裁剪、翻转、颜色扰动等，只应用于训练集。
  • 验证/测试变换：保持确定性，并与模型预训练权重要求的输入范围一致。

迁移学习时不要凭经验手写预处理。TorchVision 的权重对象同时提供匹配的推理变换：

python
from torchvision.models import ResNet18_Weights

weights = ResNet18_Weights.DEFAULT
preprocess = weights.transforms()

4. 一份可复用的分类 epoch 模板

这份实现按样本数累计损失和正确数，避免“先算每批准确率、再对批次取平均”带来的最后小批次偏差。

python
from contextlib import nullcontext
import torch


def run_epoch(model, loader, loss_fn, device, optimizer=None):
    training = optimizer is not None
    model.train(training)

    loss_sum = 0.0
    correct = 0
    sample_count = 0
    grad_context = nullcontext() if training else torch.inference_mode()

    with grad_context:
        for inputs, targets in loader:
            inputs = inputs.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)

            if training:
                optimizer.zero_grad(set_to_none=True)

            logits = model(inputs)
            loss = loss_fn(logits, targets)

            if training:
                loss.backward()
                optimizer.step()

            batch_size = targets.size(0)
            loss_sum += loss.item() * batch_size
            correct += (logits.argmax(dim=1) == targets).sum().item()
            sample_count += batch_size

    if sample_count == 0:
        raise ValueError("loader contains no samples")

    return {
        "loss": loss_sum / sample_count,
        "accuracy": correct / sample_count,
    }

注意：

  • CrossEntropyLoss 接收未归一化的 logits，不要在模型末尾先做 softmax()。
  • model.train(training) 会影响 Dropout 和 BatchNorm，但不负责开关自动微分。
  • 回归、多标签分类或分割任务需要替换指标计算，不能直接复用 argmax()。

5. 训练、验证和选择最佳模型

python
import copy
import torch

best_val_loss = float("inf")
best_state = None

for epoch in range(num_epochs):
    train_metrics = run_epoch(
        model, train_loader, loss_fn, device, optimizer=optimizer
    )
    val_metrics = run_epoch(
        model, val_loader, loss_fn, device
    )

    if val_metrics["loss"] < best_val_loss:
        best_val_loss = val_metrics["loss"]
        best_state = copy.deepcopy(model.state_dict())

    print(
        f"epoch={epoch + 1:03d} "
        f"train_loss={train_metrics['loss']:.4f} "
        f"val_loss={val_metrics['loss']:.4f} "
        f"val_acc={val_metrics['accuracy']:.4f}"
    )

if best_state is None:
    raise RuntimeError("no valid checkpoint was selected")
model.load_state_dict(best_state)
test_metrics = run_epoch(model, test_loader, loss_fn, device)

如果直接保存到磁盘：

python
torch.save(model.state_dict(), "best_model.pth")

state_dict = torch.load(
    "best_model.pth",
    map_location=device,
    weights_only=True,
)
model.load_state_dict(state_dict)
model.eval()

只保存权重时要同时保存足够的模型配置，例如类别数、模型版本、输入变换和标签映射，否则无法可靠重建模型。

完整训练检查点

需要断点续训时，至少保存：

python
checkpoint = {
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "best_val_loss": best_val_loss,
    "config": config,
}
torch.save(checkpoint, "checkpoint.pth")

学习率调度器、混合精度缩放器和随机数状态如参与训练，也应按恢复需求纳入检查点。

6. 学习率调度器的调用位置

  • StepLR 等按 epoch 更新：通常在每个 epoch 训练结束后 scheduler.step()。
  • ReduceLROnPlateau 依赖验证指标：scheduler.step(val_loss)。
  • OneCycle 等按 step 更新：通常每个训练 batch 调用一次。

不要把所有 scheduler 都套用同一种调用方式。

7. 监控什么

最小可用记录：

  • epoch、global step、学习率；
  • 训练和验证 loss；
  • 与任务匹配的指标，如 accuracy、F1、ROC-AUC；
  • 最佳验证 epoch 和对应测试结果；
  • 数据划分、随机种子、模型配置和运行环境。

诊断逻辑：

现象 优先检查
训练 loss 不下降 标签、损失、学习率、梯度是否存在、输入尺度
训练好而验证差 数据泄漏、过拟合、划分偏差、增强不一致
loss 为 NaN/Inf 数据异常、学习率、除零、数值范围、梯度爆炸
GPU 利用率低 数据加载、batch 太小、CPU 预处理、同步点
复现实验失败 数据划分、随机种子、非确定性算子、版本差异

8. 从“跑通”到“验收”

一次运行结束不等于训练完成。交付前至少检查：

  • [ ] 数据集划分无交叉，测试集未参与调参。
  • [ ] 训练、验证、测试的变换符合各自职责。
  • [ ] 指标按样本而不是按 batch 等权累计。
  • [ ] 最佳检查点确实来自验证集。
  • [ ] 重新加载检查点后结果可复现。
  • [ ] 保存了类别映射、输入尺寸、归一化参数和配置。
  • [ ] 对类别不平衡任务补充混淆矩阵、macro-F1 等指标。
  • [ ] 记录失败实验和数据异常，而不是只保留最好数字。

来源映射

相关文章

CNN：LeNet 与 ResNet →PyTorch：张量与自动微分 →