| import math | |
| from torch.optim.lr_scheduler import _LRScheduler, CosineAnnealingLR | |
| class ConstantLRScheduler(_LRScheduler): | |
| def __init__( | |
| self, | |
| optimizer, | |
| last_epoch: int = -1, | |
| verbose: bool = False, | |
| init_lr: float = 0.0, | |
| ): | |
| """ | |
| This is an implementation of constant learning rate scheduler. | |
| Args: | |
| optimizer: Optimizer | |
| last_epoch: The index of last epoch. Default: -1 | |
| verbose: If ``True``, prints a message to stdout for each update. Default: ``False`` | |
| init_lr: Initial learning rate | |
| """ | |
| self.init_lr = init_lr | |
| super().__init__(optimizer, last_epoch, verbose) | |
| def state_dict(self): | |
| state_dict = {k: v for k, v in self.__dict__.items() if k not in ["optimizer"]} | |
| return state_dict | |
| def load_state_dict(self, state_dict): | |
| self.__dict__.update(state_dict) | |
| def get_lr(self): | |
| if not self._get_lr_called_within_step: | |
| raise RuntimeError( | |
| "To get the last learning rate computed by the scheduler, use " | |
| "get_last_lr()" | |
| ) | |
| return [self.init_lr for group in self.optimizer.param_groups] | |
| class CosineAnnealingLRScheduler(_LRScheduler): | |
| def __init__( | |
| self, | |
| optimizer, | |
| last_epoch: int = -1, | |
| verbose: bool = False, | |
| init_lr: float = 4.0e-5, | |
| max_lr: float = 4e-4, | |
| final_lr: float = 4e-5, | |
| warmup_steps: int = 2000, | |
| cosine_steps: int = 10000, | |
| ): | |
| """ | |
| This is an implementation of cosine annealing learning rate scheduler. | |
| Args: | |
| optimizer: Optimizer | |
| last_epoch: The index of last epoch. Default: -1 | |
| verbose: If ``True``, prints a message to stdout for each update. Default: ``False`` | |
| init_lr: Initial learning rate | |
| max_lr: Maximum learning rate after warmup | |
| final_lr: Final learning rate after decay | |
| warmup_steps: Number of steps for warmup | |
| cosine_steps: Number of steps for cosine annealing | |
| """ | |
| self.init_lr = init_lr | |
| self.max_lr = max_lr | |
| self.final_lr = final_lr | |
| self.warmup_steps = warmup_steps | |
| self.cosine_steps = cosine_steps | |
| super(CosineAnnealingLRScheduler, self).__init__(optimizer, last_epoch, verbose) | |
| def state_dict(self): | |
| state_dict = {k: v for k, v in self.__dict__.items() if k not in ["optimizer"]} | |
| return state_dict | |
| def load_state_dict(self, state_dict): | |
| self.__dict__.update(state_dict) | |
| def get_lr(self): | |
| if not self._get_lr_called_within_step: | |
| raise RuntimeError( | |
| "To get the last learning rate computed by the scheduler, use " | |
| "get_last_lr()" | |
| ) | |
| step_no = self.last_epoch | |
| if step_no <= self.warmup_steps: | |
| lr = self.init_lr + step_no / self.warmup_steps * ( | |
| self.max_lr - self.init_lr | |
| ) | |
| else: | |
| lr = self.final_lr + 0.5 * (self.max_lr - self.final_lr) * ( | |
| 1 | |
| + math.cos(math.pi * (step_no - self.warmup_steps) / self.cosine_steps) | |
| ) | |
| return [lr for group in self.optimizer.param_groups] | |
| class Esm2LRScheduler(_LRScheduler): | |
| def __init__( | |
| self, | |
| optimizer, | |
| last_epoch: int = -1, | |
| verbose: bool = False, | |
| init_lr: float = 4e-5, | |
| max_lr: float = 4e-4, | |
| final_lr: float = 4e-5, | |
| warmup_steps: int = 2000, | |
| start_decay_after_n_steps: int = 500000, | |
| end_decay_after_n_steps: int = 5000000, | |
| on_use: bool = True, | |
| ): | |
| """ | |
| An implementation of ESM2's learning rate scheduler. | |
| Args: | |
| optimizer: Optimizer | |
| last_epoch: The index of last epoch. Default: -1 | |
| verbose: If ``True``, prints a message to stdout for each update. Default: ``False`` | |
| init_lr: Initial learning rate | |
| max_lr: Maximum learning rate after warmup | |
| final_lr: Final learning rate after decay | |
| warmup_steps: Number of steps for warmup | |
| start_decay_after_n_steps: Start decay after this number of steps | |
| end_decay_after_n_steps: End decay after this number of steps | |
| on_use: Whether to use this scheduler. If ``False``, the scheduler will not change the learning rate | |
| and will only use the ``init_lr``. Default: ``True`` | |
| """ | |
| self.init_lr = init_lr | |
| self.max_lr = max_lr | |
| self.final_lr = final_lr | |
| self.warmup_steps = warmup_steps | |
| self.start_decay_after_n_steps = start_decay_after_n_steps | |
| self.end_decay_after_n_steps = end_decay_after_n_steps | |
| self.on_use = on_use | |
| super(Esm2LRScheduler, self).__init__(optimizer, last_epoch, verbose) | |
| def state_dict(self): | |
| state_dict = {k: v for k, v in self.__dict__.items() if k not in ["optimizer"]} | |
| return state_dict | |
| def load_state_dict(self, state_dict): | |
| self.__dict__.update(state_dict) | |
| def get_lr(self): | |
| if not self._get_lr_called_within_step: | |
| raise RuntimeError( | |
| "To get the last learning rate computed by the scheduler, use " | |
| "get_last_lr()" | |
| ) | |
| step_no = self.last_epoch | |
| if not self.on_use: | |
| return [base_lr for base_lr in self.base_lrs] | |
| if step_no <= self.warmup_steps: | |
| lr = self.init_lr + step_no / self.warmup_steps * ( | |
| self.max_lr - self.init_lr | |
| ) | |
| elif step_no <= self.start_decay_after_n_steps: | |
| lr = self.max_lr | |
| elif step_no <= self.end_decay_after_n_steps: | |
| portion = (step_no - self.start_decay_after_n_steps) / ( | |
| self.end_decay_after_n_steps - self.start_decay_after_n_steps | |
| ) | |
| lr = self.max_lr - portion * (self.max_lr - self.final_lr) | |
| else: | |
| lr = self.final_lr | |
| return [lr for group in self.optimizer.param_groups] | |