LinearWarmupCyclicAnnealing

LinearWarmupCyclicAnnealing#

class stable_ssl.optim.LinearWarmupCyclicAnnealing(optimizer, total_steps, start_factor=0.01, peak_step=0.1)[source]#

Bases:

Combine linear warmup with cyclic cosine annealing.

This function creates a scheduler that combines linear warmup with cyclic cosine annealing. The cyclic annealing provides multiple learning rate cycles which can help escape local minima during training.

Parameters:
  • optimizer (torch.optim.Optimizer) – The optimizer to schedule.

  • total_steps (int) – Total number of training steps.

  • start_factor (float, optional) – Initial learning rate factor for warmup. Defaults to 0.01.

  • peak_step (float, optional) – Step at which warmup ends (as fraction of total_steps). Defaults to 0.1.

Returns:

Combined warmup and cyclic annealing scheduler.

Return type:

torch.optim.lr_scheduler.SequentialLR

Example

>>> optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
>>> scheduler = LinearWarmupCyclicAnnealing(optimizer, total_steps=1000)