学习率调度器

本页包含 timm 中提供的学习率调度器的 API 参考文档，并概述了它们的创建和步进方式。

用法

调度器通过名称使用 create_scheduler_v2 创建，该函数返回调度器以及要训练的总 epoch 数（包括任何预热前缀、周期和冷却期，详见下文）。与 torch.optim.lr_scheduler 不同，timm 调度器不会自行计算步数。它们被设计为无状态的：学习率是调度配置和显式 epoch 或更新次数的函数，因此恢复训练运行只需要 epoch，且检查点中不存储任何调度器状态。训练循环会显式传入 epoch 或更新次数：

在每个 epoch 末尾调用 scheduler.step(epoch)，传入下一个 epoch 的索引（epoch + 1）
在每次优化器更新之后，调用 scheduler.step_update(num_updates)，并传入迄今为止的更新次数

调度器按 epoch 或按更新步长（step_on_epochs）进行步进，另一个调用为空操作，因此训练循环可以同时调用两者以支持这两种方式。

import timm
from timm.optim import create_optimizer_v2
from timm.scheduler import create_scheduler_v2

model = timm.create_model('resnet50')
optimizer = create_optimizer_v2(model, opt='adamw', lr=1e-3, weight_decay=0.05)
updates_per_epoch = len(loader)  # optimizer updates per epoch

# Cosine decay w/ 5 epochs of linear warmup, stepped per update for a smooth schedule.
scheduler, num_epochs = create_scheduler_v2(
    optimizer,
    sched='cosine',
    num_epochs=100,
    warmup_epochs=5,
    min_lr=1e-6,
    step_on_epochs=False,
    updates_per_epoch=updates_per_epoch,
)

num_updates = 0
for epoch in range(num_epochs):
    for batch in loader:
        ...  # forward, backward, optimizer.step()
        num_updates += 1
        scheduler.step_update(num_updates=num_updates)
    scheduler.step(epoch + 1)  # pass the eval metric too for 'plateau', scheduler.step(epoch + 1, metric)

持续时间（num_epochs、warmup_epochs、decay_epochs、……）始终以 epoch 为单位给出，当 step_on_epochs=False 时，工厂会将它们转换为更新次数。也可以直接从其类构造调度器，此时持续时间的单位即为其步进单位（t_in_epochs）。

训练脚本（train.py）通过以下标志暴露工厂参数：

--sched 用于选择调度策略。
--epochs 用于设置训练轮数。
--sched-on-updates 按更新次数而非按 epoch 设置步数。
--warmup-epochs 用于设置预热时长。
--warmup-lr 用于设置初始预热学习率。
--warmup-prefix 会在调度开始前添加预热，而不是与其起始部分重叠。
--min-lr 用于设置最小学习率。
--decay-epochs 用于设置 WSD / rsqrt 的步数衰减间隔或最终衰减时长。
--decay-milestones 用于设置多步学习率调度中的衰减轮次。
--decay-rate 用于设置衰减因子，即多项式调度的幂次。
--cooldown-epochs 在调度结束后，以最低学习率增加训练轮数。
--patience-epochs 用于设置 plateau 调度器的 patience。
--lr-noise 用于选择学习率噪声所在的调度部分。
--lr-noise-pct 用于限制噪声幅度。
--lr-noise-std 用于设置噪声的标准差。
--lr-cycle-mul 用于设置周期长度乘数。
--lr-cycle-decay 用于设置周期之间的峰值学习率乘数。
--lr-cycle-limit 用于设置循环次数。
--lr-k-decay 用于设置余弦/多项式调度中的 k-decay。
--lr-decay-fn 用于选择 WSD / rsqrt 的最终衰减函数。
--lr-stable-timescale 用于设置 rsqrt 稳定阶段的 timescale。

计划

`sched`	类	日程
`cosine`	`CosineLRScheduler`	从基础学习率到 `min_lr` 的余弦退火（SGDR），可选的重启（周期）和 k 衰减
`tanh`	`TanhLRScheduler`	双曲正切衰减（论文），可选的重启
`poly`	`PolyLRScheduler`	衰减至 `min_lr` 的多项式衰减，幂次由 `decay_rate` 指定，支持可选的重启和 k-衰减
`step`	`StepLRScheduler`	每 `decay_epochs` 将学习率乘以 `decay_rate`
`multistep`	`MultiStepLRScheduler`	在 `decay_milestones` 个 epoch 中的每一个，学习率乘以 `decay_rate`
`plateau`	`PlateauLRScheduler`	已弃用。当评估指标在 `patience_epochs` 个 epoch 内没有改善时，将学习率乘以 `decay_rate`，仅按 epoch 步进
`wsd`	`WSDLRScheduler`	Warmup-stable-decay（MiniCPM），先保持恒定学习率，然后在最后 `decay_epochs` 个 epoch 内衰减至 `min_lr`，默认采用 `1-sqrt` 衰减（论文）
`rsqrt`	`WSDLRScheduler`	在最后的 `decay_epochs` 期间采用线性冷却的倒数平方根衰减（Big Vision）

常用选项

Warmup. warmup_epochs 线性预热，从 warmup_lr 到基础学习率，所有调度策略均支持。
- warmup_prefix：默认情况下，预热阶段与调度的起始部分重叠（调度从 epoch 0 开始计算，预热会替换其最初的几个 epoch），因此 num_epochs 表示总长度。当设置 warmup_prefix=True 时，调度在预热结束后才开始，预热阶段会被额外计入总长度。
- 对于 wsd 和 rsqrt，稳定阶段始终在预热之后开始。如果没有 warmup_prefix，稳定阶段和衰减阶段会被压缩到 num_epochs - warmup_epochs 中，因此 num_epochs 仍然是总轮数。
min_lr：衰减调度（cosine、tanh、poly、wsd、rsqrt）在结束时达到的学习率，以及 plateau 的下限。

plateau 调度已被弃用，并将在未来的版本中移除。它从未契合上述无状态设计：其学习率取决于评估指标的历史记录，而非 epoch，且该历史记录（最佳指标、patience 和 cooldown 计数器）并未保存在检查点中，因此恢复的运行会重新开始 plateau 逻辑。不会添加对检查点的支持。验证 plateau 也是关于何时退火的一个较弱信号；像 cosine 或 wsd 这样的固定调度是可复现的，并能精确恢复，而 WSD 扫描（参见超参数扫描）可以从一次稳定的运行中比较多个预算下的 cooldown。

cooldown_epochs：在调度器结束（cosine、tanh、poly、wsd、rsqrt）后，以 min_lr 保持训练的 epoch 数，该数值会加到返回的 num_epochs 中。由于调度器在其最后一步之后的下一步才达到 min_lr，因此冷却 epoch 也是一种在按 epoch 步进时以 min_lr 进行训练的方式。
循环（cosine、tanh、poly）：当 cycle_limit > 1 时，调度器会重新开始，每个循环的时长是前一个的 cycle_mul 倍，且其峰值学习率相对于前一个循环乘以 cycle_decay。返回的 num_epochs 涵盖所有循环。
k_decay (cosine、poly): k-decay 会改变衰减的形状 (t ** k)，k > 1 会使学习率保持在较高水平更长时间。
噪声：noise 在调度器的一段范围内添加随机的乘法学习率噪声，该范围以 num_epochs 的比例给出（一个起始值，或一个起始值和一个结束值），并受 noise_pct 限制。

WSD 和 rsqrt

WSDLRScheduler 包含三个阶段：线性预热、稳定阶段以及衰减至 min_lr。

decay_epochs 设置最终的衰减时长。取值在 (0, 1) 之间时，表示稳定期 + 衰减期总长度的比例，默认值为 20%。
decay_fn（--lr-decay-fn）可选择 1-sqrt（wsd 的默认值）、linear（rsqrt 的默认值）、cosine、power 或 exp。
stable_timescale（--lr-stable-timescale，以 epoch 为单位）控制 rsqrt 稳定阶段，该阶段从 warmup 结束时开始，按 (1 + t / timescale) ** -0.5 衰减。其默认值为 warmup 长度，在没有 warmup 时必须手动设置。

该类还支持广义的power稳定阶段和自定义衰减函数。

最终的退火阶段是 WSD 调度获得最大收益的地方，因此请使用按更新次数步进（--sched-on-updates），或者在衰减阶段仅持续几个 epoch 时添加一个冷却 epoch。如果按 epoch 步进，短衰减阶段的最后一个 epoch 会明显高于 min_lr（对于 5 个衰减 epoch 中的最后一个，使用 1-sqrt 时约为基础学习率的 10%）。

# 300 epochs, 5 epoch warmup, constant lr, then a 1-sqrt decay over the final 20%
scheduler, num_epochs = create_scheduler_v2(
    optimizer, sched='wsd', num_epochs=300, warmup_epochs=5, decay_epochs=0.2,
    step_on_epochs=False, updates_per_epoch=updates_per_epoch,
)

参数组学习率缩放

带有 lr_scale 字段的参数组（例如来自 param_groups_layer_decay 的逐层学习率衰减组）会将调度后的学习率乘以其缩放系数。创建调度器时，每个参数组的基础学习率会作为 initial_lr 记录在该组中。

恢复

随着步数的传入，恢复运行的调度器会按照原始运行的方式创建，并将其步进到恢复点，即 scheduler.step(start_epoch) 或 scheduler.step_update(start_epoch * updates_per_epoch)。state_dict() / load_state_dict() 用于保存和恢复调度器的属性。

API

工厂函数

timm.scheduler.create_scheduler

< source >

( argsoptimizer: Optimizerupdates_per_epoch: int = 0 )

timm.scheduler.create_scheduler_v2

< source >

( optimizer: Optimizersched: str = 'cosine'num_epochs: int = 300decay_epochs: typing.Optional[float] = Nonedecay_milestones: typing.List[int] = (90, 180, 270)cooldown_epochs: int = 0patience_epochs: int = 10decay_rate: float = 0.1min_lr: float = 0.0warmup_lr: float = 1e-05warmup_epochs: int = 0warmup_prefix: bool = Falsenoise: typing.Union[float, typing.List[float]] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42cycle_mul: float = 1.0cycle_decay: float = 0.1cycle_limit: int = 1k_decay: float = 1.0decay_fn: typing.Optional[str] = Nonestable_timescale: typing.Optional[float] = Noneplateau_mode: str = 'max'step_on_epochs: bool = Trueupdates_per_epoch: int = 0 )

参数

optimizer — 要调度的优化器
sched — 调度器名称，可选值包括 ‘cosine’、‘tanh’、‘step’、‘multistep’、‘poly’、‘wsd’、‘rsqrt’、‘constant’、‘plateau’（已弃用）或 ‘none’（表示不使用调度器，返回 None，优化器的学习率保持不变）。 ‘constant’ 在预热阶段结束后保持学习率不变，并应用参数组的学习率缩放（layer-wise lr decay）。
num_epochs — 调度中的 epoch 数量（若使用 warmup_prefix，则不包括 warmup；不包括 cooldown）
decay_epochs — 学习率衰减间隔（以步数为单位），或衰减阶段（wsd / rsqrt）的长度（以 epoch 为单位），取值在 (0, 1) 之间时表示占整个调度长度的比例。None = 90（步数），0.2（wsd / rsqrt）
decay_milestones — 用于衰减学习率的 epoch（多步）
cooldown_epochs — 在调度结束后将学习率保持在 min_lr 的 epoch 数
patience_epochs — 在衰减学习率之前，没有改进的 epoch 数（plateau，已弃用）
decay_rate — 学习率衰减率（step、multistep、plateau），多项式幂次（poly）
min_lr — 衰减至该值的调度器的学习率下限
warmup_lr — 预热开始时的学习率
warmup_epochs — 线性预热轮数
warmup_prefix — 将预热期排除在 num_epochs 之外（预热期会额外计入总时长）；wsd / rsqrt 始终带有预热前缀，当 warmup_prefix=False 时，其稳定阶段和衰减阶段将包含在 num_epochs 内
noise — 学习率噪声开启/关闭点占调度总长度的比例
noise_pct — 学习率噪声限制百分比
noise_std — 学习率噪声的标准差
noise_seed — 学习率噪声种子
cycle_mul — 周期长度乘数（cosine、tanh、poly）
cycle_decay — 每个周期的学习率衰减（余弦、tanh、多项式）
cycle_limit — 周期数（cosine、tanh、poly）
k_decay — k 衰减（余弦、多项式）
decay_fn — 衰减阶段函数（wsd / rsqrt），None 表示 wsd 使用 ‘1-sqrt’，rsqrt 使用 ‘linear’
stable_timescale — 以 epoch 为单位的 rsqrt 时间尺度（rsqrt），None = warmup_epochs
plateau_mode — ‘max’ 或 ‘min’，指标方向（plateau）
step_on_epochs — 按 epoch 步进调度器，否则按更新步进（需要 updates_per_epoch）
updates_per_epoch — 每个 epoch 的优化器更新次数（step_on_epochs=False）

按名称创建学习率调度器。

调度器类

类 timm.scheduler.CosineLRScheduler

< source >

( optimizer: Optimizert_initial: intlr_min: float = 0.0cycle_mul: float = 1.0cycle_decay: float = 1.0cycle_limit: int = 1warmup_t: int = 0warmup_lr_init: float = 0.0warmup_prefix: bool = Falset_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42k_decay: float = 1.0initialize: bool = True )

带重启的余弦衰减。这在论文 https://arxiv.org/abs/1608.03983 中有描述。

灵感来自 https://github.com/allenai/allennlp/blob/master/allennlp/training/learning_rate_schedulers/cosine.py

基于k-decay: A New Method For Learning Rate Schedule的 k-decay 选项 - https://arxiv.org/abs/2004.05909

class timm.scheduler.MultiStepLRScheduler

< source >

( optimizer: Optimizerdecay_t: typing.List[int]decay_rate: float = 1.0warmup_t: int = 0warmup_lr_init: float = 0.0warmup_prefix: bool = Truet_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42initialize: bool = True )

带有预热和噪声的多步学习率调度。

在预热阶段之外，衰减里程碑会与 t + 1 进行比较，其中当 warmup_prefix=True 时，t 不包含预热阶段。例如，在没有预热的情况下，值为 10 的里程碑会在调度索引 9 处生效。

这一长期沿用的惯例与 PyTorch 的 MultiStepLR 不同，为了与现有的训练计划保持兼容而予以保留。

class timm.scheduler.PlateauLRScheduler

< source >

( optimizer: Optimizerdecay_rate: float = 0.1patience_t: int = 10threshold: float = 0.0001cooldown_t: int = 0warmup_t: int = 0warmup_lr_init: float = 0.0lr_min: float = 0.0mode: str = 'max'noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_type: str = 'normal'noise_pct: float = 0.67noise_std: float = 1.0noise_seed: typing.Optional[int] = Noneinitialize: bool = True )

每当验证损失趋于平稳时，按一定因子衰减学习率。

已弃用：此调度器将在未来版本中移除。timm 调度器被设计为关于 epoch 或更新次数的无状态函数，而此指标驱动的调度器从未如此。其状态（最佳指标、patience 和 cooldown 计数器）不会保存在训练检查点中，因此恢复的运行会重新开始 plateau 逻辑，且不会添加检查点支持。建议改用固定调度器，例如 cosine 或 wsd，sweep.py 中的 WSD cooldown 分支使用可复现的调度器在多个预算下比较退火效果。

class timm.scheduler.PolyLRScheduler

< source >

( optimizer: Optimizert_initial: intpower: float = 0.5lr_min: float = 0.0cycle_mul: float = 1.0cycle_decay: float = 1.0cycle_limit: int = 1warmup_t: int = 0warmup_lr_init: float = 0.0warmup_prefix: bool = Falset_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42k_decay: float = 1.0initialize: bool = True )

带预热、噪声和 k 衰减的多项式学习率调度器

基于k-decay: A New Method For Learning Rate Schedule的 k-decay 选项 - https://arxiv.org/abs/2004.05909

类 timm.scheduler.StepLRScheduler

< source >

( optimizer: Optimizerdecay_t: floatdecay_rate: float = 1.0warmup_t: int = 0warmup_lr_init: float = 0.0warmup_prefix: bool = Truet_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42initialize: bool = True )

类 timm.scheduler.TanhLRScheduler

< source >

( optimizer: Optimizert_initial: intlb: float = -7.0ub: float = 3.0lr_min: float = 0.0cycle_mul: float = 1.0cycle_decay: float = 1.0cycle_limit: int = 1warmup_t: int = 0warmup_lr_init: float = 0.0warmup_prefix: bool = Falset_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42initialize: bool = True )

带重启的双曲正切衰减。这在论文 https://arxiv.org/abs/1806.01593 中有描述。

class timm.scheduler.WSDLRScheduler

< source >

( optimizer: Optimizert_initial: intdecay_t: intlr_min: float = 0.0stable_fn: str = 'constant'stable_timescale: typing.Optional[float] = Nonestable_power: float = 0.5decay_fn: typing.Union[str, typing.Callable[[float], float]] = '1-sqrt'decay_power: float = 0.5decay_ratio: float = 0.1warmup_t: int = 0warmup_lr_init: float = 0.0t_in_epochs: bool = Truenoise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = Nonenoise_pct: float = 0.67noise_std: float = 1.0noise_seed: int = 42initialize: bool = True )

参数

optimizer — 要调度的优化器
t_initial — 稳定阶段和衰减阶段的长度（以 epoch 或更新次数计，参见 t_in_epochs）
decay_t — 衰减阶段的长度，即 t_initial 的最后 decay_t 步
lr_min — 衰减阶段结束时的学习率
stable_fn — 稳定阶段函数，取值为 ‘constant’、‘rsqrt’ 或 ‘power’
stable_timescale — ‘rsqrt’ / ‘power’ 稳定函数的时间尺度，默认值为 warmup_t
stable_power — “power”稳定函数的指数
decay_fn — 衰减阶段函数，可以是 ‘linear’、‘1-sqrt’、‘cosine’、‘power’、‘exp’ 之一，或一个可调用对象
decay_power — ‘power’衰减函数的指数
decay_ratio — ‘exp’ 衰减函数在归一化之前的最终比例
warmup_t — 预热长度，t_initial 的前缀
warmup_lr_init — 预热开始时的学习率
t_in_epochs — t 以 epoch 为单位（step()），否则以更新次数为单位（step_update()）
noise_range_t — 学习率噪声范围，参见 Scheduler
noise_pct — 学习率噪声限制百分比
noise_std — 学习率噪声的标准差
noise_seed — 学习率噪声种子
initialize — 根据当前学习率初始化 initial_lr 参数组字段

Warmup-Stable-Decay (WSD) 学习率调度，可选择稳定的衰减函数。

三个阶段：在 warmup_t 步内线性预热，一个稳定阶段，以及在最后 decay_t 步内衰减至 lr_min。t_initial 是稳定阶段与衰减阶段的总长度，由于预热始终作为前缀，因此调度器的总长度为 warmup_t + t_initial（参见 get_cycle_length()），之后保持 lr_min 不变。

NOTE: 与其他调度器一样，lr_min 在 t = warmup_t + t_initial 时达到，即最后一个计划步骤之后的一步，因此最后一个计划步骤比它少一个衰减步骤。如果衰减阶段较短且按 epoch 步进，这一点会很重要（例如，5 个衰减 epoch 中最后一个使用 ‘1-sqrt’ 时，其学习率约为稳定学习率的 10%）。最终退火对 WSD 很重要，因此请按更新步进（t_in_epochs=False，train.py 中的 --sched-on-updates）或添加一个冷却 epoch（保持在 lr_min，--cooldown-epochs）。

稳定函数基于预热结束后的时间（t_s)进行评估，并在 t_s = 0 时等于 1，因此预热阶段始终以基础学习率进行交接。衰减函数将衰减进度 frac（取值范围为 [0, 1]）映射为从 1 到 0 的系数，并与稳定值相乘，因此具有衰减特性的稳定函数会在衰减阶段持续衰减（例如 Big Vision 的 rsqrt + linear cooldown 调度策略）。

lr(t) = lr_min + (base_lr stable(t_s) - lr_min) decay(frac)

稳定的函数：

constant：1，WSD 调度
rsqrt：(1 + t_s / timescale) ** -0.5，当 stable_timescale == warmup_t（默认值）时，Big Vision 的倒数平方根调度
power：(1 + t_s / timescale) ** -stable_power，rsqrt 广义形式（Power 调度器）

衰减函数：

1-sqrt：1 - sqrt(frac)，在 https://arxiv.org/abs/2405.18392 中表现最佳的冷却策略（默认值）
linear：1 - frac，即“梯形”调度以及 Big Vision 冷却
cosine：0.5 (1 + cos(pi frac))
power：1 - frac ** decay_power，0.5 == ‘1-sqrt’，1.0 == ‘linear’，2.0 为 ‘square’ 衰减
exp：指数衰减至稳定值的 decay_ratio 比例，并归一化以到达 lr_min
其他任何可调用对象 fn(frac) -> factor

Update on GitHub