import math import warnings from collections import Counter from bisect import bisect_right import torch from torch.optim.lr_scheduler import _LRScheduler class LinearWarmupCosineAnnealingLR(_LRScheduler): """ Sets the learning rate of each parameter group to follow a linear warmup schedule between warmup_start_lr and base_lr followed by a cosine annealing schedule between base_lr and eta_min. .. warning:: It is recommended to call :func:`.step()` for :class:`LinearWarmupCosineAnnealingLR` after each iteration as calling it after each epoch will keep the starting lr at warmup_start_lr for the first epoch which is 0 in most cases. .. warning:: passing epoch to :func:`.step()` is being deprecated and comes with an EPOCH_DEPRECATION_WARNING. It calls the :func:`_get_closed_form_lr()` method for this scheduler instead of :func:`get_lr()`. Though this does not change the behavior of the scheduler, when passing epoch param to :func:`.step()`, the user should call the :func:`.step()` function before calling train and validation methods. Example: >>> layer = nn.Linear(10, 1) >>> optimizer = Adam(layer.parameters(), lr=0.02) >>> scheduler = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=10, max_epochs=40) >>> # >>> # the default case >>> for epoch in range(40): ... # train(...) ... # validate(...) ... scheduler.step() >>> # >>> # passing epoch param case >>> for epoch in range(40): ... scheduler.step(epoch) ... # train(...) ... # validate(...) """ def __init__( self, optimizer, warmup_epochs, max_epochs, warmup_start_lr = 0.0, eta_min = 1e-8, last_epoch = -1, ): """ Args: optimizer (Optimizer): Wrapped optimizer. warmup_epochs (int): Maximum number of iterations for linear warmup max_epochs (int): Maximum number of iterations warmup_start_lr (float): Learning rate to start the linear warmup. Default: 0. eta_min (float): Minimum learning rate. Default: 0. last_epoch (int): The index of last epoch. Default: -1. """ self.warmup_epochs = warmup_epochs self.max_epochs = max_epochs self.warmup_start_lr = warmup_start_lr self.eta_min = eta_min super(LinearWarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch) def get_lr(self): """ Compute learning rate using chainable form of the scheduler """ if not self._get_lr_called_within_step: warnings.warn( "To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.", UserWarning, ) if self.last_epoch == 0: return [self.warmup_start_lr] * len(self.base_lrs) elif self.last_epoch < self.warmup_epochs: return [ group["lr"] + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1) for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups) ] elif self.last_epoch == self.warmup_epochs: return self.base_lrs elif (self.last_epoch - 1 - self.max_epochs) % (2 * (self.max_epochs - self.warmup_epochs)) == 0: return [ group["lr"] + (base_lr - self.eta_min) * (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs))) / 2 for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups) ] return [ (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs))) / ( 1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs - 1) / (self.max_epochs - self.warmup_epochs)) ) * (group["lr"] - self.eta_min) + self.eta_min for group in self.optimizer.param_groups ] def _get_closed_form_lr(self): """ Called when epoch is passed as a param to the `step` function of the scheduler. """ if self.last_epoch < self.warmup_epochs: return [ self.warmup_start_lr + self.last_epoch * (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1) for base_lr in self.base_lrs ] return [ self.eta_min + 0.5 * (base_lr - self.eta_min) * (1 + math.cos(math.pi * (self.last_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs))) for base_lr in self.base_lrs ] class LinearWarmupMultiStepLR(_LRScheduler): """ Sets the learning rate of each parameter group to follow a linear warmup schedule between warmup_start_lr and base_lr followed by a multi-step schedule that decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones. .. warning:: It is recommended to call :func:`.step()` for :class:`LinearWarmupCosineAnnealingLR` after each iteration as calling it after each epoch will keep the starting lr at warmup_start_lr for the first epoch which is 0 in most cases. .. warning:: passing epoch to :func:`.step()` is being deprecated and comes with an EPOCH_DEPRECATION_WARNING. It calls the :func:`_get_closed_form_lr()` method for this scheduler instead of :func:`get_lr()`. Though this does not change the behavior of the scheduler, when passing epoch param to :func:`.step()`, the user should call the :func:`.step()` function before calling train and validation methods. """ def __init__( self, optimizer, warmup_epochs, milestones, warmup_start_lr = 0.0, gamma = 0.1, last_epoch = -1, ): """ Args: optimizer (Optimizer): Wrapped optimizer. warmup_epochs (int): Maximum number of iterations for linear warmup max_epochs (int): Maximum number of iterations milestones (list): List of epoch indices. Must be increasing. warmup_start_lr (float): Learning rate to start the linear warmup. Default: 0. gamma (float): Multiplicative factor of learning rate decay. Default: 0.1. last_epoch (int): The index of last epoch. Default: -1. """ self.warmup_epochs = warmup_epochs self.warmup_start_lr = warmup_start_lr self.milestones = Counter(milestones) self.gamma = gamma super(LinearWarmupMultiStepLR, self).__init__(optimizer, last_epoch) def get_lr(self): """ Compute learning rate using chainable form of the scheduler """ if not self._get_lr_called_within_step: warnings.warn("To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.", UserWarning) if self.last_epoch == 0: # starting warm up return [self.warmup_start_lr] * len(self.base_lrs) elif self.last_epoch < self.warmup_epochs: # linear warm up (0 ~ self.warmup_epochs -1) return [ group["lr"] + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1) for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups) ] elif self.last_epoch == self.warmup_epochs: # end of warm up (reset to base lrs) return self.base_lrs elif (self.last_epoch - self.warmup_epochs) not in self.milestones: # in between the steps return [group['lr'] for group in self.optimizer.param_groups] return [ group['lr'] * self.gamma ** self.milestones[self.last_epoch - self.warmup_epochs] for group in self.optimizer.param_groups ] def _get_closed_form_lr(self): """ Called when epoch is passed as a param to the `step` function of the scheduler. """ if self.last_epoch < self.warmup_epochs: return [ self.warmup_start_lr + self.last_epoch * (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1) for base_lr in self.base_lrs ] milestones = list(sorted(self.milestones.elements())) return [base_lr * self.gamma ** bisect_right(milestones, self.last_epoch - self.warmup_epochs) for base_lr in self.base_lrs]