训练一个图像分类模型，跑到第 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 三种常见打法对比

常被提起的热身方案主要有三种：常量预热、线性预热、指数预热。它们的区别在于学习率在预热阶段按什么轨迹爬升。

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_scheduler=None): 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, epoch=None): 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）分两种情况。第一种是 multiplier=1，学习率从 0 线性爬到 base_lr；第二种是 multiplier 大于 1，学习率从 base_lr 爬到 multiplier×base_lr。两者的区别关键在于 optimizer 创建时 lr 该填什么：

交接段（last_epoch >= total_epoch）做一件事：把 after_scheduler 的 base_lrs 从原来的 base_lr 改成 base_lr×multiplier。这一步很关键——如果 CosineAnnealingLR 创建时记录的初始学习和你要的峰值不一致，后面的余弦曲线起点就是错的。改完 base_lrs 之后再返回 after_scheduler 自己算出来的学习率。

另外有个细节值得注意：multiplier=1 的写法里，第一次 scheduler.step() 后 last_epoch=0，得到的学习率是 0，也就是说第一个 batch 其实没有更新。如果你在意这一步，可以把公式改成 (last_epoch+1)/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 步还是 K+1 步。不要纠结哪个更对，选定一种后在计算总步数和断点恢复时保持一致就行。

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_max=total_train_steps - warmup_steps ) scheduler = GradualWarmupScheduler( optimizer, multiplier=1, total_epoch=warmup_steps, after_scheduler=after_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(f"step {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_factor=0.01, # 从 1% 起步 end_factor=1.0, # 预热结束到 100% total_iters=warmup_steps ) cosine = CosineAnnealingLR( optimizer, T_max=total_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, multiplier=1, total_epoch=warmup_steps, after_scheduler=CosineAnnealingLR( optimizer, T_max=total_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_epoch=5，然后训练循环里每个 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_max=total_train_steps - warmup_steps) scheduler = GradualWarmupScheduler( optimizer, multiplier=1, total_epoch=warmup_steps, after_scheduler=after_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_scheduler

ReduceLROnPlateau 的 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 分场景推荐区间

6.3 怎么判断设置是否合理

训练时盯三个信号。第一，预热结束前 loss 不应该出现大幅反弹，如果有，说明预热步数偏短或者学习率峰值偏高。第二，交接点前后 loss 曲线不能有突然的斜率转折，出现转折多半是主调度器的 base_lrs 没接对。第三，稳定训练一段时间后，如果 loss 下降明显比无预热基线慢，说明预热比例太高，可以砍一半再试。

我自己默认“5% 总步数 + 余弦退火 + multiplier=1”的配置，先把损失曲线压稳，再根据第一轮结果微调，很少在这个参数上翻车。用 Adam 类优化器时，预热步数可以比 SGD 场景略短一些，因为自适应矩估计本身就有缓冲效果；但大 batch 训练反而要把预热拉长，因为大 batch 下有效学习率更高，一小段预热不够让 BN 统计量稳定下来。先跑一版，把前 500 个 step 的 lr 和 loss 打印出来看看，比任何经验值都靠谱。