Merge pull request #80 from kozistr/feature/ranger21-optimizer

[Feature] Ranger21 with AdamD

Author: kozistr

pytorch_optimizer/optimizer/ranger21.py

@@ -11,15 +11,6 @@
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.utils import normalize_gradient, unit_norm
 
-__AUTHORS__ = [
-    '@lessw2020',
-    '@NestorDemeure',
-    # with contributions from :
-    '@BrianPugh',
-    '@Kayuksel',
-    '@TheZothen',
-]
-
 
 class Ranger21(Optimizer, BaseOptimizer):
     """
@@ -38,7 +29,7 @@ class Ranger21(Optimizer, BaseOptimizer):
           optimizer.step()
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=R0913
         self,
         params: PARAMETERS,
         num_iterations: int,
@@ -58,6 +49,7 @@ def __init__(
         lookahead_blending_alpha: float = 0.5,
         weight_decay: float = 1e-4,
         norm_loss_factor: float = 1e-4,
+        adamd_debias_term: bool = False,
         eps: float = 1e-8,
     ):
         """Ranger21 optimizer
@@ -76,6 +68,7 @@ def __init__(
         :param lookahead_blending_alpha: float. blending alpha
         :param weight_decay: float. weight decay (L2 penalty)
         :param norm_loss_factor: float. norm loss factor
+        :param adamd_debias_term: bool.Only correct the denominator to avoid inflating step sizes early in training
         :param eps: float. term added to the denominator to improve numerical stability
         """
         self.lr = lr
@@ -91,6 +84,7 @@ def __init__(
         self.lookahead_blending_alpha = lookahead_blending_alpha
         self.weight_decay = weight_decay
         self.norm_loss_factor = norm_loss_factor
+        self.adamd_debias_term = adamd_debias_term
         self.eps = eps
 
         self.validate_parameters()
@@ -108,6 +102,7 @@ def __init__(
             betas=betas,
             eps=eps,
             weight_decay=weight_decay,
+            adamd_debias_term=adamd_debias_term,
         )
         super().__init__(params, defaults)
 
@@ -240,6 +235,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 variance_ma_sum += (variance_ma / bias_correction2).sum()
 
         # stable weight decay
+        if param_size == 0:
+            raise ValueError('[-] size of parameter is 0')
+
         variance_normalized = math.sqrt(variance_ma_sum / param_size)
         if math.isnan(variance_normalized):
             raise RuntimeError('hit nan for variance_normalized')
@@ -299,7 +297,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 noise_norm: float = math.sqrt((1.0 + beta2) ** 2 + beta2 ** 2)
 
-                step_size: float = lr / bias_correction1
+                step_size: float = lr
+                if not group['adamd_debias_term']:
+                    step_size /= bias_correction1
 
                 if self.use_softplus:
                     de_nom = F.softplus(de_nom, beta=self.beta_softplus)

tests/test_optimizer_parameters.py

@@ -1,7 +1,9 @@
 from typing import List
 
 import pytest
+import torch
 from torch import nn
+from torch.nn import functional as F
 
 from pytorch_optimizer import SAM, AdamP, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizer
 from tests.utils import Example
@@ -205,7 +207,29 @@ def test_safe_fp16_methods():
     assert optimizer.loss_scale == 2.0 ** (15 - 1)
 
 
-def test_ranger21_methods():
+def test_ranger21_warm_methods():
     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_size_of_parameter():
+    model: nn.Module = nn.Linear(1, 1, bias=False)
+    model.requires_grad_(False)
+
+    with pytest.raises(ValueError):
+        Ranger21(model.parameters(), 100).step()
+
+
+def test_ranger21_closure():
+    model: nn.Module = Example()
+    optimizer = Ranger21(model.parameters(), num_iterations=100, betas=(0.9, 1e-9))
+
+    loss_fn = nn.BCEWithLogitsLoss()
+
+    def closure():
+        loss = loss_fn(torch.ones((1, 1)), model(torch.ones((1, 1))))
+        loss.backward()
+        return loss
+
+    optimizer.step(closure)

tests/test_optimizers.py

@@ -97,6 +97,7 @@
     (RaLamb, {'lr': 1e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 500),
     (RAdam, {'lr': 1e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 200),
     (Ranger, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 200),
+    (Ranger21, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 200),
 ]
 
 
@@ -247,11 +248,14 @@ def closure():
 def test_adamd_optimizers(optimizer_adamd_config):
     (x_data, y_data), model, loss_fn = build_environment()
 
-    optimizer_class, config, iterations = optimizer_adamd_config
+    optimizer_class, config, num_iterations = optimizer_adamd_config
+    if optimizer_class.__name__ == 'Ranger21':
+        config.update({'num_iterations': num_iterations})
+
     optimizer = optimizer_class(model.parameters(), **config)
 
     init_loss, loss = np.inf, np.inf
-    for _ in range(iterations):
+    for _ in range(num_iterations):
         optimizer.zero_grad()
 
         y_pred = model(x_data)