diff options
author | Jordan Gong <jordan.gong@protonmail.com> | 2022-03-15 17:29:40 +0800 |
---|---|---|
committer | Jordan Gong <jordan.gong@protonmail.com> | 2022-03-15 17:29:40 +0800 |
commit | d7751c46315a6126950a283d8a32a8d14686ded2 (patch) | |
tree | 05f2d96e0a34b09d7d1fe45735d96c9d075e600e /supervised/scheduler.py |
Add supervised baseline on CIFAR10
Diffstat (limited to 'supervised/scheduler.py')
-rw-r--r-- | supervised/scheduler.py | 29 |
1 files changed, 29 insertions, 0 deletions
diff --git a/supervised/scheduler.py b/supervised/scheduler.py new file mode 100644 index 0000000..828e547 --- /dev/null +++ b/supervised/scheduler.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] |