新手在 PyTorch 里写完模型定义卡住的第一个地方往往是优化器到底是怎么把损失函数变成参数更新的说实话我刚接触的时候也疑惑过zero_grad()是不是多此一举Adam 的参数更新为什么内部还要分好几个步骤。这篇内容就是围绕PyTorch 优化器参数更新步骤展开的从优化器的设计理念、zero_grad/backward/step的完整生命周期到 SGD、Momentum、Adam、AdamW 的更新公式差异再到实战里的学习率调度和梯度异常排查。不管你是刚入门 PyTorch 的小白还是已经能跑通训练但总感觉调参靠运气的同学这篇都能给你一套可以“抄作业”的实操思路。我会尽量用大白话把原理讲清楚中间穿插我实际踩过的一些坑希望能帮你少走弯路。1. 优化器到底在干什么不只是“下一步往哪走”1.1 从一个最小闭环理解梯度下降先做一个最简单的类比。假设你现在站在一座山坡上想尽快下到山脚策略很简单往坡度最陡的方向跨一步再重新判断方向继续迈步。在训练神经网络时“山脚”是损失函数的最小值点“坡度”是损失对参数的梯度模型参数就是你的位置。更新参数就是迈出一步这一步的方向是梯度的反方向步长由学习率lr控制。数学上最朴素的形式是这样parameter_new parameter_old - lr * gradient这就是原始梯度下降。PyTorch 里所有优化器不管内部逻辑多复杂本质都是在“怎么走这一步”上做文章。区别在于是要不要带历史动量要不要对不同参数用不同步长要不要对参数本身加约束这一层想清楚了后面看源码会轻松很多。1.2 PyTorch 优化器的协作方式引用而不是复制PyTorch 优化器的第一个设计关键是它跟模型参数之间存在引用关系不是值拷贝。当你写下面这行代码时optimizer torch.optim.SGD(model.parameters(), lr0.01)优化器拿到的是model中每个参数对象的引用并把这些引用存放在param_groups里。当你调用optimizer.step()优化器会直接原地修改param.data不会返回一组新参数再让你重新赋值模型。这一点和 TensorFlow 早期那种“每次更新返回新参数”的流程不同刚开始切换框架的人容易懵。另一个容易忽略的点是优化器只改param.data不改param.grad。梯度在backward()里被计算出来后存到param.gradstep()全程读取这个值但不会主动清空它。清空梯度的责任在zero_grad()上。所以训练循环里的顺序一定不能乱optimizer.zero_grad() # 先清空上一次残留的梯度 loss.backward() # 前向传播后计算梯度 optimizer.step() # 用梯度更新参数这三步是 PyTorch 训练循环里出现频率最高的代码片段无数报错都源于顺序颠倒或者漏了其中一步。尤其zero_grad()不是可有可无的“仪式”它是为了保证“本 batch 的梯度干净”。2. 核心机制一个 batch 里参数更新的完整生命周期2.1zero_grad()为什么梯度默认是累加的PyTorch 的梯度默认累加也就是说如果不手动清零多次backward()会把梯度不断叠加。第一次损失backward()之后param.grad是第一次的梯度第二次再对另一个 loss 执行backward()param.grad变成两次梯度之和。这个设计不是失误而是特意留出来的“梯度累积”能力当显存不够支撑一个大 batch 时你可以拆成小 batch连续计算几次梯度、累加后再统一更新一次参数模拟大 batch 效果。但在常规训练里这对新手是个隐藏炸弹。最常见的错误就是写了backward()和step()却忘了zero_grad()结果每个 batch 的梯度都叠加到上一次的梯度上。损失曲线会变得非常奇怪要么一直下跌要么突然爆炸。排查这类问题时先确认有没有在每个 batch 开始时清空梯度。2.2backward()梯度是怎么“落”到参数上的loss.backward()背后是一条反向传播链。PyTorch 在张量操作时构建了一个动态计算图loss 是计算图末端的一个节点。backward()从 loss 出发按链式法则逐层往前传播梯度每经过一个带requires_gradTrue的参数就把计算出的梯度累加到该参数的.grad属性上。值得注意的点只有叶子张量的.grad会被保存。很多中间量是计算过程产生的临时张量反向传播时用过就被释放掉了。所以你在step()之后如果打印某个中间层输出的grad大概率是None这是正常现象。真正需要检查的是模型最外层的参数比如model.fc.weight.grad。另外backward()执行完后计算图默认会被释放。这意味着如果你在backward()之后再去调用loss相关的量可能会报“尝试访问已释放张量”的错误。想保留计算图继续二次求导需要给backward(retain_graphTrue)但普通训练里极少用到。2.3step()优化器内部在做什么optimizer.step()是所有优化器实现的核心。不同类型的优化器在里面做的事情不一样但整体流程可以归纳为五步遍历param_groups里的每一组参数遍历该组里每个参数param读取param.grad根据优化器类型计算参数更新量执行param.data - update_amount。以 PyTorch 自带的 SGD 为例无动量时step()内部大致等价于for group in optimizer.param_groups: lr group[lr] for p in group[params]: if p.grad is None: continue p.data.add_(p.grad, alpha-lr)这里用add_(other, alpha...)实现的是p.data p.data - lr * p.grad。注意它直接操作p.data而不是p本身。因为如果直接对p做运算可能会破坏参数张量的 autograd 追踪状态。进阶玩家如果自己写优化器这一点要格外小心。Adam 的step()就更复杂了。它会为每个参数维护两个状态一阶动量exp_avg和二阶动量exp_avg_sq。每次step()读入梯度g更新动量再计算偏差校正最后组合出更新量exp_avg beta1 * exp_avg (1 - beta1) * g exp_avg_sq beta2 * exp_avg_sq (1 - beta2) * g * g bias_correction1 1 - beta1 ** t bias_correction2 1 - beta2 ** t step_size lr * sqrt(bias_correction2) / bias_correction1 parameter_update step_size * exp_avg / (sqrt(exp_avg_sq) eps)这些状态都存在optimizer.state[p]里。每个参数都对应一个字典里面是exp_avg、exp_avg_sq等张量。这也是为什么 Adam 优化器的显存占用几乎是模型参数的两倍。后面讲内存排查时还会再提到。3. 常用优化器更新公式拆解与选型3.1 SGD、Momentum、RMSProp 的更新公式朴素 SGD我上面已经写了它就是θ θ - lr * g。优点是简单、稳定缺点是收敛慢而且在“峡谷形”的损失面上容易来回震荡。SGDMomentumMomentumSGD在朴素 SGD 上加了“惯性”。先维护一个速度变量vv_t momentum * v_{t-1} g θ_t θ_{t-1} - lr * v_t注意 PyTorch 实现里v的初值是零向量momentum通常取 0.9。这个“惯性”能让优化器在梯度方向连续时加快步伐在梯度方向反复摆动时互相抵消。实际效果就是训练前期收敛快不容易卡在局部极小点。我在图像分类任务里最常用的就是 SGD 加 momentum。RMSProp为解决不同参数更新步长不一致的问题RMSProp 维护一个梯度的滑动平方平均E[g^2]_t beta * E[g^2]_{t-1} (1 - beta) * g^2 θ_t θ_{t-1} - lr * g / (sqrt(E[g^2]_t) eps)它的效果是梯度大的方向除以一个较大的数步长被压小梯度小的方向除以较小的数步长相对变大。这个思路对非平稳目标函数很有效也是后来 Adam 的一只“轮子”。3.2 Adam 与 AdamWTransformer 时代的事实标准Adam 可以理解为 Momentum 加上 RMSProp一阶动量负责平滑梯度方向二阶动量负责自适应调节每个参数的学习步长。它的完整更新公式我在 2.3 节已经列出来了。Adam 最大的优点是超参数少、默认配置通常就能跑出不错的结果lr1e-3betas(0.9, 0.999)eps1e-8。我在刚开始做 NLP 任务时几乎无脑用 Adam。ButAdam 有个被反复讨论的问题它把权重衰减weight decay和 L2 正则混在一起。在标准 Adam 实现里L2 正则项会先被加进梯度再被一阶、二阶动量“自适应”地处理这会导致正则效果被削弱。AdamW 则是把权重衰减直接从梯度里分离出来在更新参数时单独执行θ_t θ_{t-1} - lr * weight_decay * θ_{t-1} - adaptive_update这个修正对于大模型和 Transformer 效果很关键这也是为什么 Hugging Face 的 Trainer 默认推荐 AdamW。现在很多人说“用 Adam”其实在代码里写的是AdamW这点在组会讨论时特别容易起争执。3.3 优化器选型的经验法则选优化器这事没有绝对唯一答案但根据任务类型可以先圈定范围场景推荐优化器默认参数经验入门教学、简单线性模型、小数据集SGD / SGDMomentumlr 0.01 到 0.1momentum 0.9图像分类、目标检测CNN 主流任务SGD Momentum / AdamWSGD lr 0.01~0.1AdamW lr 1e-3 到 1e-4NLP 模型、Transformer、LLM 微调AdamWlr 1e-5 到 1e-4配合 warmup强化学习、非平稳目标RMSProp / AdamRMSProp lr 1e-3Adam 默认显存紧张、参数超大AdamW 关闭 grad 缓存或用 8bit 优化器结合梯度检查点我个人有一条很朴素的判断标准如果任务的环境稳定、训练数据多SGD 系列往往能收敛到更好的极小值而且泛化性更强如果任务环境不稳定、需要快速迭代出结果Adam 系列更省心。不要迷信“Adam 永远最好”我在一个图像分类项目上做过对比同样的模型和批大小SGD Momentum 在训练后期损失降得比 Adam 更稳。4. 实战从训练循环到学习率调度4.1 一次完整训练循环的实现直接给出一段可以在小数据集上跑的简化代码。这里用 MNIST 结构做说明一个简单 CNN、交叉熵损失、SGD 优化器。重点是看清optimizer在循环里的位置。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader model SimpleCNN() optimizer optim.SGD(model.parameters(), lr0.05, momentum0.9) criterion nn.CrossEntropyLoss() train_loader DataLoader(train_dataset, batch_size128, shuffleTrue) model.train() for epoch in range(10): total_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() print(fepoch {epoch1}, loss{total_loss/len(train_loader):.4f})这段代码里有几个容易出问题的地方。第一model.train()必须记得调它会影响 Dropout 和 BatchNorm 的行为如果忘了训练和推理结果都会乱。第二optimizer.zero_grad()一定在backward()之前调用我习惯把它放在每个 batch 开头因为我对“忘记清零”有阴影。第三只要用 loader 遍历数据每个 epoch 结束时模型的参数已经更新了一轮新的 epoch 直接继续。4.2 学习率调度器与优化器的联动学习率不是越固定越好。训练初期希望学习率大一点快速走出起点训练后期希望学习率小一点细腻地逼近极小值。PyTorch 提供了torch.optim.lr_scheduler。最常用的组合是scheduler optim.lr_scheduler.StepLR(optimizer, step_size3, gamma0.1) for epoch in range(10): train_one_epoch() scheduler.step()StepLR每step_size个 epoch 把学习率乘以gamma。还有两个特别常用的CosineAnnealingLR按照余弦曲线从初始学习率降到接近 0适合训练时间长的任务ReduceLROnPlateau当某个监控指标通常是验证集损失连续多个 epoch 不下降时自动降低学习率。初学者最容易犯的错误是调用时机错误。scheduler.step()应该放在每个 epoch 结束后调用而不是每个 batch 后。只有OneCycleLR这类特殊调度器被设计为每 batch 调用一次。如果你混用了场景会发现学习率曲线完全错乱训练过程也会忽快忽慢。4.3 如何验证优化器是不是“真正在工作”很多人调参时只盯着 loss 看其实在训练早期更好的做法是同时观察梯度和参数范数。这里给出一个轻量级检查方法total_grad_norm 0.0 for p in model.parameters(): if p.grad is not None: total_grad_norm p.grad.norm().item() ** 2 total_grad_norm total_grad_norm ** 0.5 print(fgrad norm: {total_grad_norm:.4f}, lr: {optimizer.param_groups[0][lr]:.5f})梯度范数反映的是整个网络反向传播的“能量大小”。如果梯度范数为 0说明梯度没有传到参数上或者网络某些层的输出恒为常数。如果梯度范数呈指数级爆炸那大概率是学习率过大或者模型结构里存在初始化不当。参数范数则能侧面反映权重是否出现了异常膨胀。另外一个更直接的验证思路是让模型在几十个样本上先“强行记住”。如果一个模型连一个 batch 都无法过拟合那问题通常不在优化器而在模型设计、数据输入或 loss 计算。我每次搭建新网络都会先把训练集缩小到 32 条试着用一个 batch 把损失降到接近理论最小值。如果做不到先别急着调优化器回头检查前向传播的 bug。5. 常见问题与排查技巧实录5.1 梯度出现 NaN 或 Inf这是训练过程中最让人头疼的问题。常见成因有学习率过大导致参数发散输入数据里混入 NaNloss 函数在特定数值下产生无穷大网络初始化不当比如使用了过大的初始化权重。排查思路我建议固定顺序来第一检查输入数据里有没有 NaN加一句assert torch.isfinite(images).all()第二loss里有没有除零操作固定eps很多时候能救回来第三把学习率降到原来的 1/10 再跑第四给梯度加裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这是临时止血不是根治。我在跑 BERT 微调时遇到过梯度爆炸加个梯度裁剪之后顺利跑完但后来发现根因是学习率设成了 1e-2。做完裁剪再定位根因是更高效的方式。5.2 损失不下降或者震荡剧烈经验法则如下如果 loss 一直是nan或者inf走 5.1 的思路如果 loss 几乎不变考虑学习率太小如果 loss 剧烈抖动考虑学习率偏大、batch size 偏小或者优化器动量和二阶矩参数不匹配。还有一类隐蔽问题model.eval()和model.train()状态切换。如果你在训练过程中开了torch.no_grad()再跑optimizer.step()param.grad可能为 None优化器会静默跳过loss 自然不动。有些代码库里的 hook 会影响requires_grad状态排查时可以打印p.requires_grad和p.grad是否正常。5.3 显存增长Adam 状态吞噬内存Adam 家族优化器为每个参数维护两个额外张量exp_avg和exp_avg_sq。如果你的模型参数有 5 亿个那么 Adam 状态会额外占用2 * 5亿 * 4 字节也就是约 4GB 内存按 float32 算。加上模型本身、梯度、激活值缓存很容易把显存打爆。解决办法有几个方向改用 SGD Momentum它只保存一个动量张量状态内存减半用torch.optim.AdamW时打开foreachTrue虽然不省内存但能提速显存实在不够时用bitsandbytes提供的 8bit Adam实验里能把优化器状态压到很小用梯度检查点减少前向传播缓存但会增加计算时间。我实际测试过一个 1B 参数规模的模型仅 AdamW 的优化器状态就占掉 8GB 显存。所以在做大规模模型时选优化器不能只看收敛效果还要看内存账。5.4 训练循环中的错误速查表错误现象常见原因解决方式每个 batch 的 loss 都一样忘了zero_grad()或模型进入eval()状态按顺序zero_grad → backward → stepstep()后参数没变param.requires_gradFalse或梯度为 Nonebackward()后打印p.grad检查梯度只有最后一层有前面全是 None网络中间层断开了传播链检查是否有detach()训练时 loss 正常测试时突然异常忘了切换model.eval()测试前调用eval()用了 CPU 和 GPU 混合输入和参数设备不一致统一.to(device)5.5 关于step()前的backward()次数如果一次step()之前调用了两次backward()比如某些 GAN 训练里要先更新判别器再更新生成器那优化器会使用两次累积的梯度。有时这是故意的梯度累积有时是失误。判断标准很简单看你希望这次参数更新基于几个 batch 的梯度。如果希望每个 batch 都独立更新就必须在每次backward()前zero_grad()。如果希望累积若干 batch就在最后一个 batch 完成后再step()。6. 调参心得我用优化器踩过的那些坑最后分享几个没有写在官方文档里的经验。第一lr的优先级永远高于优化器类型。我做过一个对比实验同一模型SGD 用 0.1 收敛得很好Adam 用 1e-3 也收敛但如果把 Adam 的学习率提到 1e-2loss 立刻发散。优化器之间可以有差异但学习率的量级差一两个数量级往往决定了训练成败。所以我每次新项目的第一件事是固定优化器先几组学习率做快速扫描而不是纠结 Adam 还是 SGD。第二weight_decay不等于一定加得越多越好。在 AdamW 里太大的weight_decay会明显抑制模型表达能力尤其在 Transformer 微调场景下。我一般先设weight_decay0.01再观察验证集指标如果波动明显降到 0.001 试试。第三动量参数betas很少需要动。Adam 默认的(0.9, 0.999)覆盖了 95% 的场景。真正需要调整betas的场景是训练数据非常稀疏、或者需要更快响应梯度变化的时候新手不要在这里浪费时间。第四用余弦退火调度器时scheduler.step()的最佳位置是每个 epoch 之后。如果搭配OneCycleLR则要放在每个 batch 之后。我曾经因为把两种调度器写进同一个循环导致学习率在 epoch 中途反跳训练曲线看起来像锯齿排查了一下午才发现是调用时机错了。优化器参数更新步骤这件事代码上只是短短三行zero_grad、backward、step但真正理解它在内部做了什么能让你在调参、排障、换模型架构时都更有底气。希望这份实操笔记能帮你少踩一些我踩过的坑。