colossalai.nn.lr_scheduler.delayed

class colossalai.nn.lr_scheduler.delayed.DelayerScheduler(optimizer, delay_epochs, after_scheduler, last_epoch=- 1)

Starts with a flat lr schedule until it reaches N epochs then applies the specific scheduler (For example: ReduceLROnPlateau)

Parameters

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

• delay_epochs (int) – Number of epochs to keep the initial lr until starting applying the scheduler.

• after_scheduler (torch.optim.lr_scheduler) – After target_epoch, use this scheduler.

• 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.delayed.WarmupScheduler(optimizer, warmup_epochs, after_scheduler, last_epoch=- 1)

Starts with a linear warmup lr schedule until it reaches N epochs then applies the specific scheduler (For example: ReduceLROnPlateau).

Parameters

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

• warmup_epochs (int) – Number of epochs to linearly warmup lr until starting applying the scheduler.

• after_scheduler (torch.optim.lr_scheduler) – After target_epoch, use this scheduler.

• 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.delayed.WarmupDelayerScheduler(optimizer, warmup_epochs, delay_epochs, after_scheduler, last_epoch=- 1)

Starts with a linear warmup lr schedule until it reaches N epochs and a flat lr schedule until it reaches M epochs then applies the specific scheduler (For example: ReduceLROnPlateau).

Parameters

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

• warmup_epochs (int) – Number of epochs to linearly warmup lr until starting applying the scheduler.

• delay_epochs (int) – Number of epochs to keep the initial lr until starting applying the scheduler.

• after_scheduler (torch.optim.lr_scheduler) – After target_epoch, use this scheduler.

• 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.