Merge pull request #108 from kozistr/feature/apollo-optimizer

[Feature] Implement Apollo optimizer

Author: kozistr

README.rst

@@ -248,3 +248,11 @@ AdaFactor
 
 .. autoclass:: pytorch_optimizer.AdaFactor
     :members:
+
+.. _Apollo:
+
+Apollo
+------
+
+.. autoclass:: pytorch_optimizer.Apollo
+    :members:

docs/optimizer_api.rst

@@ -23,6 +23,7 @@
 from pytorch_optimizer.optimizer.adan import Adan
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
 from pytorch_optimizer.optimizer.agc import agc
+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
@@ -94,6 +95,7 @@
     DAdaptSGD,
     AdamS,
     AdaFactor,
+    Apollo,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/__init__.py

@@ -42,6 +42,13 @@ def validate_weight_decay(weight_decay: float):
         if weight_decay < 0.0:
             raise ValueError(f'[-] weight_decay {weight_decay} must be non-negative')
 
+    @staticmethod
+    def validate_weight_decay_type(weight_decay_type: str):
+        if weight_decay_type not in ('l2', 'decoupled', 'stable'):
+            raise ValueError(
+                f'[-] weight_decay_type {weight_decay_type} must be one of (\'l2\', \'decoupled\', \'stable\')'
+            )
+
     @staticmethod
     def validate_weight_decay_ratio(weight_decay_ratio: float):
         if not 0.0 <= weight_decay_ratio < 1.0:
@@ -99,6 +106,11 @@ def validate_norm(norm: float):
         if norm < 0.0:
             raise ValueError(f'[-] norm {norm} must be positive')
 
+    @staticmethod
+    def validate_rebound(rebound: str):
+        if rebound not in ('constant', 'belief'):
+            raise ValueError(f'[-] rebound {rebound} must be one of (\'constant\' or \'belief\')')
+
     @abstractmethod
     def validate_parameters(self):
         raise NotImplementedError

pytorch_optimizer/base/optimizer.py

@@ -0,0 +1,157 @@
+from typing import Optional
+
+import numpy as np
+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 CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class Apollo(Optimizer, BaseOptimizer):
+    r"""An Adaptive Parameter-wise Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization.
+
+    :param params: PARAMETERS. iterable of parameters to optimize or dicts defining parameter groups.
+    :param lr: float. learning rate.
+    :param init_lr: Optional[float]. initial learning rate (default lr / 1000).
+    :param beta: float. coefficient used for computing running averages of gradient.
+    :param rebound: str. rectified bound for diagonal hessian. (constant, belief).
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decay_type: str. type of weight decay. (l2, decoupled, stable).
+    :param warmup_steps: int. number of warmup steps.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        init_lr: Optional[float] = None,
+        beta: float = 0.9,
+        rebound: str = 'constant',
+        weight_decay: float = 0.0,
+        weight_decay_type: str = 'l2',
+        warmup_steps: int = 500,
+        eps: float = 1e-4,
+    ):
+        self.lr = lr
+        self.beta = beta
+        self.rebound = rebound
+        self.weight_decay = weight_decay
+        self.weight_decay_type = weight_decay_type
+        self.warmup_steps = warmup_steps
+        self.eps = eps
+
+        self.validate_parameters()
+
+        self.init_lr: float = init_lr if init_lr is not None else lr / 1000.0
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'init_lr': self.init_lr,
+            'beta': beta,
+            'weight_decay': weight_decay,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def validate_parameters(self):
+        self.validate_learning_rate(self.lr)
+        self.validate_beta(self.beta)
+        self.validate_rebound(self.rebound)
+        self.validate_weight_decay(self.weight_decay)
+        self.validate_weight_decay_type(self.weight_decay_type)
+        self.validate_epsilon(self.eps)
+
+    @property
+    def __str__(self) -> str:
+        return 'Apollo'
+
+    @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_grad'] = torch.zeros_like(p)
+                state['approx_hessian'] = torch.zeros_like(p)
+                state['update'] = 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:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            current_lr: float = (
+                group['lr']
+                if group['step'] >= self.warmup_steps
+                else (self.lr - group['init_lr']) * group['step'] / self.warmup_steps + group['init_lr']
+            )
+
+            weight_decay, eps = group['weight_decay'], group['eps']
+
+            bias_correction: float = 1.0 - group['beta'] ** group['step']
+            alpha: float = (1.0 - group['beta']) / bias_correction
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(self.__str__)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['exp_avg_grad'] = torch.zeros_like(p)
+                    state['approx_hessian'] = torch.zeros_like(p)
+                    state['update'] = torch.zeros_like(p)
+
+                exp_avg_grad, b, d_p = state['exp_avg_grad'], state['approx_hessian'], state['update']
+
+                if weight_decay > 0.0 and self.weight_decay_type == 'l2':
+                    grad.add_(p, alpha=weight_decay)
+
+                delta_grad = grad - exp_avg_grad
+                if self.rebound == 'belief':
+                    rebound = delta_grad.norm(p=np.inf)
+                else:
+                    rebound = 1e-2
+                    eps /= rebound
+
+                exp_avg_grad.add_(delta_grad, alpha=alpha)
+
+                de_nom = d_p.norm(p=4).add(eps)
+                d_p.div_(de_nom)
+
+                v_sq = d_p.mul(d_p)
+                delta = delta_grad.div_(de_nom).mul_(d_p).sum().mul(-alpha) - b.mul(v_sq).sum()
+
+                b.addcmul_(v_sq, delta)
+
+                de_nom = b.abs().clamp_(min=rebound)
+                if self.rebound == 'belief':
+                    de_nom.add_(eps / alpha)
+
+                d_p.copy_(exp_avg_grad.div(de_nom))
+
+                if weight_decay > 0.0 and self.weight_decay_type != 'l2':
+                    if self.weight_decay_type == 'stable':
+                        weight_decay /= de_nom.mean().item()
+
+                    d_p.add_(p, alpha=weight_decay)
+
+                p.add_(d_p, alpha=-current_lr)
+
+        return loss

pytorch_optimizer/optimizer/apollo.py

@@ -89,12 +89,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
         k = self.state['k']
 
         for group in self.param_groups:
-            eps = group['eps']
+            weight_decay, momentum, eps = group['weight_decay'], group['momentum'], group['eps']
             lr = group['lr'] + eps
-            weight_decay = group['weight_decay']
-            momentum = group['momentum']
 
-            ck: float = 1.0 - momentum
             _lambda = lr * math.pow(k + 1, 0.5)
 
             for p in group['params']:
@@ -113,8 +110,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if momentum > 0.0 and grad.is_sparse:
                     raise NoSparseGradientError(self.__str__, note='momentum > 0.0')
 
-                grad_sum_sq = state['grad_sum_sq']
-                s = state['s']
+                grad_sum_sq, s = state['grad_sum_sq'], state['s']
 
                 if weight_decay > 0.0 and not self.decouple_decay:
                     if grad.is_sparse:
@@ -176,7 +172,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                         p.copy_(x0.addcdiv(s, rms, value=-1))
                     else:
                         z = x0.addcdiv(s, rms, value=-1)
-                        p.mul_(1.0 - ck).add_(z, alpha=ck)
+                        p.mul_(momentum).add_(z, alpha=1.0 - momentum)
 
                     if weight_decay > 0.0 and self.decouple_decay:
                         p.add_(p_old, alpha=-lr * weight_decay)

pytorch_optimizer/optimizer/madgrad.py

@@ -3,6 +3,7 @@
 from pytorch_optimizer import (
     LARS,
     MADGRAD,
+    OPTIMIZERS,
     PNM,
     SGDP,
     AdaBelief,
@@ -13,6 +14,7 @@
     AdamS,
     Adan,
     AdaPNM,
+    Apollo,
     DAdaptAdaGrad,
     DAdaptAdam,
     DAdaptSGD,
@@ -32,63 +34,7 @@
 ADAPTIVE_FLAGS: List[bool] = [True, False]
 PULLBACK_MOMENTUM: List[str] = ['none', 'reset', 'pullback']
 
-SPARSE_OPTIMIZERS: List[str] = [
-    'madgrad',
-    'dadaptadagrad',
-]
-NO_SPARSE_OPTIMIZERS: List[str] = [
-    'adamp',
-    'sgdp',
-    'madgrad',
-    'ranger',
-    'ranger21',
-    'radam',
-    'adabound',
-    'adabelief',
-    'diffgrad',
-    'diffrgrad',
-    'lamb',
-    'ralamb',
-    'lars',
-    'shampoo',
-    'scalableshampoo',
-    'nero',
-    'adan',
-    'adai',
-    'adapnm',
-    'pnm',
-    'dadaptadam',
-    'dadaptsgd',
-    'adams',
-    'adafactor',
-]
-VALID_OPTIMIZER_NAMES: List[str] = [
-    'adamp',
-    'adan',
-    'sgdp',
-    'madgrad',
-    'ranger',
-    'ranger21',
-    'radam',
-    'adabound',
-    'adabelief',
-    'diffgrad',
-    'diffrgrad',
-    'lamb',
-    'ralamb',
-    'lars',
-    'shampoo',
-    'scalableshampoo',
-    'pnm',
-    'adapnm',
-    'nero',
-    'adai',
-    'dadaptadagrad',
-    'dadaptadam',
-    'dadaptsgd',
-    'adams',
-    'adafactor',
-]
+VALID_OPTIMIZER_NAMES: List[str] = list(OPTIMIZERS.keys())
 INVALID_OPTIMIZER_NAMES: List[str] = [
     'asam',
     'sam',
@@ -97,6 +43,12 @@
     'adamd',
     'lookahead',
 ]
+
+SPARSE_OPTIMIZERS: List[str] = ['madgrad', 'dadaptadagrad']
+NO_SPARSE_OPTIMIZERS: List[str] = [
+    optimizer for optimizer in VALID_OPTIMIZER_NAMES if optimizer not in SPARSE_OPTIMIZERS
+]
+
 BETA_OPTIMIZER_NAMES: List[str] = [
     'adabelief',
     'adabound',
@@ -126,9 +78,7 @@
     'CyclicLR',
     'OneCycleLR',
 ]
-INVALID_LR_SCHEDULER_NAMES: List[str] = [
-    'dummy',
-]
+INVALID_LR_SCHEDULER_NAMES: List[str] = ['dummy']
 
 OPTIMIZERS: List[Tuple[Any, Dict[str, Union[float, bool, int]], int]] = [
     (build_lookahead, {'lr': 5e-1, 'weight_decay': 1e-3}, 10),
@@ -205,6 +155,9 @@
     (AdamS, {'lr': 1.0, 'weight_decay': 1e-3}, 30),
     (AdamS, {'lr': 1.0, 'weight_decay': 1e-3, 'amsgrad': True}, 30),
     (AdaFactor, {'lr': 5e-1, 'weight_decay': 1e-2, 'scale_parameter': False}, 100),
+    (Apollo, {'lr': 5e-1, 'weight_decay': 1e-3}, 10),
+    (Apollo, {'lr': 5e-1, 'weight_decay': 1e-3, 'rebound': 'belief'}, 10),
+    (Apollo, {'lr': 5e-1, 'weight_decay': 1e-3, 'weight_decay_type': 'stable', 'warmup_steps': 0}, 50),
 ]
 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}, 10),

tests/constants.py

@@ -59,7 +59,15 @@ def test_sparse_supported(sparse_optimizer):
     optimizer.zero_grad()
     optimizer.step()
 
-    optimizer = opt([param], momentum=0.0, weight_decay=1e-3)
+    if sparse_optimizer == 'madgrad':
+        optimizer = opt([param], momentum=0.0, weight_decay=1e-3, decouple_decay=False)
+        optimizer.reset()
+        optimizer.zero_grad()
+
+        with pytest.raises(NoSparseGradientError):
+            optimizer.step()
+
+    optimizer = opt([param], momentum=0.9, weight_decay=1e-3)
     optimizer.reset()
     optimizer.zero_grad()
 

tests/test_gradients.py

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