PyTorch图像分类实战：自定义Dataset到多卡训练的完整指南

从第一篇笔记写到现在，PyTorch 基础语法、张量操作、自动求导这些内容基本都过了一遍，但从"能跑通教程代码"到"能独立跑一个自己的训练任务"，中间其实还隔着一段很长的路。这篇笔记 11 我打算换个思路，不再按部就班讲 API，而是把我自己在完整训练一个图像分类模型时踩过的坑、反复调整过的细节一次性梳理出来。内容围绕自定义 Dataset、DataLoader 参数调优、训练循环里的隐性细节、过拟合处理和多卡训练这几个方向展开，偏实战，适合已经掌握 PyTorch 基础、正在尝试跑自己数据的读者参考。

1. 从张量到手写 Dataset：我为什么建议你尽早摆脱内置数据集

1.1 内置数据集很好，但真实项目没有现成的

MNIST、CIFAR-10、ImageNet 这些内置数据集对学习来说非常友好，几行代码就能加载出来，但它们最大的问题在于——太干净了。文件名有序、标签正确、图片尺寸统一、没有损坏文件，所有脏活累活都被框架处理完了。一旦你换了真实场景的数据，比如自己从文件夹读图、处理 CSV 标注文件、过滤坏图，很多人会突然发现自己连数据加载都不会写了。

这不是危言耸听。我见过不少同学在 Kaggle 比赛里直接套用内置数据集的加载方式，然后被各种报错折磨：路径不对、标签错位、图片通道数不一致、显存爆掉还不知道是 DataLoader 一次拿了太多数据。所以这篇笔记第一个想强调的观点就是：尽早学会手写 Dataset，这是从学习走向实战的第一道分水岭。

1.2 Dataset 与 DataLoader 的完整实现：图片分类为例

PyTorch 的 Dataset 类核心只需要实现三个方法：__init__、__len__ 和 __getitem__。看起来简单，但实际写起来有不少细节决定着你后面能否顺利训练。

我以最常见的图像分类场景为例：图片存放在 train/images/ 下，标签存放在 train/label.csv 中，内容大概是 filename,label 两列。一个基础但完整的 Dataset 写法如下：

python复制import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

class ImageClassificationDataset(Dataset):
    def __init__(self, img_dir, label_file, transform=None):
        self.img_dir = img_dir
        self.df = pd.read_csv(label_file)
        # 建立类别名称到索引的映射
        self.class_to_idx = {cls: idx for idx, cls in enumerate(sorted(self.df['label'].unique()))}
        self.idx_to_class = {idx: cls for cls, idx in self.class_to_idx.items()}
        self.transform = transform

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_path = os.path.join(self.img_dir, row['filename'])
        # 用 try-except 捕获坏图，避免训练中断
        try:
            image = Image.open(img_path).convert('RGB')
        except Exception:
            # 坏图就随机拿一张正常的顶替
            return self.__getitem__((idx + 1) % len(self.df))
        label = self.class_to_idx[row['label']]

        if self.transform:
            image = self.transform(image)

        return image, label

这个写法里有几个地方值得说明。第一，convert('RGB') 是必须的，因为真实数据集中经常混入灰度图、RGBA 图，不统一转 RGB 的话，后面 batch 的 shape 对不上，报错特别隐蔽。第二，坏图处理我用了一个简单粗暴的方案——跳过当前图片，取下一张。真实项目中更稳妥的做法是先在预处理阶段把所有坏图剔除掉，但训练过程中偶尔还是会遇到读取超时或文件被占用的情况，加一个 fallback 逻辑能避免整个训练任务直接挂掉。

1.3 一个容易被忽略的坑：返回的样本格式与 collate_fn 的关系

很多人第一次写自定义 Dataset 时，会在 __getitem__ 里返回一个字典：

python复制return {"image": image, "label": label, "filename": row['filename']}

这样做的优点是便于调试和扩展，但它有一个直接影响——DataLoader 默认的 collate_fn 会将字典中的每个键分别堆叠成 batch。如果你的 image 是经过 transform 后的 Tensor，那没问题；但如果你想把原始图片路径也传进 batch，collate_fn 遇到字符串列表会直接报错，因为默认的堆叠逻辑只处理数值型 Tensor。

解决办法有两个：要么避免在样本里包含字符串字段，只在 __getitem__ 返回 Tensor；要么自定义一个 collate_fn。我个人的习惯是训练时只返回 Tensor，调试时单独写一个循环去打印原始信息，不要让调试逻辑污染训练数据流。这样做还有一个好处：DataLoader 的 num_workers 在使用自定义 collate_fn 时，子进程需要序列化这个函数，如果 collate_fn 是定义在某个类内部的局部函数，可能触发 pickle 错误。保持简单，能省掉很多环境相关的诡异问题。

2. DataLoader 参数调优：同样一轮训练，为什么别人比你快 3 倍

2.1 num_workers 不是越大越好：实测对比

DataLoader 的 num_workers 估计是被误解最深的参数。很多人的直觉是：开越多进程，数据加载越快。这个直觉在 4 以内基本成立，但再往上就会遇到收益递减甚至负优化。

我在一台 12 核 CPU、单张 RTX 3090 的机器上做过一次简单对比，数据是 24GB 的图片文件夹，单张图片解码平均耗时约 15ms，模型单步前向加反向约 80ms。把 num_workers 分别设为 0、2、4、8、16，测每 100 步的平均耗时，结果如下：

16 个 worker 反而比 8 个更慢，原因是进程上下文切换和内存带宽争抢消耗了额外的资源。实际项目中我一般建议从 num_workers = 4 起步，然后观察训练时的 GPU 利用率：如果 nvidia-smi 显示 GPU-Util 稳定在 90% 以上，说明数据加载不是瓶颈；如果 GPU 使用率经常掉到 50% 以下，再尝试逐级调大 num_workers。

2.2 pin_memory 的价值与适用场景

pin_memory=True 表示把数据加载到锁页内存（pinned memory）中。CPU 默认使用的分页内存可能会被操作系统换到磁盘，而锁页内存始终留在物理内存里，GPU 从锁页内存拷贝数据走的 DMA（直接内存访问）速度比从可分页内存拷贝快很多。

但这里有个容易忽略的前提：pin_memory 只有在数据需要从 CPU 拷贝到 GPU 时才起作用。如果你用的是 CPU 训练，或者数据本来就存在 GPU 显存里，那 pin_memory=True 不会带来任何提升，反而可能增加主存占用。所以这句配置建议是：只要你有 GPU，训练时 num_workers > 0 且 pin_memory=True 的组合几乎是无脑收益的默认选择；但别把它当成万能药，CPU 训练场景下它毫无意义。

另外提醒一句，pin_memory 会显著增加显存之外的主存占用。一张 1080Ti 上跑 16GB 数据集的场景里，pin_memory=True 时主存占用大约多了 2~3GB 的固定开销。如果你的服务器主存本身就紧张，需要权衡。

2.3 prefetch_factor 与数据流水线瓶颈判断

PyTorch 1.5 之后新增了 prefetch_factor 参数，默认是 2，意思是每个 worker 预取 2 个 batch 的数据。调大这个值可以让数据流水线更深，掩盖数据加载过程中的毛刺，但也意味着需要更多内存来缓存预取的数据。

判断数据流水线是否卡顿有一个简单办法：在训练循环开头记录一下每个 batch 之间的时间间隔，如果出现明显的间隔忽大忽小，大概率是某次数据读取突然变慢。这时候可以考虑：

• 将 prefetch_factor 从 2 提到 4，但观察主存是否吃紧
• 把图片解码从 PIL.Image.open() 换为支持硬件加速的解码库，比如 opencv 或 tensorflow 的解码逻辑（虽然不推荐混用，但提速明显）
• 先做一次全量预处理器，把图片统一 resize 成较小尺寸再存成 LMDB 或 tar 归档文件，减少单次 IO 时间

我自己在项目里最喜欢用的是最后一种思路：离线把原图缩放到 512 或训练所需尺寸的整数倍，然后用 torchdata 或手写一个小脚本打包为 tar 格式。这样训练时从连续的大文件里读数据，IO 性能比读几千个小文件好非常多，num_workers 的需求也会随之降低。

3. 训练循环里的隐形瓶颈：loss.backward() 之外的细节

3.1 optimizer.zero_grad() 放在哪里最安全

关于 zero_grad() 的位置，社区里一直有两派：一派习惯在 loss.backward() 之后调用，一派习惯在 optimizer.step() 之前调用。这两种写法其实殊途同归，都能保证梯度正确累积，但有一个细节值得注意：如果你在每次迭代开始时先调用 zero_grad()，可以避免上一次迭代的梯度残留在某些未被覆盖的参数上。

我在实践中更推荐把 optimizer.zero_grad() 放在计算 loss 之前，也就是每个 batch 开始的第一行。因为如果训练循环里存在 forward 提前 continue 的分支，写在循环头部的 zero_grad() 可以确保即使中途跳过这次迭代，梯度也不会带进下一个 batch。

真正需要注意的是：loss.backward() 默认会累加梯度，而不是覆盖梯度。这意味着如果你连续多次调用 backward() 而没有清零，梯度会被反复累加，最终导致优化方向振荡甚至梯度爆炸。所以无论你把 zero_grad() 放在哪里，一定要保证它对每个需要更新的参数都执行到了。

3.2 梯度累积：小 batch 训练大模型的实用技巧

梯度累积的原理不复杂：每 accumulation_steps 个 batch 才调用一次 optimizer.step()，相当于把多个小 batch 的梯度累加后模拟一个大 batch 的效果。这个技巧在显存不足时非常实用，比如你想用 64 的 batch size 但显存只够跑 16，那就设置 accumulation_steps = 4。

python复制accumulation_steps = 4
scaler = torch.cuda.amp.GradScaler()

for batch_idx, (inputs, labels) in enumerate(train_loader):
    inputs, labels = inputs.cuda(), labels.cuda()

    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss_value = loss_fn(outputs, labels)

    # 注意这里要除以 accumulation_steps，把梯度归一化到等价于大 batch 的尺度
    scaler.scale(loss_value / accumulation_steps).backward()

    if (batch_idx + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

这里最容易犯的错就是忘记除以 accumulation_steps。不除的话，等价于用比预期大 accumulation_steps 倍的学习率在训练，收敛曲线会明显不稳。另外还有一个隐含辩证：梯度累积只是"等效"大 batch，不是完全等价，因为 BatchNorm 层的统计量是按单个小 batch 计算的，累积梯度并不会累积 BatchNorm 的均值方差估计，所以某些任务效果可能略有差异。

3.3 学习率调度的踩坑记录：scheduler 与优化器状态

PyTorch 中学习率调度器最常见的用法是 scheduler.step() 每个 epoch 调用一次，但我见过不少初学者在训练循环里每步都调用 scheduler.step()，然后发现损失曲线疯狂振荡，心态直接崩掉。

ReduceLROnPlateau 和 StepLR 的调用时机是不一样的。StepLR 每调用一次就按固定步长衰减一次，所以如果你在每个 step 调用，学习率会衰减得非常快；ReduceLROnPlateau 则根据验证集指标动态调整，通常放在每个 epoch 结束、拿到验证集 loss 或者准确率之后调用。

关于恢复训练还有个经常踩的坑：使用 torch.save(checkpoint) 保存断点时，很多人只保存了 model.state_dict()，忘了保存 optimizer.state_dict() 和 scheduler.state_dict()。如果训练中断后想恢复，优化器的动量状态和 Adam 的一阶二阶矩估计全部丢失，相当于学习率调度器也从头开始，实际训练效果会与预期偏差很大。正确的保存方式应该包含三件套：

python复制checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_acc': best_acc,
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')

4. 过拟合与模型稳定：从验证集设计到模型保存的最佳实践

4.1 验证集划分的陷阱：乱序、重复与数据泄露

很多人拿到数据之后直接 train_test_split 或者手动随机抽样，但这些做法在时间序列、医疗图像、以人为单位收集的多张图片数据中很容易引入数据泄露。比如同一个人的多张照片被同时分到训练集和验证集，模型其实是在"认人"而不是"认病"或"认物体"，验证集指标虚高，部署时立刻打回原形。

我踩过最惨的一次：某目标检测项目里，同一辆车出现在多个视频帧中，随机划分导致训练集和验证集中大量出现同一辆车的不同角度图片，验证集 mAP 高达 0.85，但真正部署到停车场识别新场景时 mAP 直接掉到 0.6 以下。

正确的做法是按最高层级的实体分组划分：如果是人脸识别，按人的 ID 划分；如果是车辆检测，按车牌或车辆 ID 划分；如果是时间序列，按时间段切分，训练集取前 80% 时间，验证集取后 20%。PyTorch 中实现分组划分非常简单，用 torch.utils.data.Subset 配合预先计算好的索引数组即可：

python复制from torch.utils.data import Subset

train_idx = []  # 由分组逻辑计算得到
val_idx = []    # 由分组逻辑计算得到

train_dataset = Subset(full_dataset, train_idx)
val_dataset = Subset(full_dataset, val_idx)

比起直接在 Dataset 里加一个 train=True 参数再去切分，Subset 的写法更灵活，也便于追踪每个样本的真实数据来源。

4.2 早停与最佳模型的保存逻辑

训练神经网络有个心照不宣的事实：验证集 loss 降到最低点之后，继续训练虽然会降低训练集 loss，但验证集效果只会越来越差，这就是过拟合。早停（Early Stopping）是应对这个问题最简单、最有效的手段之一。

实现早停的时候，我建议不仅记录验证集指标的历史最优值，还要记录达到该最优值时的模型状态和优化器状态。这样训练中途断电或崩溃时，你可以直接从最好的一版 checkpoint 继续，而不是从最后一个 epoch 恢复。

一个比较完整的早停 + 最优模型保存逻辑如下：

python复制best_val_loss = float('inf')
patience = 10
bad_epochs = 0

for epoch in range(num_epochs):
    train_loss = train_one_epoch(model, train_loader, optimizer, criterion)
    val_loss = evaluate(model, val_loader, criterion)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        bad_epochs = 0
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'val_loss': val_loss,
        }, 'best_model.pth')
    else:
        bad_epochs += 1
        if bad_epochs >= patience:
            print(f'Early stopping at epoch {epoch}')
            break

patience 设置多少取决于模型规模和数据集大小，我通常设置在 8~15 之间。太小的 patience 会在验证集指标还有回升空间时提前终止，太大又容易浪费时间在过拟合阶段。

4.3 复现性：种子设置与 cudnn 配置

深度学习实验的可复现性是一个常被忽略但极其重要的命题。你这次训练得到一个准确率 0.91 的模型，换一台机器或者重启进程之后变成 0.90，这中间除了数据加载顺序，还有一个重要变量是 CuDNN 的基准模式。

PyTorch 在 torch.backends.cudnn.benchmark = True 时会自动选择“最快的卷积算法”，但这个选择过程依赖硬件和算法库版本，可能导致复现性下降。如果你需要一次完全可复现的调试，建议用：

python复制def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

注意 deterministic=True 会牺牲一定的性能，因为部分卷积算子会放弃一些快速的非确定性实现。通常的做法是：调试和对比实验阶段开启 deterministic，保证每次结果一致；正式训练大量数据时选择 benchmark=True，追求吞吐量，不去纠结单次实验的微小随机性。

5. 从单卡到多卡：分布式训练入门的核心认知

5.1 DataParallel 与 DistributedDataParallel 的选择

单卡显存不够时，大家第一反应是试试 DataParallel。它的用法非常直接：model = nn.DataParallel(model)，然后把 batch 扔进去，框架会自动把数据拆分到多张卡上。但它的底层实现是同步的，每张卡都持有完整模型副本，主卡要负责汇总梯度，这导致它的通信开销随着卡数增长非常快。实测下来，4 张卡用 DataParallel 相比单卡可能只有 2.5 倍左右的加速，性价比不高。

DistributedDataParallel（DDP）的底层用的是 NCCL 通信原语，每个进程独立计算梯度后再做梯度同步，通信开销远低于 DataParallel。而且 DDP 支持单机多卡和多机多卡，是 PyTorch 官方推荐的正式方案。

一个最基础的 DDP 初始化流程如下：

python复制import torch.distributed as dist
import torch.multiprocessing as mp

def train_worker(rank, world_size):
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    model = model.cuda(rank)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])

    # 构建 DistributedSampler，确保数据按卡切分
    from torch.utils.data.distributed import DistributedSampler
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, sampler=sampler)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)   # 每个 epoch 打乱顺序，否则每个 epoch 数据划分完全一致
        # ... 正常训练循环

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    mp.spawn(train_worker, args=(world_size,), nprocs=world_size)

5.2 多卡训练里最容易出错的三件事

第一件是 batch_size 的语义变化。单卡时设定 batch_size=64 是每步处理 64 张图；DDP 的 batch_size 是每张卡各自处理的样本数，全局 batch 实际是 64 × 卡数。如果你不调整学习率，全局 batch 变大后收敛行为会明显变化，通常需要相应调整学习率或用 warmup 来稳定。

第二件是验证集评估。DDP 环境下最常见的问题是每张卡都在跑验证集，导致结果重复统计，或者只跑了部分数据。标准做法是把验证集的 DataLoader 也设置为 batch_size=1 或用 DistributedSampler(shuffle=False)，然后在 all_gather 收集所有卡的预测结果后再统一计算指标。很多人偷懒只在 rank 0 上跑验证，这会导致验证集上的 BatchNorm 统计量不准确，尤其是模型用了 BatchNorm 层且验证 batch 较小时，误差会更明显。

第三件是模型保存的组织。DDP 中每张卡上都有一份模型副本，但只有 rank 0 需要保存模型。如果每张卡都在写 checkpoint，会写出多个相同文件浪费磁盘，还可能因为文件锁冲突报错。规范写法是：

python复制if dist.get_rank() == 0:
    torch.save(model.module.state_dict(), 'best_model_ddp.pth')

注意 DDP 包装后的模型要用 model.module.state_dict() 而不是 model.state_dict()，除非你保存时先取了 module 属性。这个小细节不处理好，加载模型时会遇到 key 前缀不匹配的问题。

6. 给下一个阶段：定位模型问题先看曲线还是先看代码

训练完模型发现效果不理想，这是常态。问题在于很多人第一反应是改网络结构、换损失函数、调学习率，一顿操作猛如虎，最后连问题出在哪里都没搞清楚。

我的建议是先用验证集曲线定位问题区间，再回头看代码和数据。如果训练集 loss 一直降不下来，大概率是模型容量不足、学习率过大导致不收敛，或者数据预处理有问题；如果训练集 loss 能降但验证集 loss 不降，说明过拟合，优先考虑数据增强、正则化和早停；如果训练和验证的 loss 曲线震荡剧烈，先检查学习率是不是太大、batch size 是不是太小、数据中是否有异常标签。

还有一些老手才知道的快速诊断手段。比如训练前先跑一个 batch，看看 loss 是否能正常下降；如果 loss 纹丝不动，重点检查标签是否错位、损失函数是否写错、优化器参数是否初始化异常。又比如把模型参数初始化后，直接过一遍训练集最后一个 batch，打印预测分布，如果所有类别概率都集中在某一个值附近，说明初始化可能有问题，换一种初始化方式往往立竿见影。

这些诊断方法没有哪个是 PyTorch 直接告诉你的，都是靠一次次失败实验换来的。这篇笔记写的都是我在实际项目中验证过的做法，希望能帮你少走一些弯路。下一篇笔记我打算重点写一写混合精度训练的细节和常见落地坑，那也是目前从训练到部署绕不开的一环。