Merge pull request #145 from kozistr/feature/aggmo-optimizer

[Feature] Implement AggMo optimizer

Author: kozistr

README.rst

@@ -167,6 +167,8 @@ You can check the supported optimizers & lr schedulers.
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | AdaMod       | *An Adaptive and Momental Bound Method for Stochastic Learning*                                 | `github <https://github.com/lancopku/AdaMod>`__                                   | `https://arxiv.org/abs/1910.12249 <https://arxiv.org/abs/1910.12249>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| AggMo        | *Aggregated Momentum: Stability Through Passive Damping*                                        | `github <https://github.com/AtheMathmo/AggMo`__                                   | `https://arxiv.org/abs/1804.00325 <https://arxiv.org/abs/1804.00325>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -396,6 +398,8 @@ Citations
 
 `AdaMod <https://github.com/lancopku/AdaMod#citation>`__
 
+`AggMo <https://ui.adsabs.harvard.edu/abs/2018arXiv180400325L/exportcitation>`__
+
 Citation
 --------
 

docs/optimizer_api.rst

@@ -368,3 +368,11 @@ AdaMod
 
 .. autoclass:: pytorch_optimizer.AdaMod
     :members:
+
+.. _AggMo:
+
+AggMo
+-----
+
+.. autoclass:: pytorch_optimizer.AggMo
+    :members:

pytorch_optimizer/__init__.py

@@ -28,6 +28,7 @@
 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.aggmo import AggMo
 from pytorch_optimizer.optimizer.alig import AliG
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
@@ -106,6 +107,7 @@
     ScalableShampoo,
     DAdaptAdaGrad,
     Fromage,
+    AggMo,
     DAdaptAdam,
     DAdaptSGD,
     DAdaptAdan,

pytorch_optimizer/optimizer/aggmo.py

@@ -0,0 +1,97 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class AggMo(Optimizer, BaseOptimizer):
+    r"""Aggregated Momentum: Stability Through Passive Damping.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.0, 0.9, 0.99),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = False,
+    ):
+        self.lr = lr
+        self.betas = betas
+        self.weight_decay = weight_decay
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+        }
+        super().__init__(params, defaults)
+
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_betas(self.betas)
+        self.validate_weight_decay(self.weight_decay)
+
+    def __str__(self) -> str:
+        return 'AggMo'
+
+    @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['momentum_buffer'] = {beta: torch.zeros_like(p) for beta in group['betas']}
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            betas = group['betas']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['momentum_buffer'] = {beta: torch.zeros_like(p) for beta in betas}
+
+                if group['weight_decouple']:
+                    p.mul_(1.0 - group['weight_decay'] * group['lr'])
+                elif group['weight_decay'] > 0.0:
+                    grad.add_(p, alpha=group['weight_decay'])
+
+                for beta in betas:
+                    buf = state['momentum_buffer'][beta]
+                    buf.mul_(beta).add_(grad)
+
+                    p.add_(buf, alpha=-group['lr'] / len(betas))
+
+        return loss

tests/constants.py

@@ -22,6 +22,7 @@
     Adan,
     AdaNorm,
     AdaPNM,
+    AggMo,
     AliG,
     Apollo,
     DAdaptAdaGrad,
@@ -334,6 +335,8 @@
     (MSVAG, {'lr': 5e-1}, 10),
     (AdaMod, {'lr': 5e1, 'weight_decay': 1e-3}, 10),
     (AdaMod, {'lr': 5e1, 'weight_decay': 1e-3, 'weight_decouple': False}, 10),
+    (AggMo, {'lr': 5e0, 'weight_decay': 1e-3}, 5),
+    (AggMo, {'lr': 5e0, 'weight_decay': 1e-3, 'weight_decouple': True}, 5),
 ]
 ADANORM_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adanorm': True}, 10),

tests/test_general_optimizer_parameters.py

@@ -34,6 +34,7 @@ def test_epsilon(optimizer_name):
         'sgdw',
         'fromage',
         'msvag',
+        'aggmo',
     ):
         pytest.skip(f'skip {optimizer_name} optimizer')
 
@@ -134,7 +135,7 @@ def test_betas(optimizer_name):
         config1.update({'num_iterations': 100})
         config2.update({'num_iterations': 100})
 
-    if optimizer_name not in ('adapnm', 'adan', 'adamod'):
+    if optimizer_name not in ('adapnm', 'adan', 'adamod', 'aggmo'):
         with pytest.raises(ValueError):
             optimizer(None, **config1)
 

tests/test_load_optimizers.py

@@ -16,4 +16,4 @@ def test_load_optimizers_invalid(invalid_optimizer_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 39
+    assert len(get_supported_optimizers()) == 40