| 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: |
| |
| 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 - self.warmup_epochs) not in self.milestones: |
| |
| 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] |
|
|
