Merge pull request #133 from kozistr/feature/adamnorm-optimizer

[Feature] Implement AdaNorm optimizer

Author: kozistr

README.rst

@@ -140,6 +140,8 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | SM3          | *Memory-Efficient Adaptive Optimization*                                                        | `github <https://github.com/google-research/google-research/tree/master/sm3>`__   | `https://arxiv.org/abs/1901.11150 <https://arxiv.org/abs/1901.11150>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| AdaNorm`     | *Adaptive Gradient Norm Correction based Optimizer for CNNs*                                    | `github <https://github.com/shivram1987/AdaNorm>`__                               | `https://arxiv.org/abs/2210.06364 <https://arxiv.org/abs/2210.06364>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -347,6 +349,8 @@ Citations
 
 `SM3 <https://ui.adsabs.harvard.edu/abs/2019arXiv190111150A/exportcitation>`__
 
+`AdaNorm <https://github.com/shivram1987/AdaNorm/tree/main#citation>`__
+
 Citation
 --------
 

docs/optimizer_api.rst

@@ -65,14 +65,6 @@ diffGrad
 .. autoclass:: pytorch_optimizer.DiffGrad
     :members:
 
-.. _diffRGrad:
-
-diffRGrad
----------
-
-.. autoclass:: pytorch_optimizer.DiffRGrad
-    :members:
-
 .. _GC:
 
 GC
@@ -145,14 +137,6 @@ RAdam
 .. autoclass:: pytorch_optimizer.RAdam
     :members:
 
-.. _RaLamb:
-
-RaLamb
-------
-
-.. autoclass:: pytorch_optimizer.RaLamb
-    :members:
-
 .. _Ranger:
 
 Ranger
@@ -288,3 +272,11 @@ SM3
 
 .. autoclass:: pytorch_optimizer.SM3
     :members:
+
+.. _AdaNorm:
+
+AdaNorm
+-------
+
+.. autoclass:: pytorch_optimizer.AdaNorm
+    :members:

docs/util_api.rst

@@ -169,4 +169,11 @@ merge_small_dims
 
 .. autoclass:: pytorch_optimizer.merge_small_dims
     :members:
-re
+
+.. _reduce_max_except_dim:
+
+reduce_max_except_dim
+---------------------
+
+.. autoclass:: pytorch_optimizer.reduce_max_except_dim
+    :members:

pytorch_optimizer/__init__.py

@@ -23,13 +23,13 @@
 from pytorch_optimizer.optimizer.adamp import AdamP
 from pytorch_optimizer.optimizer.adams import AdamS
 from pytorch_optimizer.optimizer.adan import Adan
+from pytorch_optimizer.optimizer.adanorm import AdaNorm
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
-from pytorch_optimizer.optimizer.diffrgrad import DiffRGrad
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gsam import GSAM
@@ -43,7 +43,6 @@
 from pytorch_optimizer.optimizer.pcgrad import PCGrad
 from pytorch_optimizer.optimizer.pnm import PNM
 from pytorch_optimizer.optimizer.radam import RAdam
-from pytorch_optimizer.optimizer.ralamb import RaLamb
 from pytorch_optimizer.optimizer.ranger import Ranger
 from pytorch_optimizer.optimizer.ranger21 import Ranger21
 from pytorch_optimizer.optimizer.sam import SAM
@@ -71,6 +70,7 @@
     enable_running_stats,
     get_optimizer_parameters,
     normalize_gradient,
+    reduce_max_except_dim,
     unit_norm,
 )
 
@@ -82,14 +82,12 @@
     Adan,
     AdaPNM,
     DiffGrad,
-    DiffRGrad,
     Lamb,
     LARS,
     MADGRAD,
     Nero,
     PNM,
     RAdam,
-    RaLamb,
     Ranger,
     Ranger21,
     SGDP,
@@ -105,6 +103,7 @@
     Lion,
     AliG,
     SM3,
+    AdaNorm,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/optimizer/adabelief.py

@@ -21,7 +21,9 @@ class AdaBelief(Optimizer, BaseOptimizer):
     :param rectify: bool. perform the rectified update similar to RAdam.
     :param degenerated_to_sgd: bool. perform SGD update when variance of gradient is high.
     :param amsgrad: bool. whether to use the AMSBound variant.
-    :param adamd_debias_term: bool. Only correct the denominator to avoid inflating step sizes early in training.
+    :param r: float. EMA factor. between 0.9 ~ 0.99 is preferred.
+    :param adanorm: bool. whether to use the AdaNorm variant.
+    :param adam_debias: 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.
     """
 
@@ -37,31 +39,35 @@ def __init__(
         rectify: bool = True,
         degenerated_to_sgd: bool = True,
         amsgrad: bool = False,
-        adamd_debias_term: bool = False,
+        r: float = 0.95,
+        adanorm: bool = False,
+        adam_debias: bool = False,
         eps: float = 1e-16,
     ):
         self.lr = lr
         self.betas = betas
         self.weight_decay = weight_decay
         self.n_sma_threshold = n_sma_threshold
-        self.weight_decouple = weight_decouple
-        self.fixed_decay = fixed_decay
-        self.rectify = rectify
         self.degenerated_to_sgd = degenerated_to_sgd
-        self.adamd_debias_term = adamd_debias_term
         self.eps = eps
 
         self.validate_parameters()
 
         defaults: DEFAULTS = {
             'lr': lr,
             'betas': betas,
-            'eps': eps,
             'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'rectify': rectify,
             'amsgrad': amsgrad,
-            'adamd_debias_term': adamd_debias_term,
-            'buffer': [[None, None, None] for _ in range(10)],
+            'adanorm': adanorm,
+            'adam_debias': adam_debias,
+            'eps': eps,
         }
+        if adanorm:
+            defaults.update({'r': r})
+
         super().__init__(params, defaults)
 
     def validate_parameters(self):
@@ -76,12 +82,14 @@ def __str__(self) -> str:
     @torch.no_grad()
     def reset(self):
         for group in self.param_groups:
+            group['step'] = 0
             for p in group['params']:
                 state = self.state[p]
 
-                state['step'] = 0
                 state['exp_avg'] = torch.zeros_like(p)
                 state['exp_avg_var'] = torch.zeros_like(p)
+                if group['adanorm']:
+                    state['exp_grad_norm'] = torch.zeros((1,), dtype=p.dtype, device=p.device)
                 if group['amsgrad']:
                     state['max_exp_avg_var'] = torch.zeros_like(p)
 
@@ -93,11 +101,21 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 loss = closure()
 
         for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
             beta1, beta2 = group['betas']
-            weight_decay: float = group['weight_decay']
+            weight_decay = group['weight_decay']
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2_sq = math.sqrt(1.0 - beta2 ** group['step'])
 
-            if self.rectify:
+            if group['rectify']:
                 n_sma_max: float = 2.0 / (1.0 - beta2) - 1.0
+                beta2_t: float = beta2 ** group['step']
+                n_sma: float = n_sma_max - 2 * group['step'] * beta2_t / (1.0 - beta2_t)
 
             for p in group['params']:
                 if p.grad is None:
@@ -109,70 +127,68 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 state = self.state[p]
                 if len(state) == 0:
-                    state['step'] = 0
                     state['exp_avg'] = torch.zeros_like(p)
                     state['exp_avg_var'] = torch.zeros_like(p)
+                    if group['adanorm']:
+                        state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
                     if group['amsgrad']:
                         state['max_exp_avg_var'] = torch.zeros_like(p)
 
-                if self.weight_decouple:
-                    p.mul_(1.0 - (group['lr'] * weight_decay if not self.fixed_decay else weight_decay))
+                if group['weight_decouple']:
+                    p.mul_(1.0 - group['weight_decay'] * (1.0 if group['fixed_decay'] else group['lr']))
                 elif weight_decay > 0.0:
                     grad.add_(p, alpha=weight_decay)
 
-                state['step'] += 1
                 exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
 
-                bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
+                s_grad = grad
+                if group['adanorm']:
+                    grad_norm = torch.linalg.norm(grad)
 
-                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
-                grad_residual = grad - exp_avg
-                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2).add_(group['eps'])
+                    exp_grad_norm = state['exp_grad_norm']
+                    exp_grad_norm.mul_(group['r']).add_(grad_norm, alpha=1.0 - group['r'])
+
+                    if exp_grad_norm > grad_norm:
+                        s_grad *= exp_grad_norm / grad_norm
+
+                exp_avg.mul_(beta1).add_(s_grad, alpha=1.0 - beta1)
+                grad_residual = s_grad - exp_avg
+                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2).add_(self.eps)
 
                 if group['amsgrad']:
                     max_exp_avg_var = state['max_exp_avg_var']
                     torch.max(max_exp_avg_var, exp_avg_var, out=max_exp_avg_var)
-                    de_nom = max_exp_avg_var.sqrt()
+                    de_nom = max_exp_avg_var.add(self.eps).sqrt()
                 else:
-                    de_nom = exp_avg_var.sqrt()
-                de_nom.div_(bias_correction2_sq).add_(group['eps'])
+                    de_nom = exp_avg_var.add(self.eps).sqrt()
+
+                de_nom.div_(bias_correction2_sq).add_(self.eps)
 
-                if not self.rectify:
-                    step_size: float = group['lr'] if group['adamd_debias_term'] else group['lr'] / bias_correction1
+                if not group['rectify']:
+                    step_size: float = group['lr'] if group['adam_debias'] else group['lr'] / bias_correction1
                     p.addcdiv_(exp_avg, de_nom, value=-step_size)
                     continue
 
-                buffered = group['buffer'][state['step'] % 10]
-                if state['step'] == buffered[0]:
-                    n_sma, step_size = buffered[1], buffered[2]
+                if n_sma >= self.n_sma_threshold:
+                    step_size = math.sqrt(
+                        (1 - beta2_t)
+                        * (n_sma - 4)
+                        / (n_sma_max - 4)
+                        * (n_sma - 2)
+                        / n_sma
+                        * n_sma_max
+                        / (n_sma_max - 2)
+                    )
+                elif self.degenerated_to_sgd:
+                    step_size = 1.0
                 else:
-                    buffered[0] = state['step']
-                    beta2_t = beta2 ** state['step']
-                    n_sma = n_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
-                    buffered[1] = n_sma
-
-                    if n_sma >= self.n_sma_threshold:
-                        step_size = math.sqrt(
-                            (1 - beta2_t)
-                            * (n_sma - 4)
-                            / (n_sma_max - 4)
-                            * (n_sma - 2)
-                            / n_sma
-                            * n_sma_max
-                            / (n_sma_max - 2)
-                        )
-                        if not group['adamd_debias_term']:
-                            step_size /= bias_correction1
-                    elif self.degenerated_to_sgd:
-                        step_size = 1.0 / bias_correction1
-                    else:
-                        step_size = -1
-
-                    buffered[2] = step_size
+                    step_size = -1
+
+                if not group['adam_debias']:
+                    step_size /= bias_correction1
 
                 if n_sma >= self.n_sma_threshold:
-                    de_nom = exp_avg_var.sqrt().add_(group['eps'])
+                    de_nom = exp_avg_var.sqrt().add_(self.eps)
                     p.addcdiv_(exp_avg, de_nom, value=-step_size * group['lr'])
                 elif step_size > 0:
                     p.add_(exp_avg, alpha=-step_size * group['lr'])