get_cosine_schedule_with_warmup¶
- torchtune.modules.get_cosine_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = - 1) LambdaLR[source]¶
Create a learning rate schedule that linearly increases the learning rate from 0.0 to lr over num_warmup_steps, then decreases to 0.0 on a cosine schedule over the remaining num_training_steps-num_warmup_steps (assuming num_cycles = 0.5).
This is based on the Hugging Face implementation https://github.com/huggingface/transformers/blob/v4.23.1/src/transformers/optimization.py#L104.
- Parameters:
optimizer (torch.optim.Optimizer) – The optimizer for which to schedule the learning rate.
num_warmup_steps (int) – The number of steps for the warmup phase.
num_training_steps (int) – The total number of training steps.
num_cycles (float) – The number of waves in the cosine schedule. Defaults to 0.5 (decrease from the max value to 0 following a half-cosine).
last_epoch (int) – The index of the last epoch when resuming training. Defaults to -1
- Returns:
torch.optim.lr_scheduler.LambdaLR with the appropriate schedule.
