训练深度学习模型的时候学习率怎么调几乎是每个人绕不过去的一道坎。我接触学习率预热warmup这个概念比较晚早期都是直接上来就用一个较大的初始学习率开训结果在图像分类任务上经常遇到前几个epoch loss剧烈抖动甚至直接发散的情况。后来才意识到训练初期模型权重还处于随机状态梯度统计信息完全不可靠这时候直接上大步长等于让一个刚睁眼的婴儿直接去跑百米冲刺。GradualWarmupScheduler就是专门解决这个问题的工具它通过预热的几个step让学习率从极小值线性爬升到目标值再交还给后续的调度器接管。这篇文章我会把它的原理、代码实现和调参心得一次性讲透适合正在被训练不稳定问题困扰、想给训练脚本加上预热机制的同学参考。1. 学习率预热一个被低估的细节1.1 训练初期模型到底有多“脆弱”很多人对“训练不稳定”的理解是模型结构有问题、数据预处理不对、loss设计不合理。但实际排查下来相当一部分不稳定的根源就是学习率起步太高。神经网络在初始化之后权重分布虽然经过了精心设计比如Xavier、Kaiming初始化但本质上仍然是一组随机数。这组随机数对应的损失曲面在参数空间里往往是陡峭且凹凸不平的。你拿着一个大学习率迈出第一步很可能直接从一个高坡跳进一个深坑再也爬不出来。我做过一次印象很深的实验同一个ResNet结构同一份数据唯一区别是前5个epoch用不用预热。不用预热的那组前两个epoch loss在2.5到4.0之间剧烈震荡到第10个epoch才勉强稳定下来而用了预热的对照组loss像一个平缓的下坡曲线一路顺畅往下走最终收敛精度还高了将近1个百分点。这个实验让我彻底改变了对warmup的看法它不是锦上添花而是训练稳定性的基础保障。还有一个容易被忽略的点batch norm的统计量在训练初期也是极不稳定的。网络前几个batch的激活值分布还在剧烈变化此时如果学习率很大BN层 running mean和running variance的更新会被带偏形成连锁反应。预热阶段相当于给网络一个“睁眼适应期”让每一层的统计量先稳定下来再开始正经学特征。1.2 为什么带自适应优化器的模型也逃不过这里多说一句不少同学觉得“我用了Adam学习率自适应就不需要warmup了”。这个想法我一开始也有后来被实验结果打了脸。Adam虽然会基于梯度一阶矩和二阶矩来自适应调整每个参数的更新幅度但它在训练初期有一个非常关键的问题二阶矩估计不足。开头几步的梯度方差很小Adam会把这个小方差误判成“这个方向很平坦可以大步走”于是步长被异常放大。这个现象在很多文章里被称为Adam的初期过估计问题也是为什么早期Transformer训练必须配warmup才能稳定跑起来的原因。SGDmomentum这类优化器就更不用说了momentum会累积历史梯度的方向如果一开始方向是错的累积出一个高速冲偏的速度后面纠正起来极其费劲。预热阶段学习率小相当于每步走得小心翼翼即使一开始方向偏一点也不会形成难以刹车的惯性。所以我的建议很简单无论你用Adam还是SGD训练周期超过50个epoch的项目都建议加上warmup。它带来的成本几乎可以忽略不计却能显著降低你需要去排查训练发散问题的概率。1.3 主流预热策略横向对比预热并不是只有“线性增长”这一种做法。就我见过的项目而言常见的大概有四类预热方式学习率变化规则优点缺点适用场景Constant Warmup先用一个固定小学习率跑N步再直接跳到目标值实现极简单只要if判断切换时有突变容易造成loss波动小规模实验、快速验证Linear / Gradual Warmup学习率从初始值线性增长到目标值本文主角平滑、无突变通用性好需要额外实现或引入包装器绝大多数CV、NLP训练Cosine Warmup按余弦曲线从0升到目标值前期慢、中间快非常平滑实现稍复杂超参数增多Transformer类模型、大batch训练Exponential Warmup按指数形式逼近目标学习率上升先快后慢灵活终点处斜率不连续调试麻烦少数特定研究的实验从工程角度来说Linear/Gradual是性价比最高的方案逻辑简单、行为可预期、切换平滑。这也是GradualWarmupScheduler这个工具存在的意义它把这条线性预热逻辑封装成了一个标准的PyTorch scheduler让你不用改训练循环结构就能直接用。2. GradualWarmupScheduler的设计思路2.1 一个老问题如何“优雅地”结束预热如果你在搜索引擎里翻warmup的代码会看到大量这种写法if epoch warmup_epochs: lr base_lr * epoch / warmup_epochs else: lr cosine_annealing(epoch - warmup_epochs)这种写法确实能工作但它有几个挺别扭的地方。第一你的训练循环里多了一个if-else分支调参时每次都要改这段逻辑。第二很多人用完warmup之后忘记把后续调度器的周期减去预热步数导致整个训练步数计算错乱。第三不同项目的预热时长不同复制粘贴后容易漏改参数。GradualWarmupScheduler的核心思路是把这个逻辑封装成一个包装器wrapper。它内部持有一个真正的后续调度器after_scheduler在预热阶段输出线性增长的学习率预热结束后把控制权交还给after_scheduler。对你来说训练循环里只需要多调用一次step不需要自己写任何if判断。2.2 拆开看它的核心逻辑这个实现本身并不复杂整个类的核心就一个get_lr方法。如果把它还原成伪代码大概是这样一个过程预热阶段从last_epoch0开始每调用一次steplast_epoch加1。在这一阶段每个参数组的学习率由公式lr base_lr (base_lr * multiplier - base_lr) * last_epoch / total_epoch计算得到。当last_epoch小于等于total_epoch时学习率均匀地从base_lr增长到base_lr * multiplier。注意这里说“均匀”是因为公式是last_epoch的线性函数斜率固定。当last_epoch第一次大于total_epoch时说明预热已经结束。此时先做一件关键的事情把after_scheduler的base_lrs整体乘以multiplier然后开始调用after_scheduler的get_lr来获取后续学习率。这一步是精髓它解决了预热终点和后续调度器起点不一致的问题。如果不乘这个multiplierafter_scheduler还傻乎乎地以为base_lr是原来的值你预热的终点是base_lr * multiplier切换时就会瞬间掉回base_lr造成学习率曲线出现一个断崖。2.3 为什么用包装器而不是子类化你可能会问为什么不直接继承一个CosineAnnealingLR然后在里面加上预热逻辑这样不是更简单吗我一开始也是这么想的后来发现这样写会让代码变得非常难维护。你每换一种后续调度策略就要重写一个子类的预热带逻辑比如CosineWithWarmup、StepWithWarmup全都要各自维护一个版本。而包装器模式只需要写一个类后续策略通过参数传入组合自由度相当高。你今天想用StepLR就传StepLR明天想换CosineAnnealingLR就传CosineAnnealingLR训练脚本主体一行都不用改。还有一个好处是调试方便。你可以把包装器外面的学习率打印出来然后单独去验证after_scheduler本身的行为是否正确两者互不干扰。出现异常时排查范围一下子缩小很多。3. 代码实现与集成完整步骤3.1 环境准备与最小依赖这个实现只依赖PyTorch本身不需要安装任何额外的第三方库。我建议使用PyTorch 1.8以上的版本因为底层调度器接口在早期版本里有不少变动。我的实验环境是PyTorch 2.x Python 3.10实测下来没有任何兼容性问题。代码的核心是继承torch.optim.lr_scheduler._LRScheduler。这是一个所有PyTorch调度器共同的基类它帮你维护了last_epoch、base_lrs等基础状态。你只需要实现get_lr方法返回当前学习率列表基类会在step时自动更新optimizer里每个参数组的学习率。3.2 核心类完整代码下面给出完整代码这个版本来自社区里流传很广的一个实现我在其基础上加了详细的注释方便阅读from torch.optim.lr_scheduler import _LRScheduler class GradualWarmupScheduler(_LRScheduler): 学习率线性预热调度器。 参数: optimizer: 被包装的优化器。 multiplier: 预热结束时的学习率相对于base_lr的倍数。 常规用法设成1表示预热后回归base_lr。 total_epoch: 预热持续的步数调用step的次数。 after_scheduler: 预热结束后接管的调度器。 def __init__(self, optimizer, multiplier, total_epoch, after_schedulerNone): self.multiplier multiplier self.total_epoch total_epoch self.after_scheduler after_scheduler self.finished False self.last_epoch -1 super().__init__(optimizer) def get_lr(self): if self.last_epoch self.total_epoch: if self.after_scheduler: if not self.finished: self.after_scheduler.base_lrs [ base_lr * self.multiplier for base_lr in self.base_lrs ] self.finished True return self.after_scheduler.get_lr() return [base_lr * self.multiplier for base_lr in self.base_lrs] return [ (base_lr * self.multiplier - base_lr) / self.total_epoch * self.last_epoch base_lr for base_lr in self.base_lrs ]这段代码看着不长每个地方都有它的用意下面我逐段拆解。3.3 关键逻辑逐段拆解先看构造函数。在调用super().init(optimizer)之前self.last_epoch被手动设置成了-1。这么做的原因是PyTorch的_LRScheduler基类在初始化时会自动调用一次step使last_epoch从-1变成0。也就是说调度器对象创建完成后last_epoch其实已经是0了优化器里的学习率也已经被设置为第一次get_lr的返回值。如果你想在创建调度器后打印初始学习率你会看到它等于base_lr而不是你设定的那个很小的起始值。这个细节很容易被忽视在后面讲可视化验证时你就会明白它有多重要。再看get_lr中的分支判断。当last_epoch还小于等于total_epoch时走的是线性增长公式。举个例子假设base_lr是0.01multiplier是1total_epoch是5。那么第0次step时last_epoch0lr0.01第1次step时last_epoch1lr0.01 (0.01-0.01)*1/50.01。诶这算下来l不还是0.01吗对你发现这个细节了当multiplier1时线性公式变成了恒等于base_lr根本不会从很小的值开始增长。这是这个社区实现的常见用法或者说容易让人困惑的地方。实际上当multiplier1时上面的公式退化为常数所以更常见的做法是在外部把base_lr设置成预热结束后的目标值然后预热阶段是恒定的低学习率不对我重新再算一下刚才的公式。注意公式是这样写的lr base_lr (base_lr * multiplier - base_lr) * last_epoch / total_epoch当multiplier1时base_lr * multiplier - base_lr 0无论last_epoch是多少lr恒等于base_lr。这不符合线性预热从低到高的语义。但如果你去看社区里大量训练脚本确实是multiplier1搭配线性增长在使用。这怎么解释关键区别在于很多人在创建GradualWarmupScheduler时传入的optimizer里的初始学习率并不是预热终点而是起始低学习率。比如optimizer的初始lr设成0.001后续StepLR的base_lr是0.01然后multiplier10。这个组合下预热阶段lr从0.001线性增长到0.01正好对接multiplier10后的base_lrs。这种配置在Faster R-CNN这类检测模型里很常见。而另一种更常见的配置是multiplier1但是大家并不会真正让预热阶段恒为base_lr而是直接不管它。很多开源仓库里抄来抄去并没有仔细推敲这个公式的语义。从我个人的角度我建议你明确区分两种配置思路思路Aoptimizer初始lr预热起始值multiplier目标值/起始值预热到目标值后after_scheduler从目标值开始衰减。思路Boptimizer初始lr目标值multiplier1预热阶段lr恒为base_lr失去线性增长意义此时应手动把公式改为从更低的值开始。为了让你使用这个社区版时行为符合直觉我推荐直接把optimizer的初始学习率设成预热起点然后multiplier设为目标值/起点值。举个例子目标学习率是0.01希望从0.001开始预热10步那么optimizer torch.optim.SGD(model.parameters(), lr0.001) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max90) warmup_scheduler GradualWarmupScheduler( optimizer, multiplier10, total_epoch10, after_schedulerscheduler )这样预热阶段lr从0.001涨到0.01余弦退火从0.01开始衰减。整体曲线是连续且平滑的这也就解释了为什么multiplier参数存在它本质上是预热终点与起点的倍数。再看切换逻辑。当last_epoch第一次大于total_epoch时代码进入if分支此时做两件事把after_scheduler.base_lrs整体乘以multiplier然后返回after_scheduler.get_lr()。这里有一点需要特别注意它调用了get_lr但没有调用after_scheduler.step。也就是说从切换那一刻起after_scheduler返回的学习率会基于被放大后的base_lrs来计算但after_scheduler自身的last_epoch并不会在这一步增加。这就意味着如果你在外部同时也调用了after_scheduler.step两个调度器的状态推进节奏可能不一致。为了让状态推进保持一致标准用法是在训练循环里同时step两个调度器先step after_scheduler再step warmup_scheduler。因为warmup_scheduler的get_lr返回的是after_scheduler当前状态下的学习率所以必须先让after_scheduler的状态更新到位然后warmup_scheduler才能拿到正确的结果。3.4 在训练循环中的集成下面给一个完整的集成示例。这里我用一个简单的两层全连接网络跑MNIST分类重点展示调度器的使用方式数据加载部分省略。import torch import torch.nn as nn from torch.optim.lr_scheduler import CosineAnnealingLR class SimpleNet(nn.Module): def __init__(self): super().__init__() self.fc1 nn.Linear(28 * 28, 256) self.fc2 nn.Linear(256, 10) def forward(self, x): x x.view(x.size(0), -1) x torch.relu(self.fc1(x)) return self.fc2(x) model SimpleNet() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9) after_scheduler CosineAnnealingLR(optimizer, T_max50) warmup_scheduler GradualWarmupScheduler( optimizer, multiplier10, total_epoch10, after_schedulerafter_scheduler ) for epoch in range(60): # 这里替换成你的训练函数 train_one_epoch(model, optimizer, train_loader) # 先step后面的调度器再step预热包装器 after_scheduler.step() warmup_scheduler.step() current_lr optimizer.param_groups[0][lr] print(fepoch {epoch}, lr {current_lr:.6f})顺序为什么必须是after_scheduler.step()在前因为warmup_scheduler的get_lr在预热结束后会调用after_scheduler.get_lr()而这个get_lr返回的结果取决于after_scheduler当前的last_epoch。如果你不先step after_scheduler那么返回的lr会比预期落后一个epoch。这个顺序错位短期内看不出严重问题但会导致lr曲线和实际衰减规律对不上。另外训练60个epochwhile预热10个epoch余弦退火的T_max是50。你不妨算一下预热结束后after_scheduler的last_epoch已经被推进到10之后每个epoch继续加1从10到60共50步正好完成一个完整的余弦周期。也就是说预热步数需要从训练总步数里预先扣除而不是在训练总步数上额外增加。很多人在这个细节上犯糊涂导致训练总步数比预期多了预热的那一段。3.5 可视化验证lr曲线是否符合预期调度器代码写完怎么确认它真的工作正常直接打印lr是最笨的方式但曲线非常直观。这里分享一个小脚本它可以在不训练任何数据的情况下用随机生成的参数来模拟学习率变化import matplotlib.pyplot as plt model SimpleNet() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9) after_scheduler CosineAnnealingLR(optimizer, T_max50) warmup_scheduler GradualWarmupScheduler( optimizer, multiplier10, total_epoch10, after_schedulerafter_scheduler ) lrs [] for epoch in range(60): after_scheduler.step() warmup_scheduler.step() lrs.append(optimizer.param_groups[0][lr]) plt.plot(range(60), lrs) plt.xlabel(epoch) plt.ylabel(lr) plt.title(GradualWarmupScheduler CosineAnnealingLR) plt.show()这个脚本跑完你会看到一条非常标准的曲线前10个epoch从0.001平滑上升到0.01然后进入余弦衰减在第60个epoch附近降到接近0。如果曲线在连接点有突变说明multiplier和初始lr的搭配出了问题。把这条曲线截图保存下来每次调参时对照着看比任何日志都直观。4. 调参心得与避坑指南4.1 warmup步数怎么选预热步数没有万能公式但我可以根据自己的实验给一个大致的参考区间。对于常规数据集下的训练任务预热步数通常设置在总训练步数的5%到10%。比如你训练100个epoch预热带5到10个epoch就够。这个比例下预热带来的稳定性收益最大而对总训练进度的影响可以忽略。把batch size调大之后预热步数要适当延长。我做过一组对比batch size从256加到2048其他条件不变原来的5个epoch预热明显不够loss在前期还是出现了明显的波动延长到12个epoch之后才恢复正常。原因是batch size变大每个step的梯度方差变小但训练初期的方向偏置问题会更突出需要更多步来“校准”。如果你的总训练步数本身就少比如只有20到30个epoch那么预热比例可以控在10%以内甚至只预热3到5个epoch。关键观察点是预热结束后loss是否可以平稳过渡如果切换点附近出现明显震荡优先延长预热步数。4.2 multiplier与base_lr怎么搭配multiplier这个参数我前面已经说过它表示预热终点学习率相当于optimizer初始学习率的倍数。这里再强调一遍配置原则optimizer的初始lr请设置成预热起始的低学习率multiplier设置成目标l除以起始l。举例来说你的训练目标lr是0.01希望预热从0.001开始那就optimizer初始lr0.001multiplier10。注意一个常见误用很多人把optimizer初始lr直接设成目标lrmultiplier设成1结果预热阶段lr恒等于目标lr一点预热效果都没有。如果你不想在配置上动脑筋有一种更省心的替代方案把optimizer初始lr设成极小的值比如1e-6multiplier设成一个较大的数比如100让预热阶段从这个极小值线性爬到0.01。这样做的好处是不需要预先计算倍率但要注意浮点精度问题如果起始值太小前期学习率过低预热早期几乎不会更新权重。还有一种情况是迁移学习微调也就是加载预训练权重后只训最后的分类头。这种场景下初始lr本身就别设太大0.001到0.003就足够multiplier用1其实也无所谓因为你需要的一直是低学习率微调。预热在这里的意义主要是让最开始的几个step更平稳避免类别特征在早期被粗暴破坏。4.3 与不同优化器搭配的实测效果我带过好几个项目Box优化器都试过这里说几个典型组合的体感优化器组合预热必要性实际表现SGD momentum warmup非常必要收敛更稳精度略有提升对momentum惯性问题有明显抑制作用Adam warmup强烈建议有效抑制初期过估计尤其训练transformer类模型时几乎必需AdamW warmup强烈建议配合weight decay训练BERT类模型时是标配方案LAMB warmup必须大batch训练时如果没有预热几乎必炸这是LAMB论文里明确提过的LAMB这个例子值得多说一句。它本来就是为大batch设计的优化器在几千的batch size下第一轮迭代的梯度方差太小导致参数更新幅度被放大到离谱的程度。训练初期如果不加预热很可能一个step之后loss就变成NaN整个训练直接报废。所以在大batch分布式训练场景里warmup不是可选优化项而是生存条件。4.4 三个容易踩的坑第一个坑是检查点恢复后warmup状态丢失。训练跑了一半因为机器重启或实验中断需要resume。很多人加载完model checkpoint之后只恢复了optimizer的状态却忘了重新构建scheduler并恢复last_epoch。结果恢复训练后预热逻辑从头开始重新执行学习率从极低值重新爬升之前的优化进度全部被打乱。解决方法是保存checkpoint时同时保存warmup_scheduler.state_dict()恢复时再load进去。第二个坑是在warmup过程中频繁打印lr误以为不更新。如果你用multiplier1的配置前面算过预热阶段lr确实恒定这是配置问题而不是代码问题。如果你用multiplier10的配置但发现预热阶段lr一直没变可以检查一下是不是忘了在训练循环里调用warmup_scheduler.step()。这种低级错误排查起来最费时间所以我建议在训练脚本里每个epoch固定打印一次lr从一开始就盯着曲线看。第三个坑是after_scheduler的base_lrs在切换时被放大后如果后续重新创建调度器multiplier会被重复应用。也就是说如果你在同一个optimizer上先构建了一次warmup_scheduler跑完一轮之后没有重建optimizer又把同一个optimizer传入新的warmup_scheduler新的调度器会把已经放大过的lr再次当作base_lr来算。这个问题的表现是学习率整体变得异常大训练快速发散。规避方法很简单每个训练阶段都重新创建optimizer或者说彻底重新初始化。5. 常见问题与排查技巧实录这里整理一份问题速查表都是我在实际训练里遇到过或者帮别人排查过的典型案例。现象可能原因排查思路解决方案warmup阶段loss不降反升起始学习率太大、multiplier配置错误打印lr曲线看warmup阶段是否应线性增长把optimizer初始lr调小multiplier按目标lr/初始lr设置预热结束瞬间loss剧烈波动切换时刻lr曲线有突变打印完整的lr曲线重点观察连接点检查multiplier是否等于目标lr与预热终点的比值调整total_epoch训练后期loss震荡不收敛余弦周期T_max未扣除预热步数计算总步数和T_max是否一致设置T_max为总步数减去预热步数恢复检查点后lr曲线错乱scheduler和warmup_scheduler的state没有恢复打印恢复后的lr对比保存前的曲线保存并加载两者完整的state_dict开了warmup反而收敛变慢预热步数太长、起始lr太低观察预热后期loss下降幅度缩短预热步数到总步数的5%提高起始lr一个数量级打印的lr一直是同一个值忘了调用step或者multiplier1导致恒等检查训练循环中是否step了包装器补上step调用并重新审视multiplier配置除了这些还有一个非常隐蔽的坑。在多卡训练中使用DistributedDataParallel时PyTorch会在每个rank上独立维护optimizer和scheduler状态。如果你在rank 0上打印了lr其他rank上的lr不一定同步。加载checkpoint时最好在main rank上保存并load scheduller state再用broadcast或仅使用rank 0的state来恢复。我遇到过rank间lr不一致导致loss曲线分叉的情况折腾了很久才发现是scheduler状态没同步。我需要特别提醒一句排查lr相关问题时最快的工具就是把每个epoch的lr打印到日志或者干脆画成曲线图。不要靠猜。每次调整warmup参数后先跑10个epoch看曲线和loss走势确认无误后再继续完整训练。这个习惯能帮你省下一大半调试时间。最后再分享一个小技巧在真正的大规模训练里我习惯把warmup步数用iteration而不是epoch来计算尤其当训练集特别大且单个epoch耗时很长时。你可以在训练循环里设置一个global_step计数器把total_epoch改成total_step来用。GradualWarmupScheduler的接口没有限定单位它只看你调用了多少次step所以你完全可以用iteration步数来驱动它。我的个人体会是把warmup和训练监控绑定在全局步数上比绑在epoch上更直观也不容易在分布式环境里出现各rank进度不一的问题。