-
Notifications
You must be signed in to change notification settings - Fork 3
/
sgdr.py
77 lines (68 loc) · 2.81 KB
/
sgdr.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import torch
import logging
import numpy as np
logger = logging.getLogger(__name__)
class CosineWithRestarts(torch.optim.lr_scheduler._LRScheduler): # pylint: disable=protected-access
"""
Cosine annealing with restarts.
This is decribed in the paper https://arxiv.org/abs/1608.03983.
Parameters
----------
optimizer : ``torch.optim.Optimizer``
t_max : ``int``
The maximum number of iterations within the first cycle.
eta_min : ``float``, optional (default=0)
The minimum learning rate.
last_epoch : ``int``, optional (default=-1)
The index of the last epoch. This is used when restarting.
factor : ``float``, optional (default=1)
The factor by which the cycle length (``T_max``) increases after each restart.
"""
def __init__(self,
optimizer: torch.optim.Optimizer,
t_max: int,
eta_min: float = 0.,
last_epoch: int = -1,
factor: float = 1.) -> None:
assert t_max > 0
assert eta_min >= 0
if t_max == 1 and factor == 1:
logger.warning("Cosine annealing scheduler will have no effect on the learning "
"rate since T_max = 1 and factor = 1.")
self.t_max = t_max
self.eta_min = eta_min
self.factor = factor
self._last_restart: int = 0
self._cycle_counter: int = 0
self._cycle_factor: float = 1.
self._updated_cycle_len: int = t_max
self._initialized: bool = False
super(CosineWithRestarts, self).__init__(optimizer, last_epoch)
def get_lr(self):
"""Get updated learning rate."""
# HACK: We need to check if this is the first time ``self.get_lr()`` was called,
# since ``torch.optim.lr_scheduler._LRScheduler`` will call ``self.get_lr()``
# when first initialized, but the learning rate should remain unchanged
# for the first epoch.
if not self._initialized:
self._initialized = True
return self.base_lrs
step = self.last_epoch + 1
self._cycle_counter = step - self._last_restart
lrs = [
self.eta_min + ((lr - self.eta_min) / 2) * (
np.cos(
np.pi *
(self._cycle_counter % self._updated_cycle_len) /
self._updated_cycle_len
) + 1
)
for lr in self.base_lrs
]
if self._cycle_counter % self._updated_cycle_len == 0:
# Adjust the cycle length.
self._cycle_factor *= self.factor
self._cycle_counter = 0
self._updated_cycle_len = int(self._cycle_factor * self.t_max)
self._last_restart = step
return lrs