Merge pull request #126 from kozistr/refactor/ranger21-optimizer

[Refactor] Ranger21 optimizer

Author: kozistr

pytorch_optimizer/optimizer/ranger21.py

@@ -5,7 +5,7 @@
 from torch.nn import functional as f
 from torch.optim import Optimizer
 
-from pytorch_optimizer.base.exception import NegativeLRError, NoSparseGradientError, ZeroParameterSizeError
+from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSizeError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 from pytorch_optimizer.optimizer.agc import agc
@@ -73,6 +73,7 @@ def __init__(  # pylint: disable=R0913
         eps: float = 1e-8,
     ):
         self.lr = lr
+        self.min_lr = warm_down_min_lr
         self.beta0 = beta0
         self.betas = betas
         self.use_softplus = use_softplus
@@ -96,7 +97,6 @@ def __init__(  # pylint: disable=R0913
         # learning rate
         self.starting_lr = lr
         self.current_lr = lr
-        self.min_lr = warm_down_min_lr
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -123,6 +123,7 @@ def __init__(  # pylint: disable=R0913
 
     def validate_parameters(self):
         self.validate_learning_rate(self.lr)
+        self.validate_learning_rate(self.min_lr)
         self.validate_betas(self.betas)
         self.validate_beta0(self.beta0)
         self.validate_weight_decay(self.weight_decay)
@@ -169,7 +170,7 @@ def warm_up_dampening(self, lr: float, step: int) -> float:
 
         return new_lr
 
-    def get_warm_down(self, lr: float, iteration: int) -> float:
+    def warm_down(self, lr: float, iteration: int) -> float:
         if iteration < self.start_warm_down:
             return lr
 
@@ -182,6 +183,7 @@ def get_warm_down(self, lr: float, iteration: int) -> float:
 
         new_lr: float = self.starting_lr - self.warm_down_lr_delta * warm_down_pct
         new_lr = max(new_lr, self.min_lr)
+
         self.current_lr = new_lr
 
         return new_lr
@@ -198,6 +200,11 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         # Phase 1 - Accumulate all the variance_ma_sum to use in stable weight decay
         for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
             beta1, beta2 = group['betas']
             for p in group['params']:
                 if p.grad is None:
@@ -214,7 +221,6 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 state = self.state[p]
                 if len(state) == 0:
-                    state['step'] = 0
                     state['grad_ma'] = torch.zeros_like(p)
                     state['variance_ma'] = torch.zeros_like(p)
                     state['lookahead_params'] = torch.empty_like(p)
@@ -227,17 +233,14 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 grad = centralize_gradient(grad, gc_conv_only=False)
                 grad = normalize_gradient(grad)
 
-                state['step'] += 1
-
-                bias_correction2 = 1.0 - beta2 ** state['step']
+                bias_correction2 = 1.0 - beta2 ** group['step']
 
                 # second moment estimation
                 # using positive-negative momentum and bias correction
                 variance_ma = state['variance_ma']
                 variance_ma.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                 variance_ma_sum += (variance_ma / bias_correction2).sum()
 
-        # stable weight decay
         if param_size == 0:
             raise ZeroParameterSizeError()
 
@@ -247,42 +250,32 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         # Phase 2 - Apply weight decay and step
         for group in self.param_groups:
-            if len(self.state) == 0:
-                continue
-            p = next(iter(self.state.keys()))
-            lr = group["lr"]
-            step = self.state[p]["step"]
-
+            lr = group['lr']
             beta1, beta2 = group['betas']
-            bias_correction1 = 1.0 - beta1 ** step  # fmt: skip
-            bias_correction2_sq = math.sqrt(1.0 - beta2 ** step)  # fmt: skip
 
-            noise_norm: float = math.sqrt((1.0 + beta2) ** 2 + beta2 ** 2)  # fmt: skip
-
-            # warm up
-            lr = self.warm_up_dampening(lr, step)
-
-            # warm down
-            lr = self.get_warm_down(lr, step)
-            if lr < 0.0:
-                raise NegativeLRError(lr)
+            bias_correction1 = 1.0 - beta1 ** group['step']  # fmt: skip
+            bias_correction2_sq = math.sqrt(1.0 - beta2 ** group['step'])  # fmt: skip
 
-            # stable decay
-            decay = group['weight_decay']
-            if decay:
-                p.mul_(1.0 - decay * lr / variance_normalized)
+            noise_norm: float = math.sqrt((1.0 + beta2) ** 2 + beta2 ** 2)  # fmt: skip
 
-            # norm loss
-            u_norm = unit_norm(p)
-            correction = 2.0 * self.norm_loss_factor * (1.0 - torch.div(1, u_norm + self.eps))
-            p.mul_(1.0 - lr * correction)
+            # warm up & down
+            lr = self.warm_up_dampening(lr, group['step'])
+            lr = self.warm_down(lr, group['step'])
 
             for p in group['params']:
                 if p.grad is None:
                     continue
 
+                # stable weight decay
+                if group['weight_decay']:
+                    p.mul_(1.0 - group['weight_decay'] * lr / variance_normalized)
+
+                # norm loss
+                correction = 2.0 * self.norm_loss_factor * (1.0 - torch.div(1, unit_norm(p) + self.eps))
+                p.mul_(1.0 - lr * correction)
+
                 state = self.state[p]
-                if state['step'] % 2 == 1:
+                if group['step'] % 2 == 1:
                     grad_ma, neg_grad_ma = state['grad_ma'], state['neg_grad_ma']
                 else:
                     grad_ma, neg_grad_ma = state['neg_grad_ma'], state['grad_ma']

tests/test_optimizer_parameters.py

@@ -244,12 +244,25 @@ def test_safe_fp16_methods():
     assert optimizer.loss_scale == 2.0 ** (15 - 1)
 
 
-def test_ranger21_warm_methods():
+def test_ranger21_warm_iterations():
     assert Ranger21.build_warm_up_iterations(1000, 0.999) == 220
     assert Ranger21.build_warm_up_iterations(4500, 0.999) == 2000
     assert Ranger21.build_warm_down_iterations(1000) == 280
 
 
+def test_ranger21_warm_up_and_down():
+    param = simple_parameter(require_grad=False)
+
+    lr: float = 1e-1
+    opt = Ranger21([param], num_iterations=500, lr=lr, warm_down_min_lr=3e-5)
+
+    assert opt.warm_up_dampening(lr, 100) == 0.09090909090909091
+    assert opt.warm_up_dampening(lr, 200) == 0.1
+    assert opt.warm_up_dampening(lr, 300) == 0.1
+    assert opt.warm_down(lr, 300) == 0.1
+    assert opt.warm_down(lr, 400) == 0.07093070921985817
+
+
 @pytest.mark.parametrize('optimizer', ['ranger21', 'adai'])
 def test_size_of_parameter(optimizer):
     param = simple_parameter(require_grad=False)