colossalai.nn.lr_scheduler.torch

class colossalai.nn.lr_scheduler.torch.LambdaLR(optimizer, total_steps, lr_lambda=None, last_epoch=- 1)

Sets the learning rate of each parameter group to the initial lr times a given function. When last_epoch=-1, sets initial lr as lr.

Parameters

• optimizer (torch.optim.Optimizer) – Wrapped optimizer.

• total_steps (int) – Number of total training steps.

• lr_lambda (Union[function, list[function]]) – A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups, defaults to None.

• last_epoch (int, optional) – The index of last epoch, defaults to -1.

class colossalai.nn.lr_scheduler.torch.MultiplicativeLR(optimizer, total_steps, lr_lambda=None, last_epoch=- 1)

Multiply the learning rate of each parameter group by the factor given in the specified function. When last_epoch=-1, sets initial lr as lr.

Parameters

• optimizer (torch.optim.Optimizer) – Wrapped optimizer.

• total_steps (int) – Number of total training steps.

• lr_lambda (Union[function, list[function]]) – A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups, defaults to None.

• last_epoch (int, optional) – The index of last epoch, defaults to -1.

class colossalai.nn.lr_scheduler.torch.StepLR(optimizer, total_steps, step_size=1, gamma=0.1, last_epoch=- 1)

Decays the learning rate of each parameter group by gamma every step_size epochs. Notice that such decay can happen simultaneously with other changes to the learning rate from outside this scheduler. When last_epoch=-1, sets initial lr as lr.

Parameters

• optimizer (torch.optim.Optimizer) – Wrapped optimizer.

• total_steps (int) – Number of total training steps.

• step_size (int, optional) – Period of learning rate decay, defaults to 1.

• gamma (float, optional) – Multiplicative factor of learning rate decay, defaults to 0.1.

• last_epoch (int, optional) – The index of last epoch, defaults to -1.

class colossalai.nn.lr_scheduler.torch.ExponentialLR(optimizer, total_steps, gamma=1.0, last_epoch=- 1)

Decays the learning rate of each parameter group by gamma every epoch. When last_epoch=-1, sets initial lr as lr

Parameters

• optimizer (Union[torch.optim.Optimizer, colossalai.nn.optimizer]) – Wrapped optimizer.

• total_steps (int) – Number of total training steps.

• gamma (float, optional) – Multiplicative factor of learning rate decay, defaults to 1.0.

• last_epoch (int, optional) – The index of last epoch, defaults to -1.