| |
| |
| |
| |
| |
| |
| import math |
| from collections import Counter |
| from torch.optim.lr_scheduler import _LRScheduler |
|
|
|
|
| class MultiStepRestartLR(_LRScheduler): |
| """ MultiStep with restarts learning rate scheme. |
| |
| Args: |
| optimizer (torch.nn.optimizer): Torch optimizer. |
| milestones (list): Iterations that will decrease learning rate. |
| gamma (float): Decrease ratio. Default: 0.1. |
| restarts (list): Restart iterations. Default: [0]. |
| restart_weights (list): Restart weights at each restart iteration. |
| Default: [1]. |
| last_epoch (int): Used in _LRScheduler. Default: -1. |
| """ |
|
|
| def __init__(self, |
| optimizer, |
| milestones, |
| gamma=0.1, |
| restarts=(0, ), |
| restart_weights=(1, ), |
| last_epoch=-1): |
| self.milestones = Counter(milestones) |
| self.gamma = gamma |
| self.restarts = restarts |
| self.restart_weights = restart_weights |
| assert len(self.restarts) == len( |
| self.restart_weights), 'restarts and their weights do not match.' |
| super(MultiStepRestartLR, self).__init__(optimizer, last_epoch) |
|
|
| def get_lr(self): |
| if self.last_epoch in self.restarts: |
| weight = self.restart_weights[self.restarts.index(self.last_epoch)] |
| return [ |
| group['initial_lr'] * weight |
| for group in self.optimizer.param_groups |
| ] |
| if self.last_epoch not in self.milestones: |
| return [group['lr'] for group in self.optimizer.param_groups] |
| return [ |
| group['lr'] * self.gamma**self.milestones[self.last_epoch] |
| for group in self.optimizer.param_groups |
| ] |
|
|
| class LinearLR(_LRScheduler): |
| """ |
| |
| Args: |
| optimizer (torch.nn.optimizer): Torch optimizer. |
| milestones (list): Iterations that will decrease learning rate. |
| gamma (float): Decrease ratio. Default: 0.1. |
| last_epoch (int): Used in _LRScheduler. Default: -1. |
| """ |
|
|
| def __init__(self, |
| optimizer, |
| total_iter, |
| last_epoch=-1): |
| self.total_iter = total_iter |
| super(LinearLR, self).__init__(optimizer, last_epoch) |
|
|
| def get_lr(self): |
| process = self.last_epoch / self.total_iter |
| weight = (1 - process) |
| |
| return [weight * group['initial_lr'] for group in self.optimizer.param_groups] |
|
|
| class VibrateLR(_LRScheduler): |
| """ |
| |
| Args: |
| optimizer (torch.nn.optimizer): Torch optimizer. |
| milestones (list): Iterations that will decrease learning rate. |
| gamma (float): Decrease ratio. Default: 0.1. |
| last_epoch (int): Used in _LRScheduler. Default: -1. |
| """ |
|
|
| def __init__(self, |
| optimizer, |
| total_iter, |
| last_epoch=-1): |
| self.total_iter = total_iter |
| super(VibrateLR, self).__init__(optimizer, last_epoch) |
|
|
| def get_lr(self): |
| process = self.last_epoch / self.total_iter |
|
|
| f = 0.1 |
| if process < 3 / 8: |
| f = 1 - process * 8 / 3 |
| elif process < 5 / 8: |
| f = 0.2 |
|
|
| T = self.total_iter // 80 |
| Th = T // 2 |
|
|
| t = self.last_epoch % T |
|
|
| f2 = t / Th |
| if t >= Th: |
| f2 = 2 - f2 |
|
|
| weight = f * f2 |
|
|
| if self.last_epoch < Th: |
| weight = max(0.1, weight) |
|
|
| |
| return [weight * group['initial_lr'] for group in self.optimizer.param_groups] |
|
|
| def get_position_from_periods(iteration, cumulative_period): |
| """Get the position from a period list. |
| |
| It will return the index of the right-closest number in the period list. |
| For example, the cumulative_period = [100, 200, 300, 400], |
| if iteration == 50, return 0; |
| if iteration == 210, return 2; |
| if iteration == 300, return 2. |
| |
| Args: |
| iteration (int): Current iteration. |
| cumulative_period (list[int]): Cumulative period list. |
| |
| Returns: |
| int: The position of the right-closest number in the period list. |
| """ |
| for i, period in enumerate(cumulative_period): |
| if iteration <= period: |
| return i |
|
|
|
|
| class CosineAnnealingRestartLR(_LRScheduler): |
| """ Cosine annealing with restarts learning rate scheme. |
| |
| An example of config: |
| periods = [10, 10, 10, 10] |
| restart_weights = [1, 0.5, 0.5, 0.5] |
| eta_min=1e-7 |
| |
| It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the |
| scheduler will restart with the weights in restart_weights. |
| |
| Args: |
| optimizer (torch.nn.optimizer): Torch optimizer. |
| periods (list): Period for each cosine anneling cycle. |
| restart_weights (list): Restart weights at each restart iteration. |
| Default: [1]. |
| eta_min (float): The mimimum lr. Default: 0. |
| last_epoch (int): Used in _LRScheduler. Default: -1. |
| """ |
|
|
| def __init__(self, |
| optimizer, |
| periods, |
| restart_weights=(1, ), |
| eta_min=0, |
| last_epoch=-1): |
| self.periods = periods |
| self.restart_weights = restart_weights |
| self.eta_min = eta_min |
| assert (len(self.periods) == len(self.restart_weights) |
| ), 'periods and restart_weights should have the same length.' |
| self.cumulative_period = [ |
| sum(self.periods[0:i + 1]) for i in range(0, len(self.periods)) |
| ] |
| super(CosineAnnealingRestartLR, self).__init__(optimizer, last_epoch) |
|
|
| def get_lr(self): |
| idx = get_position_from_periods(self.last_epoch, |
| self.cumulative_period) |
| current_weight = self.restart_weights[idx] |
| nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1] |
| current_period = self.periods[idx] |
|
|
| return [ |
| self.eta_min + current_weight * 0.5 * (base_lr - self.eta_min) * |
| (1 + math.cos(math.pi * ( |
| (self.last_epoch - nearest_restart) / current_period))) |
| for base_lr in self.base_lrs |
| ] |
|
|
