colossalai.nn.lr_scheduler.multistep

class colossalai.nn.lr_scheduler.multistep.MultiStepLR(optimizer, total_steps, milestones=None, gamma=0.1, last_epoch=- 1, **kwargs)

Decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones. 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

• milestones (List[int], optional) – List of epoch indices. Must be increasing, defaults to None

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

• num_steps_per_epoch (int, optional) – Number of steps per epoch, defaults to -1

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

class colossalai.nn.lr_scheduler.multistep.MultiStepWarmupLR(optimizer, total_steps, warmup_steps=0, milestones=None, gamma=0.1, last_epoch=- 1, **kwargs)

Multi-step laerning rate scheduler with warmup.

Parameters

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

• total_steps (int) – Number of total training steps

• warmup_steps (int, optional) – Number of warmup steps, defaults to 0

• milestones (List[int], optional) – List of epoch indices. Must be increasing, defaults to None

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

• num_steps_per_epoch (int, optional) – Number of steps per epoch, defaults to -1

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