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