lr_scheduler

class mfai.pytorch.lr_scheduler.LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs, max_epochs, warmup_start_lr=0.0, eta_min=0.0, last_epoch=-1)[source]

Bases: _LRScheduler

Sets the learning rate of each parameter group to follow a linear warmup schedule between warmup_start_lr and base_lr followed by a cosine annealing schedule between base_lr and eta_min.

Warning

It is recommended to call step() for LinearWarmupCosineAnnealingLR after each iteration as calling it after each epoch will keep the starting lr at warmup_start_lr for the first epoch which is 0 in most cases.

Warning

passing epoch to step() is being deprecated and comes with an EPOCH_DEPRECATION_WARNING. It calls the _get_closed_form_lr() method for this scheduler instead of get_lr(). Though this does not change the behavior of the scheduler, when passing epoch param to step(), the user should call the step() function before calling train and validation methods.

Example

>>> import torch.nn as nn
>>> from torch.optim import Adam
>>> #
>>> layer = nn.Linear(10, 1)
>>> optimizer = Adam(layer.parameters(), lr=0.02)
>>> scheduler = LinearWarmupCosineAnnealingLR(
...     optimizer, warmup_epochs=10, max_epochs=40
... )
>>> # the default case
>>> for epoch in range(40):
...     # train(...)
...     # validate(...)
...     scheduler.step()
>>> # passing epoch param case
>>> for epoch in range(40):
...     scheduler.step(epoch)
...     # train(...)
...     # validate(...)
Parameters:
get_lr()[source]

Compute learning rate using chainable form of the scheduler.

Return type:

list[float]