colossalai.nn.lr_scheduler.poly

class colossalai.nn.lr_scheduler.poly.PolynomialLR(optimizer, total_steps, end_lr=0.0001, power=1.0, last_epoch=- 1, **kwargs)

Polynomial learning rate scheduler.

Parameters

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

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

• end_lr (float, optional) – Minimum learning rate, defaults to 0.0001.

• power (float, optional) – The power of polynomial, defaults to 1.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.

class colossalai.nn.lr_scheduler.poly.PolynomialWarmupLR(optimizer, total_steps, warmup_steps=0, end_lr=0.0001, power=1.0, last_epoch=- 1, **kwargs)

Polynomial learning 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.

• end_lr (float, optional) – Minimum learning rate, defaults to 0.0001.

• power (float, optional) – The power of polynomial, defaults to 1.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.