Merge pull request #89 from kozistr/feature/adai-optimizer

[Feature] Implement `Adai` optimizer

Author: kozistr

README.rst

@@ -110,6 +110,8 @@ Supported Optimizers
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | 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>`__                       |
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| Adai         | *Disentangling the Effects of Adaptive Learning Rate and Momentum*                     | `github <https://github.com/zeke-xie/adaptive-inertia-adai>`__                    | `https://arxiv.org/abs/2006.15815 <https://arxiv.org/abs/2006.15815>`__                       |
++--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -299,6 +301,8 @@ Citations
 
 `Adan <https://ui.adsabs.harvard.edu/abs/2022arXiv220806677X/exportcitation>`__
 
+`Adai <https://github.com/zeke-xie/adaptive-inertia-adai#citing>`__
+
 Author
 ------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.0.1"
+version = "2.1.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

@@ -6,6 +6,7 @@
 from pytorch_optimizer.lr_scheduler.cosine_anealing import CosineAnnealingWarmupRestarts
 from pytorch_optimizer.optimizer.adabelief import AdaBelief
 from pytorch_optimizer.optimizer.adabound import AdaBound
+from pytorch_optimizer.optimizer.adai import Adai
 from pytorch_optimizer.optimizer.adamp import AdamP
 from pytorch_optimizer.optimizer.adan import Adan
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
@@ -40,6 +41,7 @@
     AdaBelief,
     AdaBound,
     AdamP,
+    Adai,
     Adan,
     AdaPNM,
     DiffGrad,

pytorch_optimizer/optimizer/adai.py

@@ -0,0 +1,160 @@
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.base_optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Adai(Optimizer, BaseOptimizer):
+    """
+    Reference : https://github.com/zeke-xie/adaptive-inertia-adai
+    Example :
+        from pytorch_optimizer import Adai
+        ...
+        model = YourModel()
+        optimizer = Adai(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.1, 0.99),
+        weight_decay: float = 0.0,
+        weight_decouple: bool = False,
+        dampening: float = 1.0,
+        eps: float = 1e-3,
+    ):
+        """Adai
+        :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
+        :param dampening: float. dampening for momentum. where dampening < 1,
+            it will show some adaptive-moment behavior
+        :param eps: float. term added to the denominator to improve numerical stability
+        """
+        self.lr = lr
+        self.betas = betas
+        self.weight_decay = weight_decay
+        self.weight_decouple = weight_decouple
+        self.dampening = dampening
+        self.eps = eps
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = dict(
+            lr=lr,
+            betas=betas,
+            weight_decay=weight_decay,
+            dampening=dampening,
+            eps=eps,
+        )
+        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_sq'] = torch.zeros_like(p)
+                state['beta1_prod'] = torch.ones_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()
+
+        param_size: int = 0
+        exp_avg_sq_hat_sum: float = 0.0
+
+        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('Adai does not support sparse gradients')
+
+                param_size += p.numel()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['beta1_prod'] = torch.ones_like(p)
+
+                state['step'] += 1
+
+                exp_avg_sq = state['exp_avg_sq']
+                _, beta2 = group['betas']
+
+                bias_correction2 = 1.0 - beta2 ** state['step']
+
+                if group['weight_decay'] != 0:
+                    if self.weight_decouple:
+                        p.mul_(1.0 - group['lr'] * group['weight_decay'])
+                    else:
+                        grad.add_(p, alpha=group['weight_decay'])
+
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                exp_avg_sq_hat_sum += exp_avg_sq.sum() / bias_correction2
+
+        if param_size == 0:
+            raise ValueError('[-] param_size is 0')
+
+        exp_avg_sq_hat_mean = exp_avg_sq_hat_sum / param_size
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                state = self.state[p]
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1_prod = state['beta1_prod']
+                beta0, beta2 = group['betas']
+
+                bias_correction2 = 1 - beta2 ** state['step']
+
+                exp_avg_sq_hat = exp_avg_sq / bias_correction2
+                beta1 = (
+                    1.0 - (exp_avg_sq_hat / exp_avg_sq_hat_mean).pow(1.0 / (3 - 2 * group['dampening'])).mul(beta0)
+                ).clamp(0.0, 1 - group['eps'])
+                beta3 = (1.0 - beta1).pow(group['dampening'])
+
+                beta1_prod.mul_(beta1)
+                bias_correction1 = 1.0 - beta1_prod
+
+                exp_avg.mul_(beta1).addcmul_(beta3, grad)
+                exp_avg_hat = exp_avg / bias_correction1 * math.pow(beta0, 1.0 - group['dampening'])
+
+                p.add_(exp_avg_hat, alpha=-group['lr'])
+
+        return loss

pytorch_optimizer/optimizer/ranger21.py

@@ -236,7 +236,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         # stable weight decay
         if param_size == 0:
-            raise ValueError('[-] size of parameter is 0')
+            raise ValueError('[-] param_size is 0')
 
         variance_normalized = math.sqrt(variance_ma_sum / param_size)
         if math.isnan(variance_normalized):

tests/test_load_optimizers.py

@@ -23,8 +23,8 @@
     'pnm',
     'adapnm',
     'nero',
+    'adai',
 ]
-
 INVALID_OPTIMIZER_NAMES: List[str] = [
     'asam',
     'sam',
@@ -47,4 +47,4 @@ def test_load_optimizers_invalid(invalid_optimizer_names):
 
 
 def test_get_supported_optimizers():
-    assert len(get_supported_optimizers()) == 18
+    assert len(get_supported_optimizers()) == 19

tests/test_optimizer_parameters.py

@@ -5,7 +5,7 @@
 from torch import nn
 from torch.nn import functional as F
 
-from pytorch_optimizer import SAM, AdamP, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizer
+from pytorch_optimizer import SAM, Adai, AdamP, Lookahead, PCGrad, Ranger21, SafeFP16Optimizer, load_optimizer
 from tests.utils import Example
 
 OPTIMIZER_NAMES: List[str] = [
@@ -25,8 +25,8 @@
     'lars',
     'pnm',
     'adapnm',
+    'adai',
 ]
-
 BETA_OPTIMIZER_NAMES: List[str] = [
     'adabelief',
     'adabound',
@@ -41,6 +41,7 @@
     'pnm',
     'adapnm',
     'adan',
+    'adai',
 ]
 
 
@@ -213,12 +214,13 @@ def test_ranger21_warm_methods():
     assert Ranger21.build_warm_down_iterations(1000) == 280
 
 
-def test_ranger21_size_of_parameter():
+@pytest.mark.parametrize('optimizer', [Ranger21, Adai])
+def test_size_of_parameter(optimizer):
     model: nn.Module = nn.Linear(1, 1, bias=False)
     model.requires_grad_(False)
 
     with pytest.raises(ValueError):
-        Ranger21(model.parameters(), 100).step()
+        optimizer(model.parameters(), 100).step()
 
 
 def test_ranger21_closure():

tests/test_optimizers.py

@@ -13,6 +13,7 @@
     SGDP,
     AdaBelief,
     AdaBound,
+    Adai,
     AdamP,
     Adan,
     AdaPNM,
@@ -36,6 +37,7 @@
     dummy_closure,
     ids,
     make_dataset,
+    names,
     tensor_to_numpy,
 )
 
@@ -50,6 +52,10 @@
     (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'fixed_decay': True}, 100),
     (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'weight_decouple': False}, 100),
     (AdaBound, {'lr': 5e-1, 'gamma': 0.1, 'weight_decay': 1e-3, 'amsbound': True}, 100),
+    (Adai, {'lr': 1e-1, 'weight_decay': 0.0}, 200),
+    (Adai, {'lr': 1e-1, 'weight_decay': 0.0, 'dampening': 0.9}, 200),
+    (Adai, {'lr': 1e-1, 'weight_decay': 1e-4, 'weight_decouple': False}, 200),
+    (Adai, {'lr': 1e-1, 'weight_decay': 1e-4, 'weight_decouple': True}, 200),
     (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3}, 100),
     (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3, 'use_gc': True}, 100),
     (AdamP, {'lr': 5e-1, 'weight_decay': 1e-3, 'nesterov': True}, 100),
@@ -84,7 +90,6 @@
     (Adan, {'lr': 1e-0, 'weight_decay': 1e-3, 'use_gc': True}, 100),
     (Adan, {'lr': 1e-0, 'weight_decay': 1e-3, 'use_gc': True, 'weight_decouple': True}, 100),
 ]
-
 ADAMD_SUPPORTED_OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (build_lookahead, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 100),
     (AdaBelief, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 100),
@@ -167,6 +172,7 @@ def test_safe_f16_optimizers(optimizer_fp16_config):
         or (optimizer_name == 'Nero')
         or (optimizer_name == 'Adan' and 'weight_decay' not in config)
         or (optimizer_name == 'RAdam')
+        or (optimizer_name == 'Adai')
     ):
         pytest.skip(f'skip {optimizer_name}')
 
@@ -195,8 +201,10 @@ def test_sam_optimizers(adaptive, optimizer_sam_config):
     (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_sam_config
-    if optimizer_class.__name__ == 'Shampoo':
-        pytest.skip(f'skip {optimizer_class.__name__}')
+
+    optimizer_name: str = optimizer_class.__name__
+    if (optimizer_name == 'Shampoo') or (optimizer_name == 'Adai'):
+        pytest.skip(f'skip {optimizer_name}')
 
     optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
@@ -221,8 +229,10 @@ def test_sam_optimizers_with_closure(adaptive, optimizer_sam_config):
     (x_data, y_data), model, loss_fn = build_environment()
 
     optimizer_class, config, iterations = optimizer_sam_config
-    if optimizer_class.__name__ == 'Shampoo':
-        pytest.skip(f'skip {optimizer_class.__name__}')
+
+    optimizer_name: str = optimizer_class.__name__
+    if (optimizer_name == 'Shampoo') or (optimizer_name == 'Adai'):
+        pytest.skip(f'skip {optimizer_name}')
 
     optimizer = SAM(model.parameters(), optimizer_class, **config, adaptive=adaptive)
 
@@ -335,26 +345,31 @@ def test_no_gradients(optimizer_config):
     assert tensor_to_numpy(init_loss) >= tensor_to_numpy(loss)
 
 
-@pytest.mark.parametrize('optimizer_config', OPTIMIZERS, ids=ids)
-def test_closure(optimizer_config):
+@pytest.mark.parametrize('optimizer', set(config[0] for config in OPTIMIZERS), ids=names)
+def test_closure(optimizer):
     _, model, _ = build_environment()
 
-    optimizer_class, config, _ = optimizer_config
-    if optimizer_class.__name__ == 'Ranger21':
-        pytest.skip(f'skip {optimizer_class.__name__}')
-
-    optimizer = optimizer_class(model.parameters(), **config)
+    if optimizer.__name__ == 'Ranger21':
+        optimizer = optimizer(model.parameters(), num_iterations=1)
+    else:
+        optimizer = optimizer(model.parameters())
 
     optimizer.zero_grad()
-    optimizer.step(closure=dummy_closure)
+
+    try:
+        optimizer.step(closure=dummy_closure)
+    except ValueError:  # in case of Ranger21, Adai optimizers
+        pass
 
 
 @pytest.mark.parametrize('optimizer_config', OPTIMIZERS, ids=ids)
 def test_reset(optimizer_config):
     _, model, _ = build_environment()
 
     optimizer_class, config, _ = optimizer_config
-    optimizer = optimizer_class(model.parameters(), **config)
+    if optimizer_class.__name__ == 'Ranger21':
+        config.update({'num_iterations': 1})
 
+    optimizer = optimizer_class(model.parameters(), **config)
     optimizer.zero_grad()
     optimizer.reset()

tests/test_sparse_gradient.py

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

tests/utils.py

@@ -78,6 +78,10 @@ def ids(v) -> str:
     return f'{v[0].__name__}_{v[1:]}'
 
 
+def names(v) -> str:
+    return v.__name__
+
+
 def build_environment(use_gpu: bool = False) -> Tuple[Tuple[torch.Tensor, torch.Tensor], nn.Module, nn.Module]:
     torch.manual_seed(42)