helios.scheduler.schedulers

Classes

CosineAnnealingRestartLR	A cosine annealing with restarts LR scheduler.
MultiStepRestartLR	Multi-step with restarts LR scheduler.

Module Contents

class helios.scheduler.schedulers.CosineAnnealingRestartLR(optimizer: torch.optim.Optimizer, periods: list[int], restart_weights: list[int] | None = None, eta_min: float = 0, last_epoch: int = -1)

Bases: torch.optim.lr_scheduler.LRScheduler

A cosine annealing with restarts LR scheduler.

Example

Given

periods = [10, 10, 10, 10]
restart_weights = [1, 0.5, 0.5, 0.5]
eta_min = 1e-7

Then the scheduler will have 4 cycles of 10 iterations each. At the 10th, 20th, and 30th, the scheduler will restart with the weights in restart_weights.

Parameters:

• optimizer – the optimizer.

• periods – period for each cosine annealing cycle.

• restart_weights – (optional) restarts weights at each restart iteration.

• eta_min – The minimum lr. Defaults to 0

• last_epoch – Used in _LRScheduler. Defaults to -1.

get_lr()

Return the current learning rate.

class helios.scheduler.schedulers.MultiStepRestartLR(optimizer: torch.optim.Optimizer, milestones: list[int], gamma: float = 0.1, restarts: list[int] | None = None, restart_weights: list[int] | None = None, last_epoch: int = -1)

Bases: torch.optim.lr_scheduler.LRScheduler

Multi-step with restarts LR scheduler.

Parameters:

• optimizer – torch optimizer.

• milestones – iterations that will decrease learning rate.

• gamma – decrease ratio. Defaults to 0.1.

• restarts – (optional) restart iterations.

• restart_weights – (optional) restart weights at each restart iteration.

• last_epoch – used in _LRScheduler. Defaults to -1.

get_lr()

Return the current learning rate.