训练一个图像分类模型跑到第 200 个 batch 的时候 loss 突然从 1.0 左右蹦到 8.7再往后几个 step 直接 NaN。如果你也遇到过这种场面多半不是模型结构写错而是学习率在训练初期就顶得太满。学习率预热learning rate warmup就是专门解决这类问题的标准手段先让学习率在一个很小的值附近跑几步再逐步抬到目标值给梯度统计量和 BN 统计量一段缓冲期。GradualWarmupScheduler 是 PyTorch 训练代码里流传很广的一种自定义调度器写法专门负责“预热 主调度器交接”。这篇文章我会从原理讲到实现代码再给出完整接入方式和调参建议适合正在用 PyTorch 训练模型、又不想在第一个 epoch 就被 loss 搞心态的人。1. 训练前几百步最容易翻车学习率为什么不能一上来就顶满1.1 一个很典型的翻车现场我见过太多相似过程loss 前几十步降得挺顺突然某一步开始震荡紧接着一个尖峰冲到几十后面怎么调学习率都救不回来。很多人第一反应是调低全局学习率结果训练变慢也有人怀疑是数据或者模型结构的问题查了半天发现 optimizer 和模型都没错问题恰恰出在那个“看起来越大会收敛越快”的学习率上。这不是玄学。模型刚初始化时所有权重都是从某个随机分布里采出来的网络输出既不准确也不稳定而训练是一个高度非线性的过程初始阶段一步“大跳”就可能把参数推到梯度很小但损失很高的死角之后再想跳出来就非常困难。学习率预热就是在这一小段时间里把步子刻意放小给模型一个“先站稳再跑”的机会。1.2 背后的三个原因第一是梯度噪声大。随机初始化下前向输出的绝对值偏大反向传播容易出现极端梯度值这时用满学习率一步跨出去相当于蒙着眼冲刺。第二是 BatchNorm 的 running stats 还没跟上。前几百个 batch 里每个 batch 的均值方差统计量都偏离稳定值网络实际收到的数据分布一直在变大学习率会进一步放大这种不确定性。第三是 Adam 这类自适应优化器在最初几步里梯度的一阶矩和二阶矩估计都不可靠有效步长忽大忽小配大学习率很容易直接炸掉。我把同一个结构分别用“直接从目标学习率起步”和“先预热 5% 步数”各跑了一遍区别非常明显前者在 300 步左右出现一个 loss 尖峰虽然没到 NaN但后面收敛明显变慢后者全程曲线平滑最终精度还高了接近一个点。这说明预热不只是“防止炸掉”很多时候它直接关系到能不能找到一个更好的解。1.3 预热到底改变了什么预热做的事情本身很简单把学习率从 0 或者一个很小的值用若干步线性提高到目标值。相当于让优化器先用小步幅试水温跑顺了再放开速度。它不改变优化算法本身只是给早期训练加了一层保护。这也是为什么几乎所有大规模 Transformer 类训练都在用 warmup——模型越大、梯度越不稳定这层保护越重要。现在很多训练框架甚至把 warmup 当作默认配置而不是可选优化项。2. 三种热身方案的取舍与 GradualWarmupScheduler 的定位2.1 三种常见打法对比常被提起的热身方案主要有三种常量预热、线性预热、指数预热。它们的区别在于学习率在预热阶段按什么轨迹爬升。方案更新规则前 K 步内优点缺点常见场景常量预热lr 固定为某个小值到点直接切目标 lr实现简单没有零学习率浪费切换瞬间跳变可能引入二次尖峰手写 if 判断的简单脚本线性预热lr base × (step / K)曲线平滑易理解和余弦退火衔接自然如果从 0 起步第一个 step 不产生更新绝大多数视觉和文本训练指数预热lr base × exp(α × step)爬升前慢后快节奏激进结束点取值需要小心否则衔接突兀某些超大规模预训练配方GradualWarmupScheduler 这个名字里的 Gradual 指的就是线性拉高。相比指数预热线性在上限封顶和代码实现上都更直观所以成了社区实现里最主流的形态。2.2 为什么线性预热最常被写进工具类因为“从当前值平滑过渡到目标值”这件事线性函数最容易验证、最容易 debug。你打印前十个 lr 值一眼就能确认增长斜率对不对。指数预热虽然在一些大模型预训练配方里出现过但普通业务模型用不上那么极端的节奏。更关键的是线性预热和余弦退火几乎是天生一对。余弦退火的起点就是最高学习率而预热恰好把学习率“抬”到最高点再交给余弦去衰减。两者在交接点上都取最高值配合起来整条曲线就是一段平滑的“爬坡 下山”。如果你自己用 if-else 写预热又用 scheduler 写余弦两段逻辑很难在同一个地方统一管理一旦要改步数或者做断点续训代码会越改越乱。2.3 工具定位它不是独立调度器而是一个“交接器”必须说清楚的一点GradualWarmupScheduler 并不独自负责整条学习率曲线。它的工作是两段式的前 K 步自己接管线性提高学习率第 K 步之后把控制权交给另一个真正的调度器比如 CosineAnnealingLR由后者接管后续的下降阶段。所以你可以把它理解成一个“交接器”或者“管道”预热段是它自己写死的逻辑主调度器是它的出口。这样设计的好处是解耦——预热逻辑写一份主调度器随便换不用每个项目都重写一套 if-else。这也是它能在各种训练代码里被反复复用的根本原因。3. 手写一版可复用的 GradualWarmupScheduler逐段拆解3.1 完整实现代码from torch.optim.lr_scheduler import _LRScheduler class GradualWarmupScheduler(_LRScheduler): def __init__(self, optimizer, multiplier, total_epoch, after_schedulerNone): super().__init__(optimizer) self.multiplier multiplier self.total_epoch total_epoch self.after_scheduler after_scheduler self.finished False def get_lr(self): if self.last_epoch self.total_epoch: if self.after_scheduler is not None: 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] if self.multiplier 1.0: return [ base_lr * (self.last_epoch / self.total_epoch) for base_lr in self.base_lrs ] return [ base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch 1.0) for base_lr in self.base_lrs ] def step(self, epochNone): if self.finished and self.after_scheduler is not None: if epoch is None: self.after_scheduler.step(None) else: self.after_scheduler.step(epoch - self.total_epoch) self._last_lr self.after_scheduler.get_last_lr() else: super().step(epoch) def state_dict(self): # 不要把 after_scheduler 整个对象存进去否则 pickle 会带出一大堆东西 return { k: v for k, v in self.__dict__.items() if k not in (optimizer, after_scheduler) }这个类在社区里流传了很久具体出处已经不好考证我按自己日常用的版本整理了一下重点补了 state_dict 的过滤逻辑。核心参数只有三个multiplier、total_epoch、after_scheduler。3.2 get_lr() 里的两条计算分支预热段last_epoch total_epoch分两种情况。第一种是 multiplier1学习率从 0 线性爬到 base_lr第二种是 multiplier 大于 1学习率从 base_lr 爬到 multiplier×base_lr。两者的区别关键在于 optimizer 创建时 lr 该填什么参数组合optimizer 里的 lr 含义预热效果multiplier1目标学习率0 → 目标 lrmultiplier10base_lr0.01起步学习率0.01 → 0.1交接段last_epoch total_epoch做一件事把 after_scheduler 的 base_lrs 从原来的 base_lr 改成 base_lr×multiplier。这一步很关键——如果 CosineAnnealingLR 创建时记录的初始学习和你要的峰值不一致后面的余弦曲线起点就是错的。改完 base_lrs 之后再返回 after_scheduler 自己算出来的学习率。另外有个细节值得注意multiplier1 的写法里第一次 scheduler.step() 后 last_epoch0得到的学习率是 0也就是说第一个 batch 其实没有更新。如果你在意这一步可以把公式改成 (last_epoch1)/total_epoch或者干脆接受这个“零号热身步”。按我的经验这一步在实际训练里没有任何可感知的影响不必为它纠结。3.3 step() 的交接动作预热段直接调用 super().step(epoch)由基类把 last_epoch 加 1再触发 get_lr()。预热结束后后续每次 step() 都改为调用 after_scheduler.step()让主调度器的 last_epoch 从 0 开始计如果调用时传了 epoch还要把计数减去 total_epoch保证余弦的相位从正确位置起算。还有个容易迷惑的点交接之后wrapper 自己的 last_epoch 就不再增长了后面真正的调度进度其实由 after_scheduler 维护。如果你 debug 时去打印 wrapper 的 last_epoch看到的会是一个停在预热末端的数字这不算 bug是这类实现的正常行为。至于判断条件是 还是 不同版本差异只在于预热算 K 步还是 K1 步。不要纠结哪个更对选定一种后在计算总步数和断点恢复时保持一致就行。3.4 状态保存after_scheduler 不该被一起 pickle几乎所有沿用老版本的代码都会在断点续训时踩一次坑scheduler.state_dict() 会把整个 after_scheduler 对象带进去torch.save 时 pickle 会把里面引用的 optimizer、参数引用全链序列化文件巨大且容易出诡异错误。我在 get_lr 后面的 state_dict() 方法里做了过滤把 after_scheduler 和 optimizer 都排除掉只存纯状态字段比如 last_epoch、finished、base_lrs。这样一来断点文件干净恢复也不会污染新建的 scheduler 实例。4. 接入训练循环一份能直接改的模板4.1 实例化顺序和量纲统一先解决整个调度的“单位”问题。PyTorch 的 scheduler 并不知道你调的是 epoch 还是 batch它只在你调用 step() 的时候把 last_epoch 加 1。所以 total_epoch、T_max 这些值必须和你的调用节奏保持同一个量纲。我的建议是统一用 optimizer 更新次数来计也就是训练循环里 optimizer.step() 被调用的总次数total_train_steps epochs * len(train_loader) warmup_steps int(total_train_steps * 0.05) after_scheduler CosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ) scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerafter_scheduler )注意 T_max 不能再填 epoch 数。如果填 100 个 epoch而实际训练有 5 万个 step余弦退火会在 5000 步左右就到底后面一大段时间学习率压到 0训练白白浪费。4.2 调用位置只在 optimizer.step() 之后调一次for epoch in range(epochs): for x, y in train_loader: loss compute_loss(model(x), y) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() if global_step % 100 0: cur_lr scheduler.get_last_lr()[0] print(fstep {global_step} loss {loss.item():.4f} lr {cur_lr:.6f}) global_step 1这里最关键的一点不要在 epoch 循环末尾再补一个 scheduler.step()。预热段如果被额外 step 一次学习率会烧得比预期快一倍交接后主调度器也会跟着多退一步。我见过很多代码把 step() 同时放在 batch 循环内和 epoch 末尾最后曲线的衰减速度和纸上画的完全对不上。4.3 PyTorch 内置的替代方案如果你用的 PyTorch 版本比较新1.13 之后都有 SequentialLR其实有一个完全不用自定义类的方案from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR warmup LinearLR( optimizer, start_factor0.01, # 从 1% 起步 end_factor1.0, # 预热结束到 100% total_iterswarmup_steps ) cosine CosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ) scheduler SequentialLR( optimizer, schedulers[warmup, cosine], milestones[warmup_steps] )内置方案的好处是状态管理更省心state_dict 不会带出多余的 scheduler 对象。缺点是不如自定义类灵活比如你想在预热段用一个自定义的曲线或者想精确控制交接瞬间的 base_lrs还是要回到自定义实现。我的建议是先把自定义类的工作原理弄明白日常用内置方案完全没问题两者并不矛盾。4.4 跑训练前先干跑一遍学习率曲线大规模训练前我强烈建议先干跑一遍调度器确认曲线形态没问题lrs [] preview_scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerCosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ), ) for _ in range(total_train_steps): preview_scheduler.step() lrs.append(preview_scheduler.get_last_lr()[0]) print(lrs[:5], lrs[-5:])如果打印出来的前几个值是 [0.0, 0.001, 0.002, ...]说明预热正常如果交接附近的学习率突然跳变说明 base_lrs 没有接对趁训练没开始赶紧修。这一步花不了几秒钟但能省下后面一整天的排查时间。5. 实测中容易踩的五个坑5.1 坑一把总步数当成了总 epoch 数最常见的错误是直接写 total_epoch5然后训练循环里每个 batch 都调一次 step()结果预热只持续了 5 个 batch 就结束后面一大段训练完全没在预热状态。正确做法是先把总更新次数算出来再取百分比。epoch 只是外层循环的单位不要拿它当调度器的计数单位。5.2 坑二一个 step 内调用了两次 scheduler.step()我在 4.2 里提过batch 内一次、epoch 末尾一次这是双倍速烧学习率。尤其预热段双倍速会让整条曲线在时间轴上压缩一半实际效果和你预期的完全不是一回事。排查方法很简单打印每个 epoch 结束时的 get_last_lr()如果比预期大很多基本就是这里出了问题。5.3 坑三梯度累积场景下调度节奏错位使用梯度累积时通常不会每个 batch 都调 optimizer.step()而是攒够 N 个 batch 才更新一次权重。如果你的 scheduler.step() 仍然放在 batch 循环里学习率在单位权重更新次数上的推进速度就会快 N 倍。正确做法是让 scheduler.step() 紧跟 optimizer.step()或者用 global_step 做统一计数保证“一个权重更新对应一次调度步”。5.4 坑四断点续训时 scheduler 状态恢复错乱如果沿用老版本写法直接 save 整个 scheduler.state_dict()恢复时报错或者文件巨大多半是 after_scheduler 被一起 pickle 了。用我前面给出的自定义 state_dict 之后恢复逻辑就干净多了checkpoint torch.load(ckpt_path) model.load_state_dict(checkpoint[model]) optimizer.load_state_dict(checkpoint[optimizer]) after_scheduler CosineAnnealingLR(optimizer, T_maxtotal_train_steps - warmup_steps) scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerafter_scheduler, ) scheduler.load_state_dict(checkpoint[scheduler]) after_scheduler.load_state_dict(checkpoint[cosine_scheduler])注意恢复时 warmup_steps 和 T_max 必须和保存时完全一致否则 last_epoch 对不上学习率会从一个错误的位置继续。5.5 坑五让 ReduceLROnPlateau 当了 after_schedulerReduceLROnPlateau 的 step() 需要接收一个指标值比如验证集 loss而 GradualWarmupScheduler 的交接逻辑会调用 after_scheduler.step(None)直接触发参数错误。所以不要把“按指标触发”的调度器塞进 after_scheduler。如果确实需要在预热后走 plateau 策略就把预热单独写成 if 分支热身结束后手动切到 plateau 逻辑。6. 预热步数到底设多少我的调参经验6.1 按总迭代比例估算一个比较通用的起点是总更新步数的 1%~10%。训练步数越少比例越偏下限训练步数越多可以适当偏上限。比如总共 5 万步的训练5% 就是 2500 步大约相当于 5 个 epoch绝大多数情况下够用如果训练总共才 2000 步5% 只有 100 步也算合理。比例太低起不到稳定作用比例太高则前期学得太慢白白消耗计算资源。6.2 分场景推荐区间场景推荐预热步数说明小型分类模型从头训练总步数的 5%~10%网络浅、梯度相对温和给足缓冲即可大中型视觉模型总步数的 3%~5%常见配余弦退火使用Transformer 类 / 大规模预训练总步数的 1%~3%或固定几千步模型层数深、梯度不稳定但总训练步数巨大迁移学习 / 微调预训练权重数百步或 1~2 个 epoch权重已有先验不需要太长热身6.3 怎么判断设置是否合理训练时盯三个信号。第一预热结束前 loss 不应该出现大幅反弹如果有说明预热步数偏短或者学习率峰值偏高。第二交接点前后 loss 曲线不能有突然的斜率转折出现转折多半是主调度器的 base_lrs 没接对。第三稳定训练一段时间后如果 loss 下降明显比无预热基线慢说明预热比例太高可以砍一半再试。我自己默认“5% 总步数 余弦退火 multiplier1”的配置先把损失曲线压稳再根据第一轮结果微调很少在这个参数上翻车。用 Adam 类优化器时预热步数可以比 SGD 场景略短一些因为自适应矩估计本身就有缓冲效果但大 batch 训练反而要把预热拉长因为大 batch 下有效学习率更高一小段预热不够让 BN 统计量稳定下来。先跑一版把前 500 个 step 的 lr 和 loss 打印出来看看比任何经验值都靠谱。 
长河君 · 2026/10/23 9:44:37 · 阅读 1.2k 学习率预热与GradualWarmupScheduler：PyTorch训练防loss尖峰
训练一个图像分类模型跑到第 200 个 batch 的时候 loss 突然从 1.0 左右蹦到 8.7再往后几个 step 直接 NaN。如果你也遇到过这种场面多半不是模型结构写错而是学习率在训练初期就顶得太满。学习率预热learning rate warmup就是专门解决这类问题的标准手段先让学习率在一个很小的值附近跑几步再逐步抬到目标值给梯度统计量和 BN 统计量一段缓冲期。GradualWarmupScheduler 是 PyTorch 训练代码里流传很广的一种自定义调度器写法专门负责“预热 主调度器交接”。这篇文章我会从原理讲到实现代码再给出完整接入方式和调参建议适合正在用 PyTorch 训练模型、又不想在第一个 epoch 就被 loss 搞心态的人。1. 训练前几百步最容易翻车学习率为什么不能一上来就顶满1.1 一个很典型的翻车现场我见过太多相似过程loss 前几十步降得挺顺突然某一步开始震荡紧接着一个尖峰冲到几十后面怎么调学习率都救不回来。很多人第一反应是调低全局学习率结果训练变慢也有人怀疑是数据或者模型结构的问题查了半天发现 optimizer 和模型都没错问题恰恰出在那个“看起来越大会收敛越快”的学习率上。这不是玄学。模型刚初始化时所有权重都是从某个随机分布里采出来的网络输出既不准确也不稳定而训练是一个高度非线性的过程初始阶段一步“大跳”就可能把参数推到梯度很小但损失很高的死角之后再想跳出来就非常困难。学习率预热就是在这一小段时间里把步子刻意放小给模型一个“先站稳再跑”的机会。1.2 背后的三个原因第一是梯度噪声大。随机初始化下前向输出的绝对值偏大反向传播容易出现极端梯度值这时用满学习率一步跨出去相当于蒙着眼冲刺。第二是 BatchNorm 的 running stats 还没跟上。前几百个 batch 里每个 batch 的均值方差统计量都偏离稳定值网络实际收到的数据分布一直在变大学习率会进一步放大这种不确定性。第三是 Adam 这类自适应优化器在最初几步里梯度的一阶矩和二阶矩估计都不可靠有效步长忽大忽小配大学习率很容易直接炸掉。我把同一个结构分别用“直接从目标学习率起步”和“先预热 5% 步数”各跑了一遍区别非常明显前者在 300 步左右出现一个 loss 尖峰虽然没到 NaN但后面收敛明显变慢后者全程曲线平滑最终精度还高了接近一个点。这说明预热不只是“防止炸掉”很多时候它直接关系到能不能找到一个更好的解。1.3 预热到底改变了什么预热做的事情本身很简单把学习率从 0 或者一个很小的值用若干步线性提高到目标值。相当于让优化器先用小步幅试水温跑顺了再放开速度。它不改变优化算法本身只是给早期训练加了一层保护。这也是为什么几乎所有大规模 Transformer 类训练都在用 warmup——模型越大、梯度越不稳定这层保护越重要。现在很多训练框架甚至把 warmup 当作默认配置而不是可选优化项。2. 三种热身方案的取舍与 GradualWarmupScheduler 的定位2.1 三种常见打法对比常被提起的热身方案主要有三种常量预热、线性预热、指数预热。它们的区别在于学习率在预热阶段按什么轨迹爬升。方案更新规则前 K 步内优点缺点常见场景常量预热lr 固定为某个小值到点直接切目标 lr实现简单没有零学习率浪费切换瞬间跳变可能引入二次尖峰手写 if 判断的简单脚本线性预热lr base × (step / K)曲线平滑易理解和余弦退火衔接自然如果从 0 起步第一个 step 不产生更新绝大多数视觉和文本训练指数预热lr base × exp(α × step)爬升前慢后快节奏激进结束点取值需要小心否则衔接突兀某些超大规模预训练配方GradualWarmupScheduler 这个名字里的 Gradual 指的就是线性拉高。相比指数预热线性在上限封顶和代码实现上都更直观所以成了社区实现里最主流的形态。2.2 为什么线性预热最常被写进工具类因为“从当前值平滑过渡到目标值”这件事线性函数最容易验证、最容易 debug。你打印前十个 lr 值一眼就能确认增长斜率对不对。指数预热虽然在一些大模型预训练配方里出现过但普通业务模型用不上那么极端的节奏。更关键的是线性预热和余弦退火几乎是天生一对。余弦退火的起点就是最高学习率而预热恰好把学习率“抬”到最高点再交给余弦去衰减。两者在交接点上都取最高值配合起来整条曲线就是一段平滑的“爬坡 下山”。如果你自己用 if-else 写预热又用 scheduler 写余弦两段逻辑很难在同一个地方统一管理一旦要改步数或者做断点续训代码会越改越乱。2.3 工具定位它不是独立调度器而是一个“交接器”必须说清楚的一点GradualWarmupScheduler 并不独自负责整条学习率曲线。它的工作是两段式的前 K 步自己接管线性提高学习率第 K 步之后把控制权交给另一个真正的调度器比如 CosineAnnealingLR由后者接管后续的下降阶段。所以你可以把它理解成一个“交接器”或者“管道”预热段是它自己写死的逻辑主调度器是它的出口。这样设计的好处是解耦——预热逻辑写一份主调度器随便换不用每个项目都重写一套 if-else。这也是它能在各种训练代码里被反复复用的根本原因。3. 手写一版可复用的 GradualWarmupScheduler逐段拆解3.1 完整实现代码from torch.optim.lr_scheduler import _LRScheduler class GradualWarmupScheduler(_LRScheduler): def __init__(self, optimizer, multiplier, total_epoch, after_schedulerNone): super().__init__(optimizer) self.multiplier multiplier self.total_epoch total_epoch self.after_scheduler after_scheduler self.finished False def get_lr(self): if self.last_epoch self.total_epoch: if self.after_scheduler is not None: 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] if self.multiplier 1.0: return [ base_lr * (self.last_epoch / self.total_epoch) for base_lr in self.base_lrs ] return [ base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch 1.0) for base_lr in self.base_lrs ] def step(self, epochNone): if self.finished and self.after_scheduler is not None: if epoch is None: self.after_scheduler.step(None) else: self.after_scheduler.step(epoch - self.total_epoch) self._last_lr self.after_scheduler.get_last_lr() else: super().step(epoch) def state_dict(self): # 不要把 after_scheduler 整个对象存进去否则 pickle 会带出一大堆东西 return { k: v for k, v in self.__dict__.items() if k not in (optimizer, after_scheduler) }这个类在社区里流传了很久具体出处已经不好考证我按自己日常用的版本整理了一下重点补了 state_dict 的过滤逻辑。核心参数只有三个multiplier、total_epoch、after_scheduler。3.2 get_lr() 里的两条计算分支预热段last_epoch total_epoch分两种情况。第一种是 multiplier1学习率从 0 线性爬到 base_lr第二种是 multiplier 大于 1学习率从 base_lr 爬到 multiplier×base_lr。两者的区别关键在于 optimizer 创建时 lr 该填什么参数组合optimizer 里的 lr 含义预热效果multiplier1目标学习率0 → 目标 lrmultiplier10base_lr0.01起步学习率0.01 → 0.1交接段last_epoch total_epoch做一件事把 after_scheduler 的 base_lrs 从原来的 base_lr 改成 base_lr×multiplier。这一步很关键——如果 CosineAnnealingLR 创建时记录的初始学习和你要的峰值不一致后面的余弦曲线起点就是错的。改完 base_lrs 之后再返回 after_scheduler 自己算出来的学习率。另外有个细节值得注意multiplier1 的写法里第一次 scheduler.step() 后 last_epoch0得到的学习率是 0也就是说第一个 batch 其实没有更新。如果你在意这一步可以把公式改成 (last_epoch1)/total_epoch或者干脆接受这个“零号热身步”。按我的经验这一步在实际训练里没有任何可感知的影响不必为它纠结。3.3 step() 的交接动作预热段直接调用 super().step(epoch)由基类把 last_epoch 加 1再触发 get_lr()。预热结束后后续每次 step() 都改为调用 after_scheduler.step()让主调度器的 last_epoch 从 0 开始计如果调用时传了 epoch还要把计数减去 total_epoch保证余弦的相位从正确位置起算。还有个容易迷惑的点交接之后wrapper 自己的 last_epoch 就不再增长了后面真正的调度进度其实由 after_scheduler 维护。如果你 debug 时去打印 wrapper 的 last_epoch看到的会是一个停在预热末端的数字这不算 bug是这类实现的正常行为。至于判断条件是 还是 不同版本差异只在于预热算 K 步还是 K1 步。不要纠结哪个更对选定一种后在计算总步数和断点恢复时保持一致就行。3.4 状态保存after_scheduler 不该被一起 pickle几乎所有沿用老版本的代码都会在断点续训时踩一次坑scheduler.state_dict() 会把整个 after_scheduler 对象带进去torch.save 时 pickle 会把里面引用的 optimizer、参数引用全链序列化文件巨大且容易出诡异错误。我在 get_lr 后面的 state_dict() 方法里做了过滤把 after_scheduler 和 optimizer 都排除掉只存纯状态字段比如 last_epoch、finished、base_lrs。这样一来断点文件干净恢复也不会污染新建的 scheduler 实例。4. 接入训练循环一份能直接改的模板4.1 实例化顺序和量纲统一先解决整个调度的“单位”问题。PyTorch 的 scheduler 并不知道你调的是 epoch 还是 batch它只在你调用 step() 的时候把 last_epoch 加 1。所以 total_epoch、T_max 这些值必须和你的调用节奏保持同一个量纲。我的建议是统一用 optimizer 更新次数来计也就是训练循环里 optimizer.step() 被调用的总次数total_train_steps epochs * len(train_loader) warmup_steps int(total_train_steps * 0.05) after_scheduler CosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ) scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerafter_scheduler )注意 T_max 不能再填 epoch 数。如果填 100 个 epoch而实际训练有 5 万个 step余弦退火会在 5000 步左右就到底后面一大段时间学习率压到 0训练白白浪费。4.2 调用位置只在 optimizer.step() 之后调一次for epoch in range(epochs): for x, y in train_loader: loss compute_loss(model(x), y) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() if global_step % 100 0: cur_lr scheduler.get_last_lr()[0] print(fstep {global_step} loss {loss.item():.4f} lr {cur_lr:.6f}) global_step 1这里最关键的一点不要在 epoch 循环末尾再补一个 scheduler.step()。预热段如果被额外 step 一次学习率会烧得比预期快一倍交接后主调度器也会跟着多退一步。我见过很多代码把 step() 同时放在 batch 循环内和 epoch 末尾最后曲线的衰减速度和纸上画的完全对不上。4.3 PyTorch 内置的替代方案如果你用的 PyTorch 版本比较新1.13 之后都有 SequentialLR其实有一个完全不用自定义类的方案from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR warmup LinearLR( optimizer, start_factor0.01, # 从 1% 起步 end_factor1.0, # 预热结束到 100% total_iterswarmup_steps ) cosine CosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ) scheduler SequentialLR( optimizer, schedulers[warmup, cosine], milestones[warmup_steps] )内置方案的好处是状态管理更省心state_dict 不会带出多余的 scheduler 对象。缺点是不如自定义类灵活比如你想在预热段用一个自定义的曲线或者想精确控制交接瞬间的 base_lrs还是要回到自定义实现。我的建议是先把自定义类的工作原理弄明白日常用内置方案完全没问题两者并不矛盾。4.4 跑训练前先干跑一遍学习率曲线大规模训练前我强烈建议先干跑一遍调度器确认曲线形态没问题lrs [] preview_scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerCosineAnnealingLR( optimizer, T_maxtotal_train_steps - warmup_steps ), ) for _ in range(total_train_steps): preview_scheduler.step() lrs.append(preview_scheduler.get_last_lr()[0]) print(lrs[:5], lrs[-5:])如果打印出来的前几个值是 [0.0, 0.001, 0.002, ...]说明预热正常如果交接附近的学习率突然跳变说明 base_lrs 没有接对趁训练没开始赶紧修。这一步花不了几秒钟但能省下后面一整天的排查时间。5. 实测中容易踩的五个坑5.1 坑一把总步数当成了总 epoch 数最常见的错误是直接写 total_epoch5然后训练循环里每个 batch 都调一次 step()结果预热只持续了 5 个 batch 就结束后面一大段训练完全没在预热状态。正确做法是先把总更新次数算出来再取百分比。epoch 只是外层循环的单位不要拿它当调度器的计数单位。5.2 坑二一个 step 内调用了两次 scheduler.step()我在 4.2 里提过batch 内一次、epoch 末尾一次这是双倍速烧学习率。尤其预热段双倍速会让整条曲线在时间轴上压缩一半实际效果和你预期的完全不是一回事。排查方法很简单打印每个 epoch 结束时的 get_last_lr()如果比预期大很多基本就是这里出了问题。5.3 坑三梯度累积场景下调度节奏错位使用梯度累积时通常不会每个 batch 都调 optimizer.step()而是攒够 N 个 batch 才更新一次权重。如果你的 scheduler.step() 仍然放在 batch 循环里学习率在单位权重更新次数上的推进速度就会快 N 倍。正确做法是让 scheduler.step() 紧跟 optimizer.step()或者用 global_step 做统一计数保证“一个权重更新对应一次调度步”。5.4 坑四断点续训时 scheduler 状态恢复错乱如果沿用老版本写法直接 save 整个 scheduler.state_dict()恢复时报错或者文件巨大多半是 after_scheduler 被一起 pickle 了。用我前面给出的自定义 state_dict 之后恢复逻辑就干净多了checkpoint torch.load(ckpt_path) model.load_state_dict(checkpoint[model]) optimizer.load_state_dict(checkpoint[optimizer]) after_scheduler CosineAnnealingLR(optimizer, T_maxtotal_train_steps - warmup_steps) scheduler GradualWarmupScheduler( optimizer, multiplier1, total_epochwarmup_steps, after_schedulerafter_scheduler, ) scheduler.load_state_dict(checkpoint[scheduler]) after_scheduler.load_state_dict(checkpoint[cosine_scheduler])注意恢复时 warmup_steps 和 T_max 必须和保存时完全一致否则 last_epoch 对不上学习率会从一个错误的位置继续。5.5 坑五让 ReduceLROnPlateau 当了 after_schedulerReduceLROnPlateau 的 step() 需要接收一个指标值比如验证集 loss而 GradualWarmupScheduler 的交接逻辑会调用 after_scheduler.step(None)直接触发参数错误。所以不要把“按指标触发”的调度器塞进 after_scheduler。如果确实需要在预热后走 plateau 策略就把预热单独写成 if 分支热身结束后手动切到 plateau 逻辑。6. 预热步数到底设多少我的调参经验6.1 按总迭代比例估算一个比较通用的起点是总更新步数的 1%~10%。训练步数越少比例越偏下限训练步数越多可以适当偏上限。比如总共 5 万步的训练5% 就是 2500 步大约相当于 5 个 epoch绝大多数情况下够用如果训练总共才 2000 步5% 只有 100 步也算合理。比例太低起不到稳定作用比例太高则前期学得太慢白白消耗计算资源。6.2 分场景推荐区间场景推荐预热步数说明小型分类模型从头训练总步数的 5%~10%网络浅、梯度相对温和给足缓冲即可大中型视觉模型总步数的 3%~5%常见配余弦退火使用Transformer 类 / 大规模预训练总步数的 1%~3%或固定几千步模型层数深、梯度不稳定但总训练步数巨大迁移学习 / 微调预训练权重数百步或 1~2 个 epoch权重已有先验不需要太长热身6.3 怎么判断设置是否合理训练时盯三个信号。第一预热结束前 loss 不应该出现大幅反弹如果有说明预热步数偏短或者学习率峰值偏高。第二交接点前后 loss 曲线不能有突然的斜率转折出现转折多半是主调度器的 base_lrs 没接对。第三稳定训练一段时间后如果 loss 下降明显比无预热基线慢说明预热比例太高可以砍一半再试。我自己默认“5% 总步数 余弦退火 multiplier1”的配置先把损失曲线压稳再根据第一轮结果微调很少在这个参数上翻车。用 Adam 类优化器时预热步数可以比 SGD 场景略短一些因为自适应矩估计本身就有缓冲效果但大 batch 训练反而要把预热拉长因为大 batch 下有效学习率更高一小段预热不够让 BN 统计量稳定下来。先跑一版把前 500 个 step 的 lr 和 loss 打印出来看看比任何经验值都靠谱。