schedulers

Learning-rate schedules with a linear warm-up, complementing torch.optim.lr_scheduler.

The schedules count in steps: call scheduler.step() once per epoch for an epoch-based schedule and once per iteration for an iteration-based one; total_steps and warmup_steps are in the same unit.

Classes:

CosineWarmupLR –

Cosine decay with linear warm-up (one cycle).
PolyLR –

Polynomial decay with linear warm-up (one cycle).

CosineWarmupLR [source] ¶

CosineWarmupLR(
    optimizer: Optimizer,
    total_steps: int,
    *,
    warmup_steps: int = 0,
    warmup_lr_init: float = 0.0,
    lr_min: float = 0.0,
    warmup_prefix: bool = False,
    last_epoch: int = -1,
)

Bases: _WarmupSchedule

Cosine decay with linear warm-up (one cycle).

\[ \eta_t = \eta_\min + \tfrac{1}{2} (\eta_\text{base} - \eta_\min) \left(1 + \cos \frac{\pi\, t}{T}\right) \]

after a linear warm-up from warmup_lr_init over warmup_steps. With warmup_prefix=False (the default), \(t\) counts from step \(0\) and \(T\) is total_steps, so the warm-up takes its steps from the cosine; with warmup_prefix=True, \(t\) and \(T\) count from the end of the warm-up.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
total_steps (int) –

Length of the schedule in steps (epochs or iterations, whichever step() counts).
warmup_steps (int, default: 0 ) –

Length of the linear warm-up.
warmup_lr_init (float, default: 0.0 ) –

Learning rate at step \(0\) of the warm-up.
lr_min (float, default: 0.0 ) –

Final (and post-schedule) learning rate.
warmup_prefix (bool, default: False ) –

Count the cosine from the end of the warm-up instead of from step \(0\).
last_epoch (int, default: -1 ) –

The index of the last step, for resuming.

Example

import torch

from torch_pointcloud.optim import CosineWarmupLR

opt = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=1.0)
sched = CosineWarmupLR(opt, total_steps=10, warmup_steps=2, lr_min=0.0)
lrs = []
for _ in range(10):
    opt.step()
    sched.step()
    lrs.append(round(sched.get_last_lr()[0], 3))
lrs  # [0.5, 0.905, 0.794, 0.655, 0.5, 0.345, 0.206, 0.095, 0.024, 0.0]

PolyLR [source] ¶

PolyLR(
    optimizer: Optimizer,
    total_steps: int,
    *,
    power: float = 0.9,
    warmup_steps: int = 0,
    warmup_lr_init: float = 0.0,
    lr_min: float = 0.0,
    warmup_prefix: bool = False,
    last_epoch: int = -1,
)

Bases: _WarmupSchedule

Polynomial decay with linear warm-up (one cycle).

\[ \eta_t = \eta_\min + (\eta_\text{base} - \eta_\min) \left(1 - \frac{t}{T}\right)^{p} \]

after the warm-up; see CosineWarmupLR for how \(t\) and \(T\) are counted. The decay reaches lr_min at step total_steps; to keep a non-zero rate on the last of \(n\) steps, pass total_steps=n + 1.

Parameters:

optimizer (Optimizer) –

Wrapped optimizer.
total_steps (int) –

Length of the schedule in steps.
power (float, default: 0.9 ) –

Exponent \(p\) of the decay.
warmup_steps (int, default: 0 ) –

Length of the linear warm-up.
warmup_lr_init (float, default: 0.0 ) –

Learning rate at step \(0\) of the warm-up.
lr_min (float, default: 0.0 ) –

Final (and post-schedule) learning rate.
warmup_prefix (bool, default: False ) –

Count the decay from the end of the warm-up instead of from step \(0\).
last_epoch (int, default: -1 ) –

The index of the last step, for resuming.

Example

import torch

from torch_pointcloud.optim import PolyLR

opt = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=1.0)
sched = PolyLR(opt, total_steps=4, power=1.0)
lrs = []
for _ in range(4):
    opt.step()
    sched.step()
    lrs.append(round(sched.get_last_lr()[0], 3))
lrs  # [0.75, 0.5, 0.25, 0.0]