这一节的内容前后拆分着看了好几遍才真正吃透。前几次笔记把张量、自动求导、各类网络层过完以后很容易产生一种“马上就能独立跑训练”的错觉。可真到了自己动手把完整流程搭出来的时候反馈往往不是“代码写不出来”而是“代码能跑但结果不对”或者“昨天复现出的曲线今天完全复现不回来”。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( datasettrain_set, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue, persistent_workersTrue, prefetch_factor4, )num_workers和pin_memory是最容易被随手抄错的地方。num_workers并不是越大越好它取决于 CPU 核心数以及__getitem__的耗时。如果每个样本只需要做一次简单的张量变换开 16 个 worker 反而会花费大量时间在进程通信上CPU 占用率飙升但训练速度并没有提升。pin_memory的作用是把 CPU 侧的数据固定在页锁定内存里这样拷贝到 GPU 时可以走更快的内存传输路径。它适合 GPU 训练场景但如果你是在 CPU 上跑实验开pin_memoryTrue基本没有意义反而会增加内存成本。2.2 数据加载参数的取舍与实测经验下面是我在不同机器上实测后比较稳定的配置场景num_workerspin_memorypersistent_workersprefetch_factor小数据集、CPU 训练0 或 1FalseFalse默认中等数据集、单 GPU4 到 8TrueTrue2 到 4大数据集、多 GPU8 到 12TrueTrue4 到 8每个样本需要重 IO尽量调高TrueTrue4 以上persistent_workersTrue在 PyTorch 1.9 之后的版本里很值得开。默认情况下每个 epoch 结束之后 worker 进程会被销毁下个 epoch 重新创建。开这个参数以后worker 进程会保留下来省掉频繁创建销毁进程的开销。小数据集上感知不明显但图像、视频这类大样本集上效果立竿见影。prefetch_factor决定每个 worker 提前预取多少个 batch。调太高会让内存暴涨调太低则可能在每个 step 之间留出等待时间。我的做法是先从默认值起步然后用nvidia-smi观察 GPU 利用率。如果 GPU 利用率一直低于 90%并且 CPU 并没有满载那大概率是预取不够优先调prefetch_factor再考虑num_workers。提示shuffleTrue加num_workers0的组合下每个 worker 持有不同的随机数流。如果只给 Dataset 设置了一次全局种子多个 worker 之间可能生成完全相同的增强序列导致同一个 batch 里出现大量重复的增强样本。这个问题非常隐蔽比较稳妥的做法是设置torch.utils.data.default_collate之前手动给每个 worker 单独设置随机源或者在__getitem__内部使用实例级np.random.RandomState。2.3 验证集 / 训练集的随机性边界验证集的 DataLoader 和训练集最大的区别是验证集不要shuffleTrue。很多人想当然地认为“每次验证时打乱一下也没关系”但它在两种情况下会出问题一是评估指标需要和上一次对比比如 A/B 实验中两个版本模型的准确率差异到底是因为模型变好了还是因为这次抽到的样本碰巧简单二是做混淆矩阵或做错误样本分析时样本顺序一变排查问题的思路全乱。如果你需要固定打乱方式有一个更干净的做法显式传入一个生成器。val_loader DataLoader( val_set, batch_size64, shuffleFalse, )训练集上随机种子管理也更推荐单独生成一个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_noneTrue) loss.backward() clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step()注意optimizer.zero_grad(set_to_noneTrue)这个细节。默认的zero_grad()会把梯度张量置零set_to_noneTrue则是直接把梯度设为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_norm1.0) optimizer.step() optimizer.zero_grad(set_to_noneTrue)这里有个很容易忽略的点loss 需要除以accum_steps。因为每个小批次都会触发反向传播梯度会累加除完以后再累积等效 batch 的梯度幅度才和原来一致。梯度累积也不是银弹。它不能让每个 step 的等待时间消失只是把更新频率降下来。如果显存小得连一个小 batch 都塞不进那果断换更小的输入尺寸或更轻量的骨干网络更实际。4. 优化器与学习率调度的配合套路4.1 从 SGD 到 AdamW 的适用场景选优化器不是选“最新的那个”而是选“符合当前任务收敛特性的那个”。理解吃透了才能调好单单照着别人的代码抄很难应对变化。优化器特点适合场景SGD Momentum收敛路径更可控泛化表现可以在 CV 任务上很突出图像分类、需要精细调 lr 的任务Adam自适应学习率调参门槛低前中期收敛快文本、序列、跨模态等非 CV 任务快速出基线AdamW修正了 Adam 的权重衰减实现Transformer 类模型、大规模预训练/微调LBFGS二阶优化内存占用高不稳定小规模、非深度学习场景比如曲线拟合很多人只记得 Adam 是默认选项却忽略了它内部维护了一阶矩和二阶矩两个状态。这两个状态是从训练开始累积的如果你在中途切换优化器或者改了学习率而矩估计还在沿用旧的分布收敛行为会变得很奇怪。这也是为什么断点续训时优化器状态必须同步保存。4.2 Warmup 余弦退火的组合实践训练初期用一个很小的学习率然后逐步上升到峰值这个策略叫 warmup。它存在的目的是避免模型在参数还很随机时被过大的梯度带偏。对 Transformer 这类对学习率敏感的模型warmup 几乎是标配。对小规模 CNN不一定需要但也不会有害。warmup 余弦退火是效果比较稳的组合之一。代码层面的实现方式有很多种比较直观的写法是继承torch.optim.lr_scheduler.LRSchedulerclass WarmupCosineScheduler(LRScheduler): def __init__(self, optimizer, warmup_steps, total_steps, min_lr1e-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_locationcpu加载后再手动移动到目标设备ckpt torch.load(checkpoint.pth, map_locationcpu) 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 Falsecudnn.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, devicecuda) 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 耗时。这些数据平时不起眼等真的需要定位问题时它们就是最可靠的参照系。