Merge pull request #109 from kozistr/feature/novograd-optimizer

[Feature] Implement NovoGrad optimizer

Author: kozistr

README.rst

@@ -120,6 +120,8 @@ Supported Optimizers
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | Apollo       | *An Adaptive Parameter-wise Diagonal Quasi-Newton Method for Nonconvex Stochastic Optimization* | `github <https://github.com/XuezheMax/apollo>`__                                  | `https://arxiv.org/abs/2009.13586 <https://arxiv.org/abs/2009.13586>`__                       |
 +--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| NovoGrad     | *Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks*    | `github <https://github.com/lonePatient/NovoGrad-pytorch>`__                      | `https://arxiv.org/abs/1905.11286 <https://arxiv.org/abs/1905.11286>`__                       |
++--------------+-------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -319,6 +321,8 @@ Citations
 
 `Apollo <https://ui.adsabs.harvard.edu/abs/2020arXiv200913586M/exportcitation>`__
 
+`NovoGrad <https://ui.adsabs.harvard.edu/abs/2019arXiv190511286G/exportcitation>`__
+
 Citation
 --------
 

docs/optimizer_api.rst

@@ -256,3 +256,11 @@ Apollo
 
 .. autoclass:: pytorch_optimizer.Apollo
     :members:
+
+.. _NovoGrad:
+
+NovoGrad
+--------
+
+.. autoclass:: pytorch_optimizer.NovoGrad
+    :members:

pytorch_optimizer/__init__.py

@@ -35,6 +35,7 @@
 from pytorch_optimizer.optimizer.lookahead import Lookahead
 from pytorch_optimizer.optimizer.madgrad import MADGRAD
 from pytorch_optimizer.optimizer.nero import Nero
+from pytorch_optimizer.optimizer.novograd import NovoGrad
 from pytorch_optimizer.optimizer.pcgrad import PCGrad
 from pytorch_optimizer.optimizer.pnm import PNM
 from pytorch_optimizer.optimizer.radam import RAdam
@@ -96,6 +97,7 @@
     AdamS,
     AdaFactor,
     Apollo,
+    NovoGrad,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/optimizer/novograd.py

@@ -0,0 +1,129 @@
+import math
+
+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 NovoGrad(Optimizer, BaseOptimizer):
+    r"""Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks.
+
+    :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 grad_averaging: bool. multiply ck (1 - momentum).
+    :param adamd_debias_term: 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.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.95, 0.98),
+        weight_decay: float = 0.0,
+        grad_averaging: bool = False,
+        adamd_debias_term: bool = False,
+        eps: float = 1e-8,
+    ):
+        self.lr = lr
+        self.betas = betas
+        self.weight_decay = weight_decay
+        self.grad_averaging = grad_averaging
+        self.adamd_debias_term = adamd_debias_term
+        self.eps = eps
+
+        self.validate_parameters()
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            '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)
+
+    @property
+    def __str__(self) -> str:
+        return 'NovoGrad'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                grad = p.grad
+                g_2 = grad ** 2  # fmt: skip
+
+                state['step'] = 0
+                state['moments'] = grad.div(g_2.sqrt() + group['eps']) + group['weight_decay'] * p
+                state['grads_ema'] = g_2
+
+    @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
+
+            beta1, beta2 = group['betas']
+            weight_decay = group['weight_decay']
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2_sq = math.sqrt(1.0 - beta2 ** group['step'])
+
+            step_size: float = group['lr'] * bias_correction2_sq
+            if not self.adamd_debias_term:
+                step_size /= bias_correction1
+
+            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]
+                g_2 = grad ** 2  # fmt: skip
+
+                if len(state) == 0:
+                    state['moments'] = grad.div(g_2.sqrt() + group['eps']) + weight_decay * p
+                    state['grads_ema'] = g_2
+
+                moments, grads_ema = state['moments'], state['grads_ema']
+
+                grads_ema.mul_(beta2).add_(g_2, alpha=1.0 - beta2)
+
+                de_nom = grads_ema.sqrt().add_(group['eps'])
+                grad.div_(de_nom)
+
+                if weight_decay > 0.0:
+                    grad.add_(p, alpha=weight_decay)
+
+                if self.grad_averaging:
+                    grad.mul_(1.0 - beta1)
+
+                moments.mul_(beta1).add_(grad)
+
+                p.add_(moments, alpha=-step_size)
+
+        return loss

tests/constants.py

@@ -22,6 +22,7 @@
     DiffRGrad,
     Lamb,
     Nero,
+    NovoGrad,
     RAdam,
     RaLamb,
     Ranger,
@@ -68,6 +69,7 @@
     'dadaptadam',
     'adams',
     'adafactor',
+    'novograd',
 ]
 
 VALID_LR_SCHEDULER_NAMES: List[str] = [
@@ -158,6 +160,7 @@
     (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),
+    (NovoGrad, {'lr': 5e-1, 'weight_decay': 1e-3, 'grad_averaging': True}, 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),
@@ -172,4 +175,5 @@
     (Ranger21, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True, 'num_iterations': 200}, 200),
     (AdaPNM, {'lr': 3e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 50),
     (AdamS, {'lr': 2e1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 50),
+    (NovoGrad, {'lr': 5e-1, 'weight_decay': 1e-3, 'adamd_debias_term': True}, 50),
 ]

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()) == 26
+    assert len(get_supported_optimizers()) == 27