colossalai.nn.lr_scheduler.linear

class colossalai.nn.lr_scheduler.linear.LinearWarmupLR(optimizer, total_steps, warmup_steps=0, last_epoch=- 1, **kwargs)

Linearly warmup learning rate and then linearly decay.

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

• last_epoch (int, optional) – The index of last epoch, defaults to -1. When last_epoch=-1, the schedule is started from the beginning or When last_epoch=-1, sets initial lr as lr.