summaryrefslogtreecommitdiffstats
path: root/training
diff options
context:
space:
mode:
Diffstat (limited to 'training')
-rw-r--r--training/optimization.py6
1 files changed, 3 insertions, 3 deletions
diff --git a/training/optimization.py b/training/optimization.py
index a79944f..725599b 100644
--- a/training/optimization.py
+++ b/training/optimization.py
@@ -30,7 +30,7 @@ def get_one_cycle_schedule(
30 return min_lr + progress * (1 - min_lr) 30 return min_lr + progress * (1 - min_lr)
31 31
32 lr = 0.5 * (1.0 + math.cos(math.pi * (1 + progress))) 32 lr = 0.5 * (1.0 + math.cos(math.pi * (1 + progress)))
33 lr = lr ** warmup_exp 33 lr = lr ** (warmup_exp - (warmup_exp - 1) * progress)
34 return min_lr + lr * (1 - min_lr) 34 return min_lr + lr * (1 - min_lr)
35 35
36 if annealing == "linear": 36 if annealing == "linear":
@@ -47,11 +47,11 @@ def get_one_cycle_schedule(
47 47
48 if annealing == "half_cos": 48 if annealing == "half_cos":
49 lr = 1.0 + math.cos(math.pi * (0.5 + 0.5 * progress)) 49 lr = 1.0 + math.cos(math.pi * (0.5 + 0.5 * progress))
50 lr = lr ** annealing_exp 50 lr = lr ** (annealing_exp - (annealing_exp - 1) * progress)
51 return lr 51 return lr
52 52
53 lr = 0.5 * (1.0 + math.cos(math.pi * progress)) 53 lr = 0.5 * (1.0 + math.cos(math.pi * progress))
54 lr = lr ** annealing_exp 54 lr = lr ** (annealing_exp - (annealing_exp - 1) * progress)
55 return lr 55 return lr
56 56
57 return LambdaLR(optimizer, lr_lambda, last_epoch) 57 return LambdaLR(optimizer, lr_lambda, last_epoch)