cd.optim

class WarmUp(optimizer: Optimizer, steps: int, lr_lambda: Callable[[int, int], float] | List[Callable[[int, int], float]] | None = None, last_epoch: int = -1, verbose: bool = False)

WarmUp.

Applies a scaling factor to learning rate for steps 1 to steps. Applies no changes otherwise. WarmUp can be chained or used sequentally with other schedulers.

Notes

  • WarmUp sets learning rates based on captured initial learning rates (base_lrs).

  • Changes from other schedulers applied before WarmUp will be ovewritten.

  • Chaining WarmUp is equivalent to SequentialLR if other schedule dynamically manipulates optimizer.param_groups without relying on base_lrs.

Examples

>>> # First WarmUp, then other schedule
>>> from torch.optim.lr_scheduler import SequentialLR, CosineAnnealingLR
>>> from celldetection.optim import WarmUp
>>> warmup_steps = 512
>>> scheduler = SequentialLR(optimizer, [
...     WarmUp(optimizer, warmup_steps),  # warmup for 512 steps
...     CosineAnnealingLR(optimizer, T_max=512, eta_min=0.00001),  # after 512 steps switch to cosine ann.
... ], milestones=[warmup_steps])

>>> # Chaining WarmUp and other schedule
>>> from torch.optim.lr_scheduler import ChainedScheduler, StepLR
>>> scheduler = ChainedScheduler([
...     StepLR(optimizer, 5, gamma=.99),
...     WarmUp(optimizer, 512),  # overwrites changes of previous scheduler during warmup
... ])
Parameters:

  • optimizer – Optimizer.

  • steps – Number of warmup steps.

  • lr_lambda – A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups.

  • last_epoch – Last epoch.

  • verbose – If True, prints a message to stdout for each update. Default: False.

get_lr()