espnet2.schedulers.piecewise_linear_warmup_lr.PiecewiseLinearWarmupLR

source

class espnet2.schedulers.piecewise_linear_warmup_lr.PiecewiseLinearWarmupLR(optimizer: Optimizer, warmup_steps_list: List[float | int] = [0, 25000], warmup_lr_list: List[float] = [0.0, 0.001], last_epoch: int = -1)

Bases: _LRScheduler, AbsBatchStepScheduler

The PiecewiseLinearWarmupLR scheduler.

This scheduler is similar to the WarmupLR Scheduler except that the warmup stage is piecewise linear. It allows for a flexible learning rate schedule during the warmup phase, enabling users to define multiple warmup steps and corresponding learning rates.

warmup_steps_list

A list of warmup steps.

• Type: List[Union[int, float]]

warmup_lr_list

A list of learning rates corresponding to the warmup steps.

• Type: List[float]

• Parameters:

  • optimizer (torch.optim.Optimizer) – The optimizer for which to schedule the learning rate.
  • warmup_steps_list (List *[*Union *[*int , float ] ] , optional) – A list of steps for the warmup phase. Default is [0, 25000].
  • warmup_lr_list (List *[*float ] , optional) – A list of learning rates that correspond to the warmup steps. Default is [0.0, 0.001].
  • last_epoch (int , optional) – The index of the last epoch. Default is -1.

• Returns: The learning rate for each parameter group.

• Return type: List[float]

####### Examples

>>> import torch
>>> from torch.optim import Adam
>>> optimizer = Adam(params, lr=0.001)
>>> scheduler = PiecewiseLinearWarmupLR(
...     optimizer,
...     warmup_steps_list=[0, 1000, 5000],
...     warmup_lr_list=[0.0, 0.01, 0.001]
... )
>>> for epoch in range(10000):
...     scheduler.step()
...     print(scheduler.get_lr())

NOTE

The learning rate will not be updated if step() is not called.

get_lr()

Computes the learning rate for the current epoch based on the warmup steps.

The learning rate is determined using piecewise linear interpolation between the specified warmup steps and corresponding learning rates. If the current epoch exceeds the last specified warmup step, the learning rate is adjusted according to a decay strategy.

• Returns: A list of learning rates for each parameter group in the optimizer.
• Return type: List[float]

####### Examples

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
>>> scheduler = PiecewiseLinearWarmupLR(optimizer,
...                                       warmup_steps_list=[0, 100, 200],
...                                       warmup_lr_list=[0.0, 0.01, 0.1])
>>> scheduler.get_lr()
[0.0, 0.005, 0.01]  # Example output during warmup

NOTE

The warmup steps and learning rates must have the same length.