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 the applies a scheduler

Parameters

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

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

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

• last_epoch (int, optional) – The index of last epoch, defaults to -1

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 the applies a scheduler

Parameters

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

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

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

• last_epoch (int, optional) – The index of last epoch, defaults to -1

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 the applies a scheduler

Parameters

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

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

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

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

• last_epoch (int, optional) – The index of last epoch, defaults to -1