helios.scheduler.schedulers

Classes

CosineAnnealingRestartLR	A cosine annealing with restarts LR scheduler.
LinearWarmupScheduler	A linear warmup 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[float] | 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) restart weights at each restart iteration. Defaults to [1].

• eta_min – (optional) the minimum lr.

• last_epoch – (optional) index of the last epoch seen by the scheduler.

get_lr() → list[float | torch.Tensor]

Return the current learning rate.

class helios.scheduler.schedulers.LinearWarmupScheduler(optimizer: torch.optim.Optimizer, warmup_steps: int, scheduler: torch.optim.lr_scheduler.LRScheduler, warmup_start_factor: float = 0.0, last_epoch: int = -1)

Bases: torch.optim.lr_scheduler.LRScheduler

A linear warmup LR scheduler.

The scheduler functions as follows:

#. During the first warmup_steps, increase the learning rate lineary from base_lr * warmup_start_factor to base_lr. #. After warmup_steps, switch to the provided scheduler.

Parameters:

• optimizer – the optimizer.

• warmup_steps – number of warmup steps.

• scheduler – the post-warmup scheduler.

• warmup_start_factor – (optional) fraction of base_lr to start from.

• last_epoch – (optional) index of the last epoch seen by the scheduler.

get_lr() → list[float | torch.Tensor]

Return the current learning rate.

step(epoch: int | None = None) → None

Advance the scheduler by one step.

class helios.scheduler.schedulers.MultiStepRestartLR(optimizer: torch.optim.Optimizer, milestones: list[int], gamma: float = 0.1, restarts: list[int] | None = None, restart_weights: list[float] | 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 – (optional) decrease ratio.

• restarts – (optional) restart iterations. Defaults to [0].

• restart_weights – (optional) restart weights at each restart iteration. Defaults to [1].

• last_epoch – (optional) index of the last epoch seen by the scheduler.

get_lr() → list[float | torch.Tensor]

Return the current learning rate.