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

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