# 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`](#timm.scheduler.scheduler_factory.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. They are designed to be
stateless: the lr is a function of the schedule configuration and an explicit epoch or update count, so resuming a
run only needs the epoch, and no scheduler state is stored in checkpoints. 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.

```python
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 through these flags:

- **`--sched`** selects the schedule.
- **`--epochs`** sets the schedule length.
- **`--sched-on-updates`** steps per update instead of per epoch.
- **`--warmup-epochs`** sets the warmup duration.
- **`--warmup-lr`** sets the initial warmup learning rate.
- **`--warmup-prefix`** adds warmup before the schedule instead of overlapping its start.
- **`--min-lr`** sets the minimum learning rate.
- **`--decay-epochs`** sets the step-decay interval or the final decay duration for WSD / rsqrt.
- **`--decay-milestones`** sets the decay epochs for the multistep schedule.
- **`--decay-rate`** sets the decay factor, or the power for the polynomial schedule.
- **`--cooldown-epochs`** adds epochs at the minimum learning rate after the schedule ends.
- **`--patience-epochs`** sets the plateau scheduler's patience.
- **`--lr-noise`** selects the part of the schedule with learning rate noise.
- **`--lr-noise-pct`** limits the noise magnitude.
- **`--lr-noise-std`** sets the noise standard deviation.
- **`--lr-cycle-mul`** sets the cycle length multiplier.
- **`--lr-cycle-decay`** sets the peak learning rate multiplier between cycles.
- **`--lr-cycle-limit`** sets the number of cycles.
- **`--lr-k-decay`** sets k-decay for cosine / polynomial schedules.
- **`--lr-decay-fn`** selects the final decay function for WSD / rsqrt.
- **`--lr-stable-timescale`** sets the rsqrt stable-phase timescale.

### Schedules

| `sched` | Class | Schedule |
|---|---|---|
| `cosine` | [`CosineLRScheduler`](#timm.scheduler.cosine_lr.CosineLRScheduler) | Cosine annealing from the base lr to `min_lr` ([SGDR](https://arxiv.org/abs/1608.03983)), optional restarts (cycles) and k-decay |
| `tanh` | [`TanhLRScheduler`](#timm.scheduler.tanh_lr.TanhLRScheduler) | Hyperbolic tangent decay ([paper](https://arxiv.org/abs/1806.01593)), optional restarts |
| `poly` | [`PolyLRScheduler`](#timm.scheduler.poly_lr.PolyLRScheduler) | Polynomial decay to `min_lr`, the power is given by `decay_rate`, optional restarts and k-decay |
| `step` | [`StepLRScheduler`](#timm.scheduler.step_lr.StepLRScheduler) | lr multiplied by `decay_rate` every `decay_epochs` |
| `multistep` | [`MultiStepLRScheduler`](#timm.scheduler.multistep_lr.MultiStepLRScheduler) | lr multiplied by `decay_rate` at each of the `decay_milestones` epochs |
| `plateau` | [`PlateauLRScheduler`](#timm.scheduler.plateau_lr.PlateauLRScheduler) | **Deprecated.** lr multiplied by `decay_rate` when the eval metric hasn't improved for `patience_epochs`, per epoch stepping only |
| `wsd` | [`WSDLRScheduler`](#timm.scheduler.wsd_lr.WSDLRScheduler) | Warmup-stable-decay ([MiniCPM](https://arxiv.org/abs/2404.06395)), constant lr then a decay to `min_lr` over the final `decay_epochs`, `1-sqrt` decay by default ([paper](https://arxiv.org/abs/2405.18392)) |
| `rsqrt` | [`WSDLRScheduler`](#timm.scheduler.wsd_lr.WSDLRScheduler) | Reciprocal square root decay w/ a linear cooldown over the final `decay_epochs` ([Big Vision](https://arxiv.org/abs/2106.04560)) |

### 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`.

The `plateau` schedule is deprecated and will be removed in a future release. It never fit the stateless design
above: its lr depends on the history of the eval metric rather than on the epoch, and that history (best metric,
patience and cooldown counters) is not saved in checkpoints, so a resumed run restarts the plateau logic. Checkpoint
support will not be added. A validation plateau is also a weak signal for when to anneal; a fixed schedule such as
`cosine` or `wsd` is reproducible and resumes exactly, and a WSD sweep
(see [Hyperparameter Sweeps](../training_script#hyperparameter-sweeps)) compares cooldowns at several budgets from one
stable run.
- **`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](https://arxiv.org/abs/2004.05909) 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`](#timm.scheduler.wsd_lr.WSDLRScheduler) has three phases: a linear warmup, a stable phase and
a decay to `min_lr`.

- **`decay_epochs`** sets the final decay duration. A value in (0, 1) is a fraction of the stable + decay length,
  20% by default.
- **`decay_fn`** (`--lr-decay-fn`) selects `1-sqrt` (default for `wsd`), `linear` (default for `rsqrt`), `cosine`,
  `power` or `exp`.
- **`stable_timescale`** (`--lr-stable-timescale`, in epochs) controls the `rsqrt` stable phase, which decays as
  `(1 + t / timescale) ** -0.5` from the end of warmup. It 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`).

```python
# 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`](optimizers#timm.optim.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]]

#### timm.scheduler.create_scheduler[[timm.scheduler.create_scheduler]]

```python
timm.scheduler.create_scheduler(args, optimizer: Optimizer, updates_per_epoch: int = 0)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/scheduler_factory.py#L54)

#### timm.scheduler.create_scheduler_v2[[timm.scheduler.create_scheduler_v2]]

```python
timm.scheduler.create_scheduler_v2(optimizer: Optimizer, sched: str = 'cosine', num_epochs: int = 300, decay_epochs: typing.Optional[float] = None, decay_milestones: typing.List[int] = (90, 180, 270), cooldown_epochs: int = 0, patience_epochs: int = 10, decay_rate: float = 0.1, min_lr: float = 0.0, warmup_lr: float = 1e-05, warmup_epochs: int = 0, warmup_prefix: bool = False, noise: typing.Union[float, typing.List[float]] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, cycle_mul: float = 1.0, cycle_decay: float = 0.1, cycle_limit: int = 1, k_decay: float = 1.0, decay_fn: typing.Optional[str] = None, stable_timescale: typing.Optional[float] = None, plateau_mode: str = 'max', step_on_epochs: bool = True, updates_per_epoch: int = 0)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/scheduler_factory.py#L66)

**Parameters:**

optimizer : optimizer to schedule

sched : scheduler name, one of 'cosine', 'tanh', 'step', 'multistep', 'poly', 'wsd', 'rsqrt', 'plateau' (deprecated), 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, deprecated)

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)

**Returns:**

(scheduler, num_epochs) w/ num_epochs the total number of epochs to train including warmup prefix and cooldown

Create a learning rate scheduler by name.

### Scheduler Classes[[timm.scheduler.CosineLRScheduler]]

#### timm.scheduler.CosineLRScheduler[[timm.scheduler.CosineLRScheduler]]

```python
timm.scheduler.CosineLRScheduler(optimizer: Optimizer, t_initial: int, lr_min: float = 0.0, cycle_mul: float = 1.0, cycle_decay: float = 1.0, cycle_limit: int = 1, warmup_t: int = 0, warmup_lr_init: float = 0.0, warmup_prefix: bool = False, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, k_decay: float = 1.0, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/cosine_lr.py#L19)

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

#### timm.scheduler.MultiStepLRScheduler[[timm.scheduler.MultiStepLRScheduler]]

```python
timm.scheduler.MultiStepLRScheduler(optimizer: Optimizer, decay_t: typing.List[int], decay_rate: float = 1.0, warmup_t: int = 0, warmup_lr_init: float = 0.0, warmup_prefix: bool = True, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/multistep_lr.py#L10)

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.

#### timm.scheduler.PlateauLRScheduler[[timm.scheduler.PlateauLRScheduler]]

```python
timm.scheduler.PlateauLRScheduler(optimizer: Optimizer, decay_rate: float = 0.1, patience_t: int = 10, threshold: float = 0.0001, cooldown_t: int = 0, warmup_t: int = 0, warmup_lr_init: float = 0.0, lr_min: float = 0.0, mode: str = 'max', noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_type: str = 'normal', noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: typing.Optional[int] = None, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/plateau_lr.py#L15)

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

Deprecated: this scheduler will be removed in a future release. timm schedulers are designed to be stateless
functions of the epoch or update count, which this metric-driven schedule never was. Its state (best metric,
patience and cooldown counters) is not saved in training checkpoints, so a resumed run restarts the plateau logic,
and checkpoint support will not be added. Prefer a fixed schedule such as `cosine` or `wsd`, the WSD cooldown
branching in `sweep.py` compares annealing at several budgets with reproducible schedules.

#### timm.scheduler.PolyLRScheduler[[timm.scheduler.PolyLRScheduler]]

```python
timm.scheduler.PolyLRScheduler(optimizer: Optimizer, t_initial: int, power: float = 0.5, lr_min: float = 0.0, cycle_mul: float = 1.0, cycle_decay: float = 1.0, cycle_limit: int = 1, warmup_t: int = 0, warmup_lr_init: float = 0.0, warmup_prefix: bool = False, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, k_decay: float = 1.0, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/poly_lr.py#L19)

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

#### timm.scheduler.StepLRScheduler[[timm.scheduler.StepLRScheduler]]

```python
timm.scheduler.StepLRScheduler(optimizer: Optimizer, decay_t: float, decay_rate: float = 1.0, warmup_t: int = 0, warmup_lr_init: float = 0.0, warmup_prefix: bool = True, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/step_lr.py#L15)

#### timm.scheduler.TanhLRScheduler[[timm.scheduler.TanhLRScheduler]]

```python
timm.scheduler.TanhLRScheduler(optimizer: Optimizer, t_initial: int, lb: float = -7.0, ub: float = 3.0, lr_min: float = 0.0, cycle_mul: float = 1.0, cycle_decay: float = 1.0, cycle_limit: int = 1, warmup_t: int = 0, warmup_lr_init: float = 0.0, warmup_prefix: bool = False, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/tanh_lr.py#L19)

Hyberbolic-Tangent decay with restarts.
This is described in the paper https://arxiv.org/abs/1806.01593

#### timm.scheduler.WSDLRScheduler[[timm.scheduler.WSDLRScheduler]]

```python
timm.scheduler.WSDLRScheduler(optimizer: Optimizer, t_initial: int, decay_t: int, lr_min: float = 0.0, stable_fn: str = 'constant', stable_timescale: typing.Optional[float] = None, stable_power: float = 0.5, decay_fn: typing.Union[str, typing.Callable[[float], float]] = '1-sqrt', decay_power: float = 0.5, decay_ratio: float = 0.1, warmup_t: int = 0, warmup_lr_init: float = 0.0, t_in_epochs: bool = True, noise_range_t: typing.Union[typing.List[int], typing.Tuple[int, int], int, NoneType] = None, noise_pct: float = 0.67, noise_std: float = 1.0, noise_seed: int = 42, initialize: bool = True)
```

[Source](https://github.com/huggingface/pytorch-image-models/blob/main/timm/scheduler/wsd_lr.py#L26)

**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