Learning Rate Schedulers

This page contains the API reference documentation for the learning rate schedulers included in timm, and an overview of how they are created and stepped.

Usage

Schedulers are created by name with create_scheduler_v2, which returns the scheduler and the total number of epochs to train (including any warmup prefix, cycles and cooldown, see below). Unlike torch.optim.lr_scheduler, timm schedulers do not count steps themselves. The training loop passes the epoch or update count in explicitly:

scheduler.step(epoch) at the end of each epoch, with the index of the next epoch (epoch + 1)
scheduler.step_update(num_updates) after each optimizer update, with the number of updates so far

A scheduler is stepped either per epoch or per update (step_on_epochs), the other call is a no-op, so a training loop can make both calls and support either.

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)

Durations (num_epochs, warmup_epochs, decay_epochs, …) are always given in epochs, the factory converts them to updates when step_on_epochs=False. Schedulers can also be constructed directly from their classes, durations are then in the unit they are stepped in (t_in_epochs).

The training script (train.py) exposes the factory arguments as --sched, --epochs, --warmup-epochs, --warmup-lr, --warmup-prefix, --min-lr, --decay-epochs, --decay-milestones, --decay-rate, --cooldown-epochs, --patience-epochs, --lr-noise*, --lr-cycle-*, --lr-k-decay, --lr-decay-fn and --lr-stable-timescale. --sched-on-updates steps per update instead of per epoch.

Schedules

`sched`	Class	Schedule
`cosine`	`CosineLRScheduler`	Cosine annealing from the base lr to `min_lr` (SGDR), optional restarts (cycles) and k-decay
`tanh`	`TanhLRScheduler`	Hyperbolic tangent decay (paper), optional restarts
`poly`	`PolyLRScheduler`	Polynomial decay to `min_lr`, the power is given by `decay_rate`, optional restarts and k-decay
`step`	`StepLRScheduler`	lr multiplied by `decay_rate` every `decay_epochs`
`multistep`	`MultiStepLRScheduler`	lr multiplied by `decay_rate` at each of the `decay_milestones` epochs
`plateau`	`PlateauLRScheduler`	lr multiplied by `decay_rate` when the eval metric hasn’t improved for `patience_epochs`, per epoch stepping only
`wsd`	`WSDLRScheduler`	Warmup-stable-decay (MiniCPM), constant lr then a decay to `min_lr` over the final `decay_epochs`, `1-sqrt` decay by default (paper)
`rsqrt`	`WSDLRScheduler`	Reciprocal square root decay w/ a linear cooldown over the final `decay_epochs` (Big Vision)

Common options

Warmup. warmup_epochs of linear warmup from warmup_lr to the base lr, supported by all schedules.
- warmup_prefix: by default the warmup overlaps the start of the schedule (the schedule is computed from epoch 0, warmup replaces its first epochs), so num_epochs is the total length. With warmup_prefix=True the schedule starts after the warmup and the warmup is added to the total.
- For wsd and rsqrt the stable phase always starts after the warmup, w/o warmup_prefix the stable and decay phases are fitted into num_epochs - warmup_epochs so num_epochs remains the total.
min_lr: the lr the decaying schedules (cosine, tanh, poly, wsd, rsqrt) reach at the end, and a lower bound for plateau.
cooldown_epochs: epochs held at min_lr after the schedule ends (cosine, tanh, poly, wsd, rsqrt), added to the returned num_epochs. As the schedule reaches min_lr one step after its last step, a cooldown epoch is also a way to train at min_lr when stepping per epoch.
Cycles (cosine, tanh, poly): cycle_limit > 1 restarts the schedule, each cycle cycle_mul times longer and w/ its peak lr multiplied by cycle_decay relative to the previous. The returned num_epochs covers all cycles.
k_decay (cosine, poly): k-decay changes the shape of the decay (t ** k), k > 1 keeps the lr high for longer.
Noise: noise adds random multiplicative lr noise over a range of the schedule, given as fractions of num_epochs (a start, or a start and an end), limited by noise_pct.

WSD and rsqrt

The WSDLRScheduler has three phases, a linear warmup, a stable phase and a decay to min_lr over the final decay_epochs (a value in (0, 1) is a fraction of the stable + decay length, 20% by default). The decay function is selected w/ decay_fn (--lr-decay-fn): 1-sqrt (default for wsd), linear (default for rsqrt), cosine, power or exp. The rsqrt stable phase decays as (1 + t / timescale) ** -0.5 from the end of warmup, stable_timescale (--lr-stable-timescale, in epochs) defaults to the warmup length and must be set when there is no warmup. The class also supports a generalized power stable phase and custom decay functions.

The final anneal is where WSD schedules get most of their benefit, so step per update (--sched-on-updates), or add a cooldown epoch, when the decay phase is only a few epochs long. Stepped per epoch, the last epoch of a short decay is noticeably above min_lr (~10% of the base lr for the last of 5 decay epochs w/ 1-sqrt).

# 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,
)

Parameter group lr scales

Param groups w/ an lr_scale field (e.g. layer-wise lr decay groups from param_groups_layer_decay) get the scheduled lr multiplied by their scale. The base lr of each group is recorded as initial_lr in the group when the scheduler is created.

Resuming

As the step counts are passed in, a scheduler for a resumed run is set up by creating it as for the original run and stepping it to the resume point, scheduler.step(start_epoch) or scheduler.step_update(start_epoch * updates_per_epoch). state_dict() / load_state_dict() save and restore the scheduler’s attributes.

API

Factory functions

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 )

Parameters

optimizer — optimizer to schedule
sched — scheduler name, one of ‘cosine’, ‘tanh’, ‘step’, ‘multistep’, ‘plateau’, ‘poly’, ‘wsd’, ‘rsqrt’, or ‘none’ for a constant lr (returns None)
num_epochs — number of epochs in the schedule (excluding warmup if warmup_prefix, excluding cooldown)
decay_epochs — lr decay interval (step), or the length of the decay phase (wsd / rsqrt) in epochs, a value in (0, 1) is a fraction of the schedule length. None = 90 (step), 0.2 (wsd / rsqrt)
decay_milestones — epochs to decay the lr at (multistep)
cooldown_epochs — epochs to hold the lr at min_lr after the schedule ends
patience_epochs — epochs w/o improvement before decaying the lr (plateau)
decay_rate — lr decay rate (step, multistep, plateau), the polynomial power (poly)
min_lr — lr lower bound for schedulers that decay to it
warmup_lr — lr at the start of warmup
warmup_epochs — epochs of linear warmup
warmup_prefix — exclude the warmup period from num_epochs (added to the total), wsd / rsqrt are always warmup prefixed, with warmup_prefix=False their stable + decay phases fit within num_epochs
noise — lr noise on / off points as fractions of the schedule length
noise_pct — lr noise limit percent
noise_std — lr noise std-dev
noise_seed — lr noise seed
cycle_mul — cycle length multiplier (cosine, tanh, poly)
cycle_decay — lr decay per cycle (cosine, tanh, poly)
cycle_limit — number of cycles (cosine, tanh, poly)
k_decay — k-decay (cosine, poly)
decay_fn — decay phase function (wsd / rsqrt), None = ‘1-sqrt’ for wsd, ‘linear’ for rsqrt
stable_timescale — rsqrt timescale in epochs (rsqrt), None = warmup_epochs
plateau_mode — ‘max’ or ‘min’, the metric direction (plateau)
step_on_epochs — step the schedule per epoch, otherwise per update (requires updates_per_epoch)
updates_per_epoch — number of optimizer updates per epoch (step_on_epochs=False)

Create a learning rate scheduler by name.

Scheduler Classes

class 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 )

Cosine decay with restarts. This is described in the paper https://arxiv.org/abs/1608.03983.

Inspiration from https://github.com/allenai/allennlp/blob/master/allennlp/training/learning_rate_schedulers/cosine.py

k-decay option based on k-decay: A New Method For Learning Rate Schedule - 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 )

Multi-step learning rate schedule with warmup and noise.

Outside warmup, decay milestones are compared against t + 1, where t excludes warmup when warmup_prefix=True. For example, without warmup, a milestone of 10 applies at schedule index 9.

This longstanding convention differs from PyTorch’s MultiStepLR and is retained for compatibility with existing training schedules.

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 )

Decay the LR by a factor every time the validation loss plateaus.

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 )

Polynomial LR Scheduler w/ warmup, noise, and k-decay

k-decay option based on k-decay: A New Method For Learning Rate Schedule - https://arxiv.org/abs/2004.05909

class 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 )

class 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 )

Hyberbolic-Tangent decay with restarts. This is described in the paper 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 )

Parameters

optimizer — optimizer to schedule
t_initial — length of the stable + decay phases (epochs or updates, see t_in_epochs)
decay_t — length of the decay phase, the final decay_t steps of t_initial
lr_min — lr at the end of the decay phase
stable_fn — stable phase function, one of ‘constant’, ‘rsqrt’, ‘power’
stable_timescale — timescale of the ‘rsqrt’ / ‘power’ stable functions, defaults to warmup_t
stable_power — exponent of the ‘power’ stable function
decay_fn — decay phase function, one of ‘linear’, ‘1-sqrt’, ‘cosine’, ‘power’, ‘exp’ or a callable
decay_power — exponent of the ‘power’ decay function
decay_ratio — final fraction (before normalization) of the ‘exp’ decay function
warmup_t — warmup length, a prefix to t_initial
warmup_lr_init — lr at the start of warmup
t_in_epochs — t is in epochs (step()), otherwise updates (step_update())
noise_range_t — lr noise range, see Scheduler
noise_pct — lr noise limit percent
noise_std — lr noise std-dev
noise_seed — lr noise seed
initialize — initialize the initial_lr param group field from the current lr

Warmup-Stable-Decay (WSD) LR schedule w/ selectable stable and decay functions.

Three phases: a linear warmup over warmup_t steps, a stable phase, and a decay to lr_min over the final decay_t steps. t_initial is the length of the stable + decay phases, the warmup is always a prefix so the schedule has a total length of warmup_t + t_initial (see get_cycle_length()), after which lr_min is held.

NOTE: as for the other schedulers, lr_min is reached at t = warmup_t + t_initial, one step past the last scheduled step, so the last scheduled step is one decay step short of it. With a short decay phase stepped per epoch this is significant (e.g. the last of 5 decay epochs w/ ‘1-sqrt’ is at ~10% of the stable lr). The final anneal matters for WSD, so step per update (t_in_epochs=False, --sched-on-updates in train.py) or add a cooldown epoch (held at lr_min, --cooldown-epochs).

The stable function is evaluated on the time since the end of warmup (t_s) and equals 1 at t_s = 0, so the warmup always hands off at the base lr. The decay function maps the decay progress frac in [0, 1] to a factor from 1 to 0 and multiplies the stable value, so a decaying stable function keeps decaying through the decay phase (as in the Big Vision rsqrt + linear cooldown schedule).

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

Stable functions:

constant: 1, the WSD schedule
rsqrt: (1 + t_s / timescale) ** -0.5, the Big Vision reciprocal square root schedule when stable_timescale == warmup_t (the default)
power: (1 + t_s / timescale) ** -stable_power, rsqrt generalized (Power scheduler)

Decay functions:

1-sqrt: 1 - sqrt(frac), the best performing cooldown in https://arxiv.org/abs/2405.18392 (default)
linear: 1 - frac, the ‘trapezoidal’ schedule and the Big Vision cooldown
cosine: 0.5 (1 + cos(pi frac))
power: 1 - frac ** decay_power, 0.5 == ‘1-sqrt’, 1.0 == ‘linear’, 2.0 is a ‘square’ decay
exp: exponential decay to a fraction decay_ratio of the stable value, normalized to reach lr_min
a callable fn(frac) -> factor for anything else

Update on GitHub