From 5869d0248fa958acd3447e6bffa8761b91e8e921 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Wed, 16 Mar 2022 17:49:51 +0800 Subject: Regular refactoring --- supervised/schedulers.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 supervised/schedulers.py (limited to 'supervised/schedulers.py') diff --git a/supervised/schedulers.py b/supervised/schedulers.py new file mode 100644 index 0000000..828e547 --- /dev/null +++ b/supervised/schedulers.py @@ -0,0 +1,29 @@ +import warnings + +import numpy as np +import torch + + +class LinearWarmupAndCosineAnneal(torch.optim.lr_scheduler._LRScheduler): + def __init__(self, optimizer, warm_up, T_max, last_epoch=-1): + self.warm_up = int(warm_up * T_max) + self.T_max = T_max - self.warm_up + super().__init__(optimizer, last_epoch=last_epoch) + + def get_lr(self): + 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()`.") + + if self.last_epoch == 0: + return [lr / (self.warm_up + 1) for lr in self.base_lrs] + elif self.last_epoch <= self.warm_up: + c = (self.last_epoch + 1) / self.last_epoch + return [group['lr'] * c for group in self.optimizer.param_groups] + else: + # ref: https://github.com/pytorch/pytorch/blob/2de4f245c6b1e1c294a8b2a9d7f916d43380af4b/torch/optim/lr_scheduler.py#L493 + le = self.last_epoch - self.warm_up + return [(1 + np.cos(np.pi * le / self.T_max)) / + (1 + np.cos(np.pi * (le - 1) / self.T_max)) * + group['lr'] + for group in self.optimizer.param_groups] -- cgit v1.2.3