刚接触PyTorch时很多人都会有这样一个瞬间模型结构写好了损失函数也选定了结果到optimizer torch.optim.Adam(model.parameters(), lr0.001)这一行就卡住了——不知道学习率该设多少不知道SGD和Adam到底差在哪更不知道后面那一堆momentum、weight_decay、betas参数都是干什么用的。网上搜了一圈要么是文档翻译腔要么就是一上来就甩一堆公式。这篇文章打算换一种讲法结合我这些年用PyTorch做模型训练的实际经验把torch.optim这个库从头到尾拆一遍它到底是什么、每类优化器在干什么、代码层面怎么写才对以及哪些坑是我真的踩过的。无论你是刚入门深度学习的小白还是已经跑过不少模型但一直对优化器“知其然不知其所以然”的选手这篇应该都能给你一些启发。1. 优化器不是“一行代码”而是训练过程的“方向盘”1.1 优化器到底在解决什么问题先抛开所有术语回到最根本的问题训练一个神经网络本质上是在做什么答案是找一组参数让损失函数的值尽量小。模型里的每一个权重张量都是多维的几百万甚至几千万个参数构成一个极高维的空间损失函数在这个空间里是一个起伏不平的曲面。训练的过程就是在这个曲面上一路摸索着往下走找到那个让损失值足够低的“谷底”。这个“往下走”的动作由谁来执行就是优化器。神经网络前向传播算出损失loss.backward()算出每个参数对应的梯度——梯度指向的是损失上升最快的方向取反就是下降方向。但算出来的这些梯度只是“一个方向”具体怎么沿着这个方向走、走多远、要不要借鉴以前走过的路、要不要每个参数区别对待这些策略全部由优化器决定。torch.optim就是PyTorch官方提供的优化器集合。它不是一个单独的算法而是一个统一的管理框架你传入模型的参数它负责维护参数状态在每一步step()时根据梯度对参数做更新。1.2 从 loss.backward() 到 optimizer.step() 的完整链路我见过不少新手把loss.backward()和optimizer.step()当成一个固定搭配少写一个就报错或者干脆loss.backward()都不写模型的loss就是不降。这里把这条链路完整拉一遍前向传播输入数据经过网络每一层算出预测值和损失值。loss.backward()PyTorch的自动微分引擎从损失开始反向遍历计算图把每个参数的梯度param.grad填上。optimizer.step()优化器读取param.data和param.grad根据自己内部维护的状态比如动量缓冲、二阶矩估计按更新规则修改param.data。optimizer.zero_grad()把上一步留下的梯度清零。如果不做这一步下一次backward()时新梯度会和旧梯度累加——这在某些特殊场景比如梯度累积是故意为之但普通训练里必须清零。注意backward()和step()是严格分开的。这意味着你可以对所有参数统一求梯度然后用不同的优化策略、不同的学习率去更新不同层。这一点在后面的参数组配置中非常有用。1.3 没有优化器会怎样为了让你更直观地理解优化器干了什么活可以看一眼“没有优化器”的训练怎么写# 手动实现梯度下降不用 torch.optim learning_rate 0.01 for epoch in range(num_epochs): loss loss_fn(model(x), y) loss.backward() with torch.no_grad(): for param in model.parameters(): param - learning_rate * param.grad model.zero_grad()这段代码其实就是最朴素的SGD。你会发现它和torch.optim.SGD在做的事本质一样只不过后者把“每个参数怎么更新”的策略集中管理了。一旦你想加动量、加自适应学习率、对每层用不同学习率、在更新时加入权重衰减手写代码将迅速失控。torch.optim的价值不在于“能更新参数”而在于把各种被验证有效的优化策略都封装好了你只需要换一个类名就能切换完全不同的训练行为。2. 先搞懂梯度下降的三种形态选型才不纠结2.1 三种梯度下降的直觉理解优化器家族里所有算法本质上都是梯度下降的变体。梯度下降按“每次用多少数据来算梯度”分成三派批量梯度下降BGD每次更新用全量训练数据。这个做法在数据集很大时根本不现实——几百万张图算一次梯度显存和算力都吃不消而且更新方向固定容易卡在局部极小点附近出不来。随机梯度下降SGD每次只用一条样本算梯度。计算开销小而且单样本梯度噪声大反而有机会跳出局部极小。但问题是梯度噪声太大损失曲线会像心电图一样疯狂跳动收敛路径蜿蜒曲折。小批量梯度下降Mini-batch GD每次用一个batch的数据算梯度。这是工业界和学术界实际采用的形式。PyTorch里的DataLoader每次产出一个batch配合优化器更新本身就是小批量梯度下降。所以你在PyTorch里写的torch.optim.SGD严格来说是Mini-batch SGD。2.2 动量机制到底在做什么单纯的小批量SGD有个明显毛病在损失曲面的狭长山谷地带梯度在某个方向上来回振荡、进退两难收敛极慢。为了抑制振荡、加速收敛动量Momentum机制被引入——它不只看当前这一步的梯度还把历史梯度的“惯性”一起考虑进来。物理类比是这样的你推一个很重的铁球下坡铁球不会每走一步都完全改变方向而是保持着之前运动的一部分惯性。如果梯度方向在某个维度上反复变化惯性会让振幅减小如果持续朝一个方向走惯性会帮助加速。PyTorch里SGD的momentum参数就是控制这个“惯性”的强度经验上设在0.9附近。加了动量的收敛速度通常比普通SGD快得多也顺滑得多。2.3 学习率是最需要花心思的超参数学习率可能是整个优化器里最重要的超参数没有之一。它决定了每一步沿着梯度方向迈多大步子。步子太小训练极度缓慢而且很容易陷入极小点附近因为梯度越来越小走不动了。步子太大可能在损失曲面两侧来回震荡甚至直接“跨过”谷底Loss变成一个巨大的NaN——我在训练图像分类模型时的确见过Adam配合0.1的学习率第一个epoch结束loss直接变NaN检查了很久才发现是学习率问题。学习率的选择通常跟优化器有关SGD系一般从0.01到0.1起步配合Learning Rate Scheduler做衰减Adam系一般默认0.001在大多数任务上都算一个稳定的起点。但这个值不是绝对的需要根据loss曲线来调整。我见过有人用0.01的Adam在Transformer类模型上训练得很好也见过0.001的SGD完全跑不动。关键是要理解学习率的含义然后学会观察loss曲线而不是死记硬背哪个数值。3. torch.optim核心API拆解从创建到参数组的完全指南3.1 构造优化器的两种写法torch.optim的构造方式非常简单把需要优化的参数传进去再指定学习率import torch import torch.nn as nn model nn.Linear(10, 2) # 方式一传入所有参数 optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9) # 方式二传入参数组列表后面细说 optimizer torch.optim.Adam([ {params: model.features.parameters()}, {params: model.classifier.parameters(), lr: 0.0001} ], lr0.001)第一种写法就是全模型统一策略。第二种则把网络不同部分拆成不同的“参数组”每组可以有自己的学习率、权重衰减等设置——这在迁移学习和微调场景中几乎是必备技能。比如用预训练模型做下游任务时通常希望backbone的学习率小一些保留已学特征而新加的分类头学习率大一些更快适应新任务。3.2 param_groups优化器内部的数据结构优化器内部维护一个optimizer.param_groups列表每个元素是一个dict描述一组参数的优化策略。常见配置项包括params这个参数组包含的参数列表。lr该组的学习率。momentum、betas、eps动量或自适应算法的参数。weight_decay权重衰减系数L2正则化的等价实现。amsgradAdam变体是否使用AMSGrad。这个结构最实用的地方在于训练中途可以动态修改学习率不用重新创建优化器。# 训练到一半把第一个参数组的lr从0.001降到0.0001 optimizer.param_groups[0][lr] 0.0001许多学习率调度器torch.optim.lr_scheduler内部就是靠修改param_groups里的lr字段来实现衰减的。理解了这一点你对调度器的工作原理就基本有数了。3.3 三个必须记住的方法zero_grad、step、state_dictoptimizer.zero_grad()清除梯度。注意它默认是set_to_noneFalse也就是把梯度清零时实际上是把梯度张量置零不是置None。因为梯度是稀疏的直接置None更省内存所以现在有个写法是optimizer.zero_grad(set_to_noneTrue)。在一些模型里把set_to_noneTrue能减少不必要的计算。optimizer.step()执行一次参数更新。绝大多数情况直接调用即可。但有一个特殊场景叫“梯度累积”gradient accumulation意思是当显存不够、无法用大的batch_size时可以用多个小batch累积梯度再更新。做法是# 每4个batch做一次更新 for i, batch in enumerate(dataloader): loss loss_fn(model(batch)) loss loss / 4 # 平均化累积的梯度 loss.backward() if (i 1) % 4 0: optimizer.step() optimizer.zero_grad()这种写法非常实用——相当于用多个小batch模拟一个大batch的效果代价是训练总时间略增。optimizer.state_dict()和optimizer.load_state_dict()保存和加载优化器的内部状态。注意这里保存的不只是学习率还包括动量缓冲、Adam的一阶/二阶矩估计等。如果训练中断要在断点续训光保存模型权重是不够的还要保存优化器状态否则恢复训练后loss曲线的行为会出现明显的“跳变”。3.4 weight_decay与L2正则化的那点事weight_decay是每个优化器都有的参数它实现的效果是在每次更新时把权重朝着0的方向拉一点。数学上它等价于在原始损失函数上加上一个L2惩罚项。很多人担心一个问题Adam里加了weight_decay是不是就等价于L2正则化严格来说传统Adam的实现是把weight_decay直接加到梯度上再进入自适应计算流程这和真正的L2正则存在细微差别而AdamW换了一种实现——把权重衰减从自适应流程中解耦直接对权重做衰减。这就是为什么现在预训练大模型普遍用AdamW而不是Adam它对权重衰减的处理更“干净”在实际任务里表现也更稳定。经验值上weight_decay在图像分类任务里常用1e-4到5e-4在AdamW里常用0.01到0.1。但这个值对任务很敏感我不建议无脑套最好在验证集上做个小范围搜索。4. 主流优化器横向对比SGD、RMSprop、Adam、AdamW的实测选型逻辑4.1 SGD与SGDMomentum简单可靠但需要调参PyTorch里的torch.optim.SGD实际是带有很多选项的SGD族参数说明经验值lr学习率0.01~0.1配合衰减率momentum动量系数0.9附近dampening动量阻尼一般0weight_decay权重衰减1e-4~5e-4nesterov是否使用Nesterov动量常TrueSGDMomentum是我在图像分类模型上的首选之一。它的优点是对超参数不那么敏感可解释性强训练到后期配合余弦退火效果很稳。缺点是需要比较精细的学习率调度起步阶段如果没有好的lr范围搜索收敛速度比Adam类慢不少。4.2 RMSprop针对不同参数“因地施策”SGD对每个参数都用同一个学习率这在网络各层梯度尺度差异很大时会很痛苦——某些参数梯度很陡某些参数梯度很平统一的步长很难兼顾。RMSprop的思路是对每个参数用它的历史梯度的平方的指数移动平均来归一化当前的更新步长。梯度大的方向步长自动缩小梯度小的方向步长自动放大。这个策略让训练更平稳尤其适合RNN这类梯度尺度波动大的结构。不过RMSprop默认学习率一般设0.001但要真的想让模型快速收敛往往得手动调到一个合适的范围。4.3 Adam与它的 betas 参数自适应学习率的典型代表AdamAdaptive Moment Estimation是RMSprop和动量机制的合体既保留了每个参数的梯度平方归一化又加入了一阶动量作为“速度”。它在绝大多数任务上开箱即用学习率默认0.001基本能跑到一个不错的水平所以成了研究者的默认选择。Adam的公式核心是维护两个状态一阶动量估计exp_avg历史梯度的指数移动平均反映梯度方向。二阶动量估计exp_avg_sq历史梯度平方的指数移动平均反映梯度尺度。betas就是这两个移动平均的衰减系数默认是(0.9, 0.999)。betas[0]0.9意味着一阶动量平均最近10步左右的梯度。betas[1]0.999意味着二阶动量平均最近1000步左右的梯度平方。eps是为了防止除零加的极小常数默认1e-8。在混合精度训练时eps可能需要调大一点比如1e-7或1e-6否则数值稳定性可能出现问题。Adam在NLP、Transformer类模型、生成模型等场景里通常是首选。它的缺点是后期训练时可能因为自适应学习率步长过大而难以收敛到极小值所以在一些CV任务里人们会先用Adam快速探索再切SGD做精调。4.4 AdamW解决Adam与权重衰减的矛盾AdamW和Adam的差异我在3.4里提过——解耦权重衰减。简单回顾一下在传统Adam里weight_decay是加在梯度上的然后和梯度一起被exp_avg_sq归一化。这意味着权重衰减力度会被自适应机制影响衰减量不再是“和参数大小成固定比例”。AdamW直接把权重衰减从梯度流程中拿出来在更新参数时单独执行param.data - lr * weight_decay * param.data # 解耦的权重衰减这种做法让权重衰减的行为可预测训练更稳定。事实上Hugging Face的transformers库训练BERT、GPT系列时默认优化器就是AdamW学习率调度大多用带warmup的线性衰减。如果是做NLP或大模型相关的工作请优先考虑AdamW。4.5 我的选型经验总结场景推荐优化器为什么图像分类/经典CV训练时间充足SGD(Momentum0.9)配合CosineAnnealingLR收敛位置好泛化强NLP/Transformer/大模型预训练AdamW数值稳定配合warmup效果好快速验证一个模型能否跑通Adam开箱即用参数少收敛快强化学习/时序模型RMSprop或Adam对梯度噪声和尺度变化更鲁棒显存有限、batch较小Adam/AdamW自适应机制能容忍噪声梯度这份表格不是“唯一正确答案”它只是我经历过的项目里最常走通的组合。优化器选型没有银弹关键是你理解了自己的模型和数据然后愿意花时间看loss曲线。5. 我踩过的优化器相关的坑以及排查链路5.1 Loss变成NaN的第一现场排查几乎每个炼丹师都经历过训练到某个steploss突然变成NaN之后所有参数也跟着变NaN整个训练彻底报废。我遇到的最常见原因是学习率过大。以Adam为例如果设置lr0.1在前面提到的自适应机制下某些参数的方向步长会剧烈震荡数值溢出是分分钟的事。排查链路应该是先把学习率降到1e-4甚至1e-5看是否还出现NaN。检查输入数据是否有NaN或无穷值torch.isnan(x).any()。检查梯度是否有NaN在optimizer.step()前打印param.grad。检查损失函数是否对数值不稳定比如计算softmax时对数值溢出可以考虑换用log_softmax加NLLLoss或者直接用CrossEntropyLoss。如果用了混合精度检查scaler.scale(loss)的操作顺序以及scaler.step(optimizer)和scaler.update()是否配对。实际上一旦loss变成NaN模型的所有参数基本都污染了。想恢复几乎不可能最好的办法是从上一个正常保存的checkpoint重新开始训练并修正超参数。所以我的习惯是训练时每几个epoch保存一次checkpoint并且把“最近一个正常loss的样本”记录到日志里方便判断是从哪个点开始出问题的。5.2 梯度消失与梯度爆炸查看参数梯度的范数优化器调参时我强烈建议在训练前几步加一段调试代码打印梯度范数total_norm 0.0 for p in model.parameters(): if p.grad is not None: param_norm p.grad.data.norm(2) total_norm param_norm.item() ** 2 total_norm total_norm ** 0.5 print(fgradient norm: {total_norm:.6f})梯度范数能直观地反映训练是否健康。如果梯度范数突然飙升到1e5以上那就是梯度爆炸常见的处理手段有降低学习率。使用梯度裁剪gradient clippingPyTorch里内置了torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。检查是否有残差连接、归一化层避免深层网络信号逐层放大。如果梯度范数小到1e-8以下那就是梯度消失了常见原因包括激活函数选择不当比如深层网络全用sigmoid、初始化方式不合适。RAdam、AdamW这类自适应优化器能在一定程度上缓解梯度消失——因为自适应机制会放大梯度较小的方向的更新步长。5.3 用错了参数组微调模型效果还不如随机初始化有一段时间我微调BERT做文本分类直接把整个模型的所有参数用同一个学习率喂给AdamW跑了20个epoch准确率反而比从头训练的基线差。后来排查才发现问题不小BERT底层的预训练表示其实已经很好了学习率太大会把它们“洗掉”而新加的分类头因为随机初始化本来就需要更大的学习率快速收敛。用一个统一的学习率要么底被洗坏要么头部学习太慢。解决办法就是参数组optimizer torch.optim.AdamW([ {params: model.encoder.parameters(), lr: 2e-5}, {params: model.classifier.parameters(), lr: 2e-4} ], weight_decay0.01)这样backbone用较小的学习率微调分类头用较大的学习率快速适配。经验上backbone学习率如果设成1e-5到3e-5头部设成1e-4到3e-4绝大多数微调任务都比较稳。如果你对这块不太确定可以在param_groups里设置不同组然后打印每一组的lr确认生效。5.4 保存和恢复优化器状态断点续训的坑有一次训练一个GAN大概跑到第40个epoch时服务器被重启了。我从第30个epoch的checkpoint恢复只加载了model权重没有加载optimizer的state_dict。结果恢复后的loss曲线和之前完全不连续甚至出现了反复的周期性波动。原因在于Adam内部的一阶/二阶动量状态丢失了。从头开始等于“失忆”自适应学习率需要重新积累历史信息所以恢复后会有一个明显的适应期。正确的断点保存方式是torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict() if scheduler else None, loss: loss.item(), }, checkpoint_path)加载时checkpoint torch.load(checkpoint_path) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) if scheduler and checkpoint.get(scheduler_state_dict): scheduler.load_state_dict(checkpoint[scheduler_state_dict])另外还要注意如果在保存状态下变化了学习率比如已经跑了几个epoch的CosineAnnealingLR加载后调度器的状态也会影响后续学习率变化轨迹。所以调度器状态最好一并保存。6. 可直接抄作业的训练模板与学习率搭配方案6.1 一个完整的训练代码骨架下面这个模板基本覆盖了我日常训练的需求从优化器创建到学习率调度、断点保存一步到位import torch import torch.nn as nn from torch.optim.lr_scheduler import CosineAnnealingLR model nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)) optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay0.01) # 使用余弦退火总训练轮数设为50 scheduler CosineAnnealingLR(optimizer, T_max50, eta_min1e-5) criterion nn.CrossEntropyLoss() num_epochs 50 start_epoch 0 # 如果有checkpoint就加载进来 checkpoint_path checkpoint.pt try: ckpt torch.load(checkpoint_path) model.load_state_dict(ckpt[model_state_dict]) optimizer.load_state_dict(ckpt[optimizer_state_dict]) scheduler.load_state_dict(ckpt[scheduler_state_dict]) start_epoch ckpt[epoch] 1 except FileNotFoundError: pass for epoch in range(start_epoch, num_epochs): model.train() train_loss 0.0 for x, y in train_loader: optimizer.zero_grad(set_to_noneTrue) out model(x) loss criterion(out, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() train_loss loss.item() * x.size(0) scheduler.step() avg_loss train_loss / len(train_loader.dataset) print(fEpoch {epoch1}/{num_epochs} | Loss: {avg_loss:.6f} | LR: {optimizer.param_groups[0][lr]:.2e}) if (epoch 1) % 5 0: torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), scheduler_state_dict: scheduler.state_dict(), loss: avg_loss, }, checkpoint_path)这个模板有几个细节值得注意optimizer.zero_grad(set_to_noneTrue)比默认的置零方式省内存、省时间。梯度裁剪放在backward()之后、step()之前防止梯度爆炸。scheduler.step()在epoch结束时调用但像StepLR、ReduceLROnPlateau这类调度器可能需要传入验证指标或调整调用位置务必看清楚文档。学习率打印用的是optimizer.param_groups[0][lr]这样任何时候想确认实际生效的学习率都很方便。6.2 学习率调度器与优化器的搭配思路优化器和学习率调度器是“前轮后轮”的关系优化器决定怎么走调度器决定每一步走多快。常见搭配有StepLR/MultiStepLR每隔固定epoch把学习率乘以一个因子比如0.1。适合SGD类简单粗暴。ReduceLROnPlateau当验证集指标不再下降时降低学习率。适合大多数任务尤其没有确定总epoch数的场景。CosineAnnealingLR/CosineAnnealingWarmRestarts余弦退火训练后期学习率平滑趋近一个最小值。配合SGD表现很好也是很多CV比赛的标配。带warmup的线性衰减大模型预训练标配。早期学习率从0线性升到一个指定值再线性衰减到0。PyTorch里可以用LinearLR配合SequentialLR或LambdaLR手动实现。我这里想额外提醒一下warmup不是只有大模型才需要。如果你的模型比较深或者batch比较大直接用大学习率起步很容易在最初几个step把参数推到一个很差的位置之后再怎么调都难以回来。加一个几百步的warmup让学习率从小到大过渡训练的稳定性会显著提升。# 一个简单的线性warmup 余弦退火组合 def warmup_cosine_lr(epoch, warmup_epochs5, total_epochs50, lr_max1e-3, lr_min1e-5): if epoch warmup_epochs: return lr_max * (epoch 1) / warmup_epochs else: progress (epoch - warmup_epochs) / (total_epochs - warmup_epochs) return lr_min 0.5 * (lr_max - lr_min) * (1 torch.cos(torch.tensor(progress * 3.1415926))) scheduler torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambdalambda e: warmup_cosine_lr(e) / lr_max)注意LambdaLR的lr_lambda返回的是学习率缩放因子而不是直接的学习率。所以我把目标学习率除以lr_max让它相对于初始学习率做缩放。写成这样之后每个epoch调度器会按这个函数更新学习率。6.3 我的实际经验和一点习惯最后分享几个我自己的习惯大家参考第一个训练前先跑过拟合检查。用一小部分训练数据比如1到2个batch把模型训到loss趋近于0。如果这一步都做不了先别怪优化器模型结构或者数据管道大概率有问题。第二个关注学习率在训练中的实际走势。不要只盯着初始学习率。很多问题藏在调度器里比如ReduceLROnPlateau的patience设得太大导致学习率下降过慢或者CosineAnnealingLR的T_max设得比实际epoch数小导致学习率提前降到最低点最后几个epoch在无意义地“滑行”。第三个也是最重要的优化器和学习率要一起调不能只动其中一个。换优化器不是换一行代码那么简单——从Adam切到SGD时学习率必须相应改变学习率调度策略也要重新设计甚至权重衰减值也要重新考虑。我见过有人把Adam换成SGD其他参数原封不动结果loss曲线乱得没法看然后得出结论说“SGD真难用”。第四个保存checkpoint时多保留一份“最近一次正常状态”。训练中途出现NaN或者loss异常时可以用它快速回滚而不是从头再来。这个习惯帮我在无数次实验事故里省下了大量时间。