LinearWarmupThreeStepsAnnealing

LinearWarmupThreeStepsAnnealing#

class stable_ssl.optim.LinearWarmupThreeStepsAnnealing(optimizer, total_steps, start_factor=0.001, gamma=0.3, peak_step=0.05)[source]#

Bases:

Combine linear warmup with a three-step learning rate annealing.

This function creates a scheduler that combines linear warmup with a three-step annealing schedule. The annealing reduces the learning rate at three predefined milestones, which can help with fine-tuning and convergence.

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.001.

  • gamma (float, optional) – Multiplicative factor for learning rate reduction. Defaults to 0.3.

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

Returns:

Combined warmup and three-step annealing scheduler.

Return type:

torch.optim.lr_scheduler.SequentialLR

Example

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