PyTorch训练工程化：DataLoader、优化器与Checkpoint复现实践

这一节的内容，前后拆分着看了好几遍才真正吃透。前几次笔记把张量、自动求导、各类网络层过完以后，很容易产生一种“马上就能独立跑训练”的错觉。可真到了自己动手把完整流程搭出来的时候，反馈往往不是“代码写不出来”，而是“代码能跑，但结果不对”；或者“昨天复现出的曲线，今天完全复现不回来”。PyTorch 深度学习（GPT 重译）（五）这一篇，恰好就落在这个断层上：从“能把模型搭出来”到“能把一次实验跑得又快、又稳、又可复现”。如果你正卡在这个过渡阶段，这篇笔记会帮你把数据供给、训练循环、优化器配置、检查点管理、复现性这些细节点挨个串起来，顺带分享几个我只有在真正踩坑之后才会理解的设置。

1. 这一节的位置：训练流程比网络结构更值得抠

1.1 前四篇的底子和第五篇的缺口

前面几篇的内容相对集中在“表达能力”上：张量怎么创建、自动求导怎么运作、卷积和循环结构大概长什么样。这些东西学起来很顺畅，因为每个概念都有非常标准的例题，照着跑就能看到输出。

但到了真正要训练一个像样的模型时，问题就变了。模型结构只是静态的骨架，训练才是让骨架“活”过来的过程。训练涉及的东西全部串在一起：数据怎么喂、梯度怎么算、每步更新多大、中间崩溃了怎么恢复、实验结果怎么保证可复现。这些内容通常散落在各种文档的角落里，很少有一本书能集中在同一章讲完。

第五篇的价值就在这里。它没有继续堆新模型，而是把注意力放到了训练本身的工程细节上。我自己读完最大的感受是：网络结构决定模型的表达能力上限，但训练流程决定最终结果能离这个上限多近。很多人模型结构选得没问题，结果不理想，问题几乎都出在数据管道、学习率策略、状态保存这些“不性感”的环节上。

1.2 实战中真正消耗时间的不是模型，而是数据与训练环节

我在本地复现某个图像分类项目时，光是调数据加载就花了超过半天。模型本身半小时就写完了，效果一直上不去的核心原因是：训练数据读取太慢，GPU 长期处于饥饿状态；验证集没有固定随机种子，导致每个 epoch 评测的样本顺序都在变化；还有一个更隐蔽的问题——训练中断之后，我直接从保存的权重恢复了，却忘了恢复优化器状态，结果训练曲线直接断崖式下跌。

这些问题没有一个是“模型结构”引起的，但任何一个都能让项目进度停滞。所以如果你正自己搭训练流程，建议把注意力从“再换一个更花哨的架构”转移到“现有这套训练流程是不是真的够工程化”。第五篇讲的正是这个层面。

2. DataLoader 与数据管道：训练还没开始就已经决定上限

2.1 Dataset 很简单，DataLoader 配置才是坑

绝大多数教程里，Dataset 的写法都很简单：继承torch.utils.data.Dataset，实现__len__和__getitem__，任务就算完成了。但这只是第一步。真正影响训练效率的是 DataLoader 的参数，而不是 Dataset 本身。

我见过最多的配置如下：

train_loader = DataLoader( dataset=train_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True, prefetch_factor=4, )

num_workers和pin_memory是最容易被随手抄错的地方。num_workers并不是越大越好，它取决于 CPU 核心数以及__getitem__的耗时。如果每个样本只需要做一次简单的张量变换，开 16 个 worker 反而会花费大量时间在进程通信上，CPU 占用率飙升，但训练速度并没有提升。

pin_memory的作用是把 CPU 侧的数据固定在页锁定内存里，这样拷贝到 GPU 时可以走更快的内存传输路径。它适合 GPU 训练场景，但如果你是在 CPU 上跑实验，开pin_memory=True基本没有意义，反而会增加内存成本。

2.2 数据加载参数的取舍与实测经验

下面是我在不同机器上实测后比较稳定的配置：

persistent_workers=True在 PyTorch 1.9 之后的版本里很值得开。默认情况下，每个 epoch 结束之后 worker 进程会被销毁，下个 epoch 重新创建。开这个参数以后，worker 进程会保留下来，省掉频繁创建销毁进程的开销。小数据集上感知不明显，但图像、视频这类大样本集上，效果立竿见影。

prefetch_factor决定每个 worker 提前预取多少个 batch。调太高会让内存暴涨，调太低则可能在每个 step 之间留出等待时间。我的做法是：先从默认值起步，然后用nvidia-smi观察 GPU 利用率。如果 GPU 利用率一直低于 90%，并且 CPU 并没有满载，那大概率是预取不够，优先调prefetch_factor，再考虑num_workers。

提示：shuffle=True加num_workers>0的组合下，每个 worker 持有不同的随机数流。如果只给 Dataset 设置了一次全局种子，多个 worker 之间可能生成完全相同的增强序列，导致同一个 batch 里出现大量重复的增强样本。这个问题非常隐蔽，比较稳妥的做法是设置torch.utils.data.default_collate之前手动给每个 worker 单独设置随机源，或者在__getitem__内部使用实例级np.random.RandomState。

2.3 验证集 / 训练集的随机性边界

验证集的 DataLoader 和训练集最大的区别是：验证集不要shuffle=True。很多人想当然地认为“每次验证时打乱一下也没关系”，但它在两种情况下会出问题：一是评估指标需要和上一次对比，比如 A/B 实验中两个版本模型的准确率差异到底是因为模型变好了，还是因为这次抽到的样本碰巧简单；二是做混淆矩阵或做错误样本分析时，样本顺序一变，排查问题的思路全乱。

如果你需要固定打乱方式，有一个更干净的做法：显式传入一个生成器。

val_loader = DataLoader( val_set, batch_size=64, shuffle=False, )

训练集上随机种子管理也更推荐单独生成一个torch.Generator实例，而不是直接设置全局种子。这样可以保证“跑数据”和“初始化模型”两件事的随机源互不影响，后面想单独锁定数据打乱方式时也更容易操作。

3. 手写训练循环的五个关键动作

3.1 为什么不只用高阶封装

PyTorch 社区里既提供Trainer风格的高阶封装，也鼓励用户手写训练循环。第五篇给的建议非常务实：刚开始实验阶段，手写训练循环更值得。原因是手写循环能看到每一步的真实状态，比如 loss 变化、梯度范数、学习率当前值，而不是只能盯着封装好的进度条。

高阶封装适合已经确定训练流程不变、只想快速出结果的场景。而前中期实验阶段通常需要频繁调整策略，手写 30 行训练循环比学会去配置一个庞大的 Trainer 体系要快得多，也更不容易踩“封装太深导致问题无法定位”的坑。

3.2 逐行拆解训练循环

一个基础但完整的训练循环大概长这样：

for epoch in range(epochs): model.train() for step, (x, y) in enumerate(train_loader): x = x.to(device) y = y.to(device) pred = model(x) loss = criterion(pred, y) optimizer.zero_grad(set_to_none=True) loss.backward() clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step()

注意optimizer.zero_grad(set_to_none=True)这个细节。默认的zero_grad()会把梯度张量置零，set_to_none=True则是直接把梯度设为None。前者在每次反向传播时都要执行一次内存写操作，后者省掉这次写，性能上会有小幅提升，尤其在大模型场景下更明显。

loss.backward()之前不推荐随意加torch.no_grad()，因为反向传播本身就需要计算图。真正需要no_grad()的是评估阶段：

model.eval() with torch.no_grad(): for x, y in val_loader: pred = model(x) val_loss += criterion(pred, y).item()

model.eval()会把 Dropout 和 BatchNorm 切换成推理行为，其意义不是单向的“关掉正则”，而是让模型在验证时表现得稳定。如果不切换，Dropout 在验证时依旧随机丢弃神经元，评估结果会有明显波动。

3.3 梯度累积：显存不够时的干净解决方案

碰到单卡显存不足、又不想把 batch size 调小太多时，梯度累积是常用的办法。核心思想是：用几个小批次累积梯度，攒够一个“等效大 batch”后再更新一次参数。

accum_steps = 4 for step, (x, y) in enumerate(train_loader): pred = model(x) loss = criterion(pred, y) / accum_steps loss.backward() if (step + 1) % accum_steps == 0: clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() optimizer.zero_grad(set_to_none=True)

这里有个很容易忽略的点：loss 需要除以accum_steps。因为每个小批次都会触发反向传播，梯度会累加，除完以后再累积，等效 batch 的梯度幅度才和原来一致。

梯度累积也不是银弹。它不能让每个 step 的等待时间消失，只是把更新频率降下来。如果显存小得连一个小 batch 都塞不进，那果断换更小的输入尺寸或更轻量的骨干网络更实际。

4. 优化器与学习率调度的配合套路

4.1 从 SGD 到 AdamW 的适用场景

选优化器不是选“最新的那个”，而是选“符合当前任务收敛特性的那个”。理解吃透了才能调好，单单照着别人的代码抄，很难应对变化。

很多人只记得 Adam 是默认选项，却忽略了它内部维护了一阶矩和二阶矩两个状态。这两个状态是从训练开始累积的，如果你在中途切换优化器或者改了学习率，而矩估计还在沿用旧的分布，收敛行为会变得很奇怪。这也是为什么断点续训时优化器状态必须同步保存。

4.2 Warmup + 余弦退火的组合实践

训练初期用一个很小的学习率，然后逐步上升到峰值，这个策略叫 warmup。它存在的目的是避免模型在参数还很随机时被过大的梯度带偏。对 Transformer 这类对学习率敏感的模型，warmup 几乎是标配。对小规模 CNN，不一定需要，但也不会有害。

warmup + 余弦退火是效果比较稳的组合之一。代码层面的实现方式有很多种，比较直观的写法是继承torch.optim.lr_scheduler.LRScheduler：

class WarmupCosineScheduler(LRScheduler): def __init__(self, optimizer, warmup_steps, total_steps, min_lr=1e-6): self.warmup_steps = warmup_steps self.total_steps = total_steps self.min_lr = min_lr super().__init__(optimizer) def get_lr(self): step = self.last_epoch if step < self.warmup_steps: return [base_lr * step / max(1, self.warmup_steps) for base_lr in self.base_lrs] progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps) return [self.min_lr + 0.5 * (base_lr - self.min_lr) * (1 + math.cos(math.pi * progress)) for base_lr in self.base_lrs]

直接用LambdaLR也能实现，但可读性稍差。自定义调度器看起来代码多一点，后续加 warm restart 或者分段策略都方便。

4.3 用一个小实验验证调度器的作用

我自己试过在同一个数据集上对比三种策略：固定学习率 3e-4、指数衰减、warmup + 余弦退火。前 50 个 epoch 里，固定学习率在初期 loss 下降很快，但到后期反复震荡；指数衰减后期下降太平缓；warmup + 余弦退火虽然前几个 epoch 看起来落后一点，到中后期收敛更顺滑，最终指标通常更高。

这里想强调的倒不是“某个调度器一定最好”，而是“学习率调度需要和训练步数绑定”。如果你的训练总步数只有几千步，warmup 阶段却占了 20%，那就是浪费时间；如果总步数有几万步，warmup 阶段却只有几个 step，起不到保护作用。调度器参数不是从某个模板里复制出来的，而是根据你的数据量、batch size、batch 总数计算出来的。

5. 模型 Checkpoint 的完整生命周期

5.1 Checkpoint 必须保存的状态不止权重

训练到一半机器重启、云端会话断开、Colab 超时，这些都是真实会发生的事。恢复训练时如果只加载模型权重，虽然能接着跑，但优化器的动量、二阶矩、当前学习率位置、调度器步数全部对不上，后续训练曲线大概率会异常。

更完整的 checkpoint 应该包含这些内容：

torch.save({ 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'epoch': epoch, 'best_acc': best_acc, 'rng_state': torch.get_rng_state(), 'np_rng_state': np.random.get_state(), }, 'checkpoint.pth')

保存optimizer.state_dict不是为了“把一个文件拷走”，而是为了恢复 Adam 的矩估计。跳过这一步，训练刚恢复时容易出现短暂的 loss 突跳。epoch和best_acc属于元数据，记录当前进度，方便断点续训时判断继续跑还是停。

5.2 加载时设备和训练模式的边界

加载模型时最常见的错误之一，是忽略了map_location。保存方是 GPU 训练，恢复方如果是 CPU 环境，直接torch.load会把张量加载到 GPU 上，然后抛出一堆设备不匹配的报错。稳妥的做法是无脑写map_location='cpu'，加载后再手动移动到目标设备：

ckpt = torch.load('checkpoint.pth', map_location='cpu') model.load_state_dict(ckpt['model_state_dict']) model = model.to(device) optimizer.load_state_dict(ckpt['optimizer_state_dict'])

还要注意加载后需要手动调用model.train()或model.eval()。有些模型类会在构造时默认处于训练模式，但如果你是从某个推理脚本里加载权重来继续训练，必须明确切换回来，否则 Dropout 和 BatchNorm 的行为会一直停留在推理状态。

5.3 断点续训最容易忽略的 RNG 状态

如果你对实验复现的要求很高，光保存权重和优化器还不够。数据加载器里每个 epoch 的数据顺序、数据增强的随机变换、模型初始化时的随机噪声，这些都受随机数生成器影响。断点续训时不保存随机数状态，实际训练数据分布会和你中断前的假设不一致。

保存 RNG 状态的目的是：恢复出来，继续下一轮训练时，相当于完全没有中断过。这对实验的公平性很重要。我自己做对比实验时，如果不想让随机因素干扰结果，就靠这个办法把数据路径固定住。需要注意的是torch.get_rng_state()只能拿到当前设备（CPU）的状态，如果用了 CUDA 张量，还需要保存torch.cuda.get_rng_state_all()。分布式训练里，每个进程的 RNG 状态都不同，需要分别保存。

6. 复现性与训练异常诊断

6.1 设置种子真的能复现吗

很多人以为设置了torch.manual_seed(42)就能让实验完全复现，实际跑起来，同一段代码在不同机器、不同 PyTorch 版本、不同 GPU 型号上仍会有细微差异。原因在于某些算子在不同实现里对浮点运算顺序做了优化，比如求和顺序不同就会导致结果不一致。torch.use_deterministic_algorithms(True)可以强制使用确定性的算子实现，但代价是部分算子会变慢，甚至有些算子不支持。

更实用的复现策略是：固定能控制的部分，接受无法控制的部分。能控制的部分包括数据顺序、模型初始化、优化器初始化、dropout 掩码。无法控制的部分包括某些 CUDA 算子内部的并行归约顺序。只要你的指标差异在可接受范围以内，不必为了“完全一致”而牺牲大量性能。如果确实需要严格一致，至少固定以下这些：

random.seed(42) np.random.seed(42) torch.manual_seed(42) torch.cuda.manual_seed_all(42) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False

cudnn.benchmark默认是True，它的作用是在启动时自动搜索对一个输入尺寸最有利的卷积算法。这个搜索过程本身是确定的，但在不同硬件上可能选出不同算法，跨设备复现时会带来偏差。为了严格复现，通常建议关掉它。

6.2 训练发散的快速定位链路

训练中 loss 突然变成NaN或者梯度数值爆炸，这是最令人抓狂的事情之一。一个比较好使的排查链路是这样的。

第一步，开异常检测。在初始化阶段加上一行torch.autograd.set_detect_anomaly(True)，它会在反向传播过程中检查是否有算子产生 NaN 或者无穷大，并报出具体位置。代价是训练变慢，但排错阶段值得。

第二步，看梯度范数。在每个 step 之后打印clip_grad_norm_前后的梯度范数，对比一下是稳定收敛还是在某个层暴增。如果梯度范数从早期就开始膨胀，优先检查学习率设定和是否缺少归一化层。

第三步，检查数据本身。数据里有 NaN 样本、标签越界、或者归一化参数设置错误，都会导致训练数值异常。很多项目训练前没有对数据做严格清洗，这个坑比我预期中要常见得多。

if not torch.isfinite(loss): print('step', step, 'loss is not finite') raise SystemExit('early stop for debugging')

早期在循环里插入这样一段检查，能快速定位“到底是从哪个 step 开始坏的”，比事后翻一堆日志更管用。

6.3 训练速度与数据管道瓶颈的观察

训练慢，不一定是模型慢。一个简单的方法可以判断瓶颈在哪边：把数据加载和模型计算拆开测各自的耗时。

# 只跑数据，不跑模型 start = time.time() for x, y in train_loader: pass data_time = time.time() - start # 只跑模型，不读数据（用固定随机输入） dummy_x = torch.randn(64, 3, 224, 224, device='cuda') start = time.time() for _ in range(100): model(dummy_x) model_time = time.time() - start

如果data_time占整体时间的比例很高，问题基本在数据管道；如果model_time很高，问题在模型本身。针对数据管道常见的解法是调num_workers、prefetch_factor、pin_memory，再进一步可以做缓存、TFRecord 类似的高吞吐格式，或者复用经过预处理的样本。模型慢则优先考虑用混合精度训练，然后再考虑是否真的需要那么大的输入分辨率。

关于混合精度，PyTorch 的torch.autocast用法不算复杂，只要在 forward 前包住即可。对大规模训练，既提升速度又减少显存，而且不会对最终精度造成明显损失。唯一要注意的是某些自定义层的数值稳定性，个别算子需要在float32下计算，不能在 autocast 范围内无脑全包。

我自己使用这套排查链路多年，最大的体会是：训练问题很难靠看代码一次性找出来，需要靠“分段计时 + 状态打印 + 确定性固定”三件套逐步缩小包围圈。模型结构错了会立刻报错，但数据管道和训练流程出了问题往往只是“结果差一点”，两者调试节奏完全不同。这篇根据 PyTorch 深度学习（GPT 重译）第五部分的实践内容总结，算是给同样在搭训练流程的人一份可以直接抄作业的自查清单。最后再分享一个小习惯：训练脚本里我一定会在第一个 epoch 后输出一次完整的状态快照，包括 loss、梯度范数、当前学习率、batch 耗时。这些数据平时不起眼，等真的需要定位问题时，它们就是最可靠的参照系。