Merge pull request #69 from kozistr/feature/adan

[Feature] Implement Adan optimizer

Author: kozistr

README.rst

@@ -87,6 +87,8 @@ Supported Optimizers
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | Nero         | *Learning by Turning: Neural Architecture Aware Optimisation*                          | `github <https://github.com/jxbz/nero>`__                                         | `https://arxiv.org/abs/2102.07227 <https://arxiv.org/abs/2102.07227>`__                       |
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| Adan         | * Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models*              | `github <https://github.com/sail-sg/Adan>`__                                      | `https://arxiv.org/abs/2208.06677 <https://arxiv.org/abs/2208.06677>`__                       |
++--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -504,6 +506,26 @@ Nero: Learning by Turning: Neural Architecture Aware Optimisation
       eprint={arXiv:2102.07227}
     }
 
+Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models
+
+::
+
+    @ARTICLE{2022arXiv220806677X,
+      author = {{Xie}, Xingyu and {Zhou}, Pan and {Li}, Huan and {Lin}, Zhouchen and {Yan}, Shuicheng},
+      title = "{Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models}",
+      journal = {arXiv e-prints},
+      keywords = {Computer Science - Machine Learning, Mathematics - Optimization and Control},
+      year = 2022,
+      month = aug,
+      eid = {arXiv:2208.06677},
+      pages = {arXiv:2208.06677},
+      archivePrefix = {arXiv},
+      eprint = {2208.06677},
+      primaryClass = {cs.LG},
+      adsurl = {https://ui.adsabs.harvard.edu/abs/2022arXiv220806677X},
+      adsnote = {Provided by the SAO/NASA Astrophysics Data System}
+    }
+
 Author
 ------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "1.1.4"
+version = "1.2.0"
 description = "Bunch of optimizer implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]

pytorch_optimizer/__init__.py

@@ -1,11 +1,10 @@
 # pylint: disable=unused-import
 from typing import Callable, Dict, List
 
-from torch.optim import Optimizer
-
 from pytorch_optimizer.adabelief import AdaBelief
 from pytorch_optimizer.adabound import AdaBound
 from pytorch_optimizer.adamp import AdamP
+from pytorch_optimizer.adan import Adan
 from pytorch_optimizer.adapnm import AdaPNM
 from pytorch_optimizer.agc import agc
 from pytorch_optimizer.chebyshev_schedule import get_chebyshev_schedule
@@ -39,6 +38,7 @@
     AdaBelief,
     AdaBound,
     AdamP,
+    Adan,
     AdaPNM,
     DiffGrad,
     DiffRGrad,
@@ -54,7 +54,7 @@
     SGDP,
     Shampoo,
 ]
-OPTIMIZERS: Dict[str, Optimizer] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
+OPTIMIZERS: Dict[str, Callable] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
 
 def load_optimizer(optimizer: str) -> Callable:

pytorch_optimizer/adan.py

@@ -0,0 +1,121 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base_optimizer import BaseOptimizer
+from pytorch_optimizer.gc import centralize_gradient
+from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Adan(Optimizer, BaseOptimizer):
+    """
+    Reference : x
+    Example :
+        from pytorch_optimizer import Adan
+        ...
+        model = YourModel()
+        optimizer = Adan(model.parameters())
+        ...
+        for input, output in data:
+          optimizer.zero_grad()
+          loss = loss_function(output, model(input))
+          loss.backward()
+          optimizer.step()
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.98, 0.92, 0.99),
+        weight_decay: float = 0.02,
+        use_gc: bool = False,
+        eps: float = 1e-16,
+    ):
+        """Adan optimizer
+        :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 use_gc: bool. use gradient centralization
+        :param eps: float. term added to the denominator to improve numerical stability
+        """
+        self.lr = lr
+        self.betas = betas
+        self.weight_decay = weight_decay
+        self.use_gc = use_gc
+        self.eps = eps
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+        )
+        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)
+        self.validate_epsilon(self.eps)
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            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)
+                state['exp_avg_nest'] = torch.zeros_like(p)
+                state['previous_grad'] = torch.zeros_like(p)
+
+    @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:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError('Adan does not support sparse gradients')
+
+                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)
+                    state['exp_avg_nest'] = torch.zeros_like(p)
+                    state['previous_grad'] = torch.zeros_like(p)
+
+                exp_avg, exp_avg_var, exp_avg_nest = state['exp_avg'], state['exp_avg_var'], state['exp_avg_nest']
+                prev_grad = state['previous_grad']
+
+                state['step'] += 1
+                beta1, beta2, beta3 = group['betas']
+
+                if self.use_gc:
+                    grad = centralize_gradient(grad, gc_conv_only=False)
+
+                grad_diff = grad - prev_grad
+                state['previous_grad'] = grad.clone()
+
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_var.mul_(beta2).add_(grad_diff, alpha=1.0 - beta2)
+                exp_avg_nest.mul_(beta3).add_((grad + beta2 * grad_diff) ** 2, alpha=1.0 - beta3)
+
+                step_size = group['lr'] / exp_avg_nest.add_(self.eps).sqrt_()
+
+                p.sub_(step_size * (exp_avg + beta2 * exp_avg_var))
+                p.div_(1.0 + group['weight_decay'])
+
+        return loss

tests/test_load_optimizers.py

@@ -6,6 +6,7 @@
 
 VALID_OPTIMIZER_NAMES: List[str] = [
     'adamp',
+    'adan',
     'sgdp',
     'madgrad',
     'ranger',
@@ -46,4 +47,4 @@ def test_load_optimizers_invalid(invalid_optimizer_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 17
+    assert len(get_supported_optimizers()) == 18

tests/test_optimizer_parameters.py

@@ -8,6 +8,7 @@
 
 OPTIMIZER_NAMES: List[str] = [
     'adamp',
+    'adan',
     'sgdp',
     'madgrad',
     'ranger',
@@ -37,6 +38,7 @@
     'ralamb',
     'pnm',
     'adapnm',
+    'adan',
 ]
 
 
@@ -122,16 +124,16 @@ def test_betas(optimizer_name):
     with pytest.raises(ValueError):
         if optimizer_name == 'ranger21':
             optimizer(None, num_iterations=100, betas=(-0.1, 0.1))
-        else:
+        elif optimizer not in ('adapnm', 'adan'):
             optimizer(None, betas=(-0.1, 0.1))
 
     with pytest.raises(ValueError):
         if optimizer_name == 'ranger21':
             optimizer(None, num_iterations=100, betas=(0.1, -0.1))
-        else:
+        elif optimizer not in ('adapnm', 'adan'):
             optimizer(None, betas=(0.1, -0.1))
 
-    if optimizer_name == 'adapnm':
+    if optimizer_name in ('adapnm', 'adan'):
         with pytest.raises(ValueError):
             optimizer(None, betas=(0.1, 0.1, -0.1))
 

tests/test_optimizers.py

@@ -14,6 +14,7 @@
     AdaBelief,
     AdaBound,
     AdamP,
+    Adan,
     AdaPNM,
     DiffGrad,
     DiffRGrad,
@@ -79,6 +80,8 @@
     (AdaPNM, {'lr': 3e-1, 'weight_decay': 1e-3, 'amsgrad': False}, 500),
     (Nero, {'lr': 5e-1}, 200),
     (Nero, {'lr': 5e-1, 'constraints': False}, 200),
+    (Adan, {'lr': 2e-1}, 200),
+    (Adan, {'lr': 1e-0, 'weight_decay': 1e-3, 'use_gc': True}, 500),
 ]
 
 ADAMD_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [

tests/test_sparse_gradient.py

@@ -25,6 +25,7 @@
     'lars',
     'shampoo',
     'nero',
+    'adan',
 ]