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
78
79
80
81
82
83
84
85
|
import math
from typing import NamedTuple, Literal, Callable
from functools import partial
import torch
from torch.optim.lr_scheduler import LambdaLR
class OneCyclePhase(NamedTuple):
step_min: int
step_max: int
min: float
max: float
func: Callable[[float], float]
def warmup_linear(progress: float):
return progress
def warmup_cos(exp: int, progress: float):
lr = 0.5 * (1.0 + math.cos(math.pi * (1 + progress)))
lr = lr ** (exp - (exp - 1) * progress)
return lr
def anneal_linear(progress: float):
return 1 - progress
def anneal_half_cos(exp: int, progress: float):
lr = 1.0 + math.cos(math.pi * (0.5 + 0.5 * progress))
lr = lr ** (exp - (exp - 1) * progress)
return lr
def anneal_cos(exp: int, progress: float):
lr = 0.5 * (1.0 + math.cos(math.pi * progress))
lr = lr ** (exp - (exp - 1) * progress)
return lr
def get_one_cycle_schedule(
optimizer: torch.optim.Optimizer,
num_training_steps: int,
warmup: Literal["cos", "linear"] = "cos",
annealing: Literal["cos", "half_cos", "linear"] = "cos",
warmup_exp: int = 1,
annealing_exp: int = 1,
min_lr: int = 0.04,
mid_point: int = 0.3,
last_epoch: int = -1
):
if warmup == "linear":
warmup_func = warmup_linear
else:
warmup_func = partial(warmup_cos, warmup_exp)
if annealing == "linear":
anneal_func = anneal_linear
elif annealing == "half_cos":
anneal_func = partial(anneal_half_cos, annealing_exp)
else:
anneal_func = partial(anneal_cos, annealing_exp)
thresh_up = int(num_training_steps * min(mid_point, 0.5))
if annealing == "linear":
thresh_down = thresh_up * 2
phases = [
OneCyclePhase(0, thresh_up, min_lr, 1, warmup_func),
OneCyclePhase(thresh_up, thresh_down, min_lr, 1, anneal_func),
OneCyclePhase(thresh_down, num_training_steps, 0, min_lr, anneal_func),
]
else:
phases = [
OneCyclePhase(0, thresh_up, min_lr, 1, warmup_func),
OneCyclePhase(thresh_up, num_training_steps, 0, 1, anneal_func),
]
def lr_lambda(current_step: int):
phase = [p for p in phases if current_step >= p.step_min][-1]
return phase.min + phase.func((current_step - phase.step_min) / (phase.step_max - phase.step_min)) * (phase.max - phase.min)
return LambdaLR(optimizer, lr_lambda, last_epoch)
|
