aiaccel.torch.lr_schedulers.SequentialLR

class aiaccel.torch.lr_schedulers.SequentialLR(optimizer: Optimizer, schedulers_fn: list[Callable[[Optimizer], _LRScheduler]], milestones: list[int])

Build a sequential learning rate scheduler from scheduler factory functions.

This wrapper makes it easier to define multiple schedulers in Hydra or OmegaConf-based configurations, where each entry in schedulers_fn is a partially configured callable that receives optimizer and returns a scheduler instance.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer passed to each scheduler factory.

• schedulers_fn (list[Callable[[torch.optim.Optimizer], torch.optim.lr_scheduler._LRScheduler]]) – Factory functions that create schedulers for optimizer.

• milestones (list[int]) – Epoch indices at which to switch to the next scheduler.

Example

scheduler_generator:
  _partial_: True
  _convert_: all
  _target_: aiaccel.torch.lr_schedulers.SequentialLR
  schedulers_fn:
    - _target_: torch.optim.lr_scheduler.LinearLR
      _partial_: True
      start_factor: 1.0e-3
      end_factor: 1.0
      total_iters: 5000
    - _target_: torch.optim.lr_scheduler.CosineAnnealingLR
      _partial_: True
      T_max: 95000
  milestones: [5000]

__init__(optimizer: Optimizer, schedulers_fn: list[Callable[[Optimizer], _LRScheduler]], milestones: list[int])

Methods

__init__(optimizer, schedulers_fn, milestones)
get_last_lr()	Get the most recent learning rates computed by this scheduler.
get_lr()	Compute the next learning rate for each of the optimizer's param_groups.
load_state_dict(state_dict)	Load the scheduler's state.
recursive_undo([sched])	Recursively undo any step performed by the initialisation of schedulers.
state_dict()	Return the state of the scheduler as a dict.
step()	Perform a step.