Add SGDP optimizer (#145)

yohann84L · dependabot-preview[bot] · web-flow · commit 9a9d233dfb4d · 2020-07-09T13:49:17.000-04:00
* add adamp * update test for adamp * update doc for adamp * update readme with adamp example * update method to staticmethod * add nesterov for AdamP * update viz and docs for AdamP * add check for delta & wd_ratio values * add SGDP optimizer * update tests for SGDP * update doc & readme with SGDP * Bump matplotlib from 3.2.1 to 3.2.2 Bumps [matplotlib](https://github.com/matplotlib/matplotlib) from 3.2.1 to 3.2.2. - [Release notes](https://github.com/matplotlib/matplotlib/releases) - [Commits](matplotlib/matplotlib@v3.2.1...v3.2.2) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump torch from 1.5.0 to 1.5.1 Bumps [torch](https://github.com/pytorch/pytorch) from 1.5.0 to 1.5.1. - [Release notes](https://github.com/pytorch/pytorch/releases) - [Commits](pytorch/pytorch@v1.5.0...v1.5.1) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump torchvision from 0.6.0 to 0.6.1 Bumps [torchvision](https://github.com/pytorch/vision) from 0.6.0 to 0.6.1. - [Release notes](https://github.com/pytorch/vision/releases) - [Commits](https://github.com/pytorch/vision/commits) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump mypy from 0.780 to 0.781 Bumps [mypy](https://github.com/python/mypy) from 0.780 to 0.781. - [Release notes](https://github.com/python/mypy/releases) - [Commits](python/mypy@v0.780...v0.781) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump sphinx-autodoc-typehints from 1.10.3 to 1.11.0 Bumps [sphinx-autodoc-typehints](https://github.com/agronholm/sphinx-autodoc-typehints) from 1.10.3 to 1.11.0. - [Release notes](https://github.com/agronholm/sphinx-autodoc-typehints/releases) - [Changelog](https://github.com/agronholm/sphinx-autodoc-typehints/blob/master/CHANGELOG.rst) - [Commits](agronholm/sphinx-autodoc-typehints@1.10.3...1.11.0) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump mypy from 0.781 to 0.782 Bumps [mypy](https://github.com/python/mypy) from 0.781 to 0.782. - [Release notes](https://github.com/python/mypy/releases) - [Commits](python/mypy@v0.781...v0.782) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump ipdb from 0.13.2 to 0.13.3 Bumps [ipdb](https://github.com/gotcha/ipdb) from 0.13.2 to 0.13.3. - [Release notes](https://github.com/gotcha/ipdb/releases) - [Changelog](https://github.com/gotcha/ipdb/blob/master/HISTORY.txt) - [Commits](gotcha/ipdb@0.13.2...0.13.3) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump ipython from 7.15.0 to 7.16.1 Bumps [ipython](https://github.com/ipython/ipython) from 7.15.0 to 7.16.1. - [Release notes](https://github.com/ipython/ipython/releases) - [Commits](ipython/ipython@7.15.0...7.16.1) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump numpy from 1.18.5 to 1.19.0 Bumps [numpy](https://github.com/numpy/numpy) from 1.18.5 to 1.19.0. - [Release notes](https://github.com/numpy/numpy/releases) - [Changelog](https://github.com/numpy/numpy/blob/master/doc/HOWTO_RELEASE.rst.txt) - [Commits](numpy/numpy@v1.18.5...v1.19.0) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump matplotlib from 3.2.1 to 3.2.2 Bumps [matplotlib](https://github.com/matplotlib/matplotlib) from 3.2.1 to 3.2.2. - [Release notes](https://github.com/matplotlib/matplotlib/releases) - [Commits](matplotlib/matplotlib@v3.2.1...v3.2.2) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump torch from 1.5.0 to 1.5.1 Bumps [torch](https://github.com/pytorch/pytorch) from 1.5.0 to 1.5.1. - [Release notes](https://github.com/pytorch/pytorch/releases) - [Commits](pytorch/pytorch@v1.5.0...v1.5.1) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump torchvision from 0.6.0 to 0.6.1 Bumps [torchvision](https://github.com/pytorch/vision) from 0.6.0 to 0.6.1. - [Release notes](https://github.com/pytorch/vision/releases) - [Commits](https://github.com/pytorch/vision/commits) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump mypy from 0.780 to 0.781 Bumps [mypy](https://github.com/python/mypy) from 0.780 to 0.781. - [Release notes](https://github.com/python/mypy/releases) - [Commits](python/mypy@v0.780...v0.781) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump sphinx-autodoc-typehints from 1.10.3 to 1.11.0 Bumps [sphinx-autodoc-typehints](https://github.com/agronholm/sphinx-autodoc-typehints) from 1.10.3 to 1.11.0. - [Release notes](https://github.com/agronholm/sphinx-autodoc-typehints/releases) - [Changelog](https://github.com/agronholm/sphinx-autodoc-typehints/blob/master/CHANGELOG.rst) - [Commits](agronholm/sphinx-autodoc-typehints@1.10.3...1.11.0) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump mypy from 0.781 to 0.782 Bumps [mypy](https://github.com/python/mypy) from 0.781 to 0.782. - [Release notes](https://github.com/python/mypy/releases) - [Commits](python/mypy@v0.781...v0.782) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump ipdb from 0.13.2 to 0.13.3 Bumps [ipdb](https://github.com/gotcha/ipdb) from 0.13.2 to 0.13.3. - [Release notes](https://github.com/gotcha/ipdb/releases) - [Changelog](https://github.com/gotcha/ipdb/blob/master/HISTORY.txt) - [Commits](gotcha/ipdb@0.13.2...0.13.3) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump ipython from 7.15.0 to 7.16.1 Bumps [ipython](https://github.com/ipython/ipython) from 7.15.0 to 7.16.1. - [Release notes](https://github.com/ipython/ipython/releases) - [Commits](ipython/ipython@7.15.0...7.16.1) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> * Bump numpy from 1.18.5 to 1.19.0 Bumps [numpy](https://github.com/numpy/numpy) from 1.18.5 to 1.19.0. - [Release notes](https://github.com/numpy/numpy/releases) - [Changelog](https://github.com/numpy/numpy/blob/master/doc/HOWTO_RELEASE.rst.txt) - [Commits](numpy/numpy@v1.18.5...v1.19.0) Signed-off-by: dependabot-preview[bot] <support@dependabot.com> Co-authored-by: dependabot-preview[bot] <27856297+dependabot-preview[bot]@users.noreply.github.com>
diff --git a/README.rst b/README.rst
@@ -230,7 +230,7 @@ Intuitively, this operation prevents the unnecessary update along the radial dir
 that only increases the weight norm without contributing to the loss minimization.
 
 +------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+
-| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdamP.png    |  .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdamP.png   |
+| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_AdamP.png     |  .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_AdamP.png    |
 +------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------+
 
 .. code:: python
@@ -577,6 +577,36 @@ RangerVA
 **Reference Code**: https://github.com/lessw2020/Ranger-Deep-Learning-Optimizer
 
 
+SGDP
+----
+
++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+
+| .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rastrigin_SGDP.png  |  .. image:: https://raw.githubusercontent.com/jettify/pytorch-optimizer/master/docs/rosenbrock_SGDP.png  |
++--------------------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------+
+
+.. code:: python
+
+    import torch_optimizer as optim
+
+    # model = ...
+    optimizer = optim.SGDP(
+        m.parameters(),
+        lr= 1e-3,
+        momentum=0,
+        dampening=0,
+        weight_decay=1e-2,
+        nesterov=False,
+        delta = 0.1,
+        wd_ratio = 0.1
+    )
+    optimizer.step()
+
+
+**Paper**: *Slowing Down the Weight Norm Increase in Momentum-based Optimizers.* (2020) [https://arxiv.org/abs/2006.08217]
+
+**Reference Code**: https://github.com/clovaai/AdamP
+
+
 SGDW
 ----
 
diff --git a/docs/api.rst b/docs/api.rst
@@ -89,6 +89,14 @@ RAdam
 .. autoclass:: torch_optimizer.RAdam
     :members:
 
+.. _SGDP:
+
+SGDP
+----
+
+.. autoclass:: torch_optimizer.SGDP
+    :members:
+
 .. _SGDW:
 
 SGDW
diff --git a/docs/index.rst b/docs/index.rst
@@ -74,6 +74,9 @@ Supported Optimizers
 | :ref:`RangerVA` | https://arxiv.org/abs/1908.00700v2                                            |
 +-----------------+-------------------------------------------------------------------------------+
 |                 |                                                                               |
+| :ref:`SGDP`     | https://arxiv.org/abs/2006.08217                                               |
++-----------------+-------------------------------------------------------------------------------+
+|                 |                                                                               |
 | :ref:`SGDW`     | https://arxiv.org/abs/1608.03983                                              |
 +-----------------+-------------------------------------------------------------------------------+
 |                 |                                                                               |
diff --git a/docs/rastrigin_AdamP.png b/docs/rastrigin_AdamP.png
diff --git a/docs/rastrigin_SGDP.png b/docs/rastrigin_SGDP.png
diff --git a/docs/rosenbrock_AdamP.png b/docs/rosenbrock_AdamP.png
diff --git a/docs/rosenbrock_SGDP.png b/docs/rosenbrock_SGDP.png
diff --git a/examples/viz_optimizers.py b/examples/viz_optimizers.py
@@ -166,6 +166,7 @@ def LookaheadYogi(*a, **kw):
         # Adam based
         (optim.AdaBound, -8, 0.3),
         (optim.AdaMod, -8, 0.2),
+        (optim.AdamP, -8, 0.2),
         (optim.DiffGrad, -8, 0.4),
         (optim.Lamb, -8, -2.9),
         (optim.NovoGrad, -8, -1.7),
@@ -174,6 +175,7 @@ def LookaheadYogi(*a, **kw):
         # SGD/Momentum based
         (optim.AccSGD, -8, -1.4),
         (optim.SGDW, -8, -1.5),
+        (optim.SGDP, -8, -1.5),
         (optim.PID, -8, -1.0),
         (optim.QHM, -6, -0.2),
         (optim.QHAdam, -8, 0.1),
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -60,7 +60,8 @@ def build_lookahead(*a, **kw):
     (optim.AccSGD, {'lr': 0.015}, 800),
     (build_lookahead, {'lr': 1.0}, 500),
     (optim.QHAdam, {'lr': 1.0}, 500),
-    (optim.AdamP, {'lr': 0.01, 'betas': (0.9, 0.95), 'eps': 1e-3}, 800)
+    (optim.AdamP, {'lr': 0.01, 'betas': (0.9, 0.95), 'eps': 1e-3}, 800),
+    (optim.SGDP, {'lr': 0.002, 'momentum': 0.91}, 900),
 ]
 
 
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -77,6 +77,7 @@ def build_lookahead(*a, **kw):
     optim.QHM,
     optim.RAdam,
     optim.SGDW,
+    optim.SGDP,
     optim.Yogi,
     build_lookahead,
     optim.Ranger,
diff --git a/tests/test_optimizer_with_nn.py b/tests/test_optimizer_with_nn.py
@@ -68,7 +68,8 @@ def build_lookahead(*a, **kw):
     (optim.RangerQH, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
     (optim.RangerVA, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
     (optim.Shampoo, {'lr': 0.1, 'weight_decay': 1e-3, 'momentum': 0.8}, 200),
-    (optim.AdamP, {'lr': 1.0, 'weight_decay': 1e-3}, 200)
+    (optim.AdamP, {'lr': 1.0, 'weight_decay': 1e-3}, 200),
+    (optim.SGDP, {'lr': 1.0, 'weight_decay': 1e-3}, 200)
 ]
 
 
diff --git a/tests/test_param_validation.py b/tests/test_param_validation.py
@@ -45,6 +45,7 @@ def test_sparse_not_supported(optimizer_class):
     optim.QHM,
     optim.RAdam,
     optim.SGDW,
+    optim.SGDP,
     optim.Shampoo,
     optim.Yogi,
 ]
@@ -68,6 +69,7 @@ def test_learning_rate(optimizer_class):
     optim.NovoGrad,
     optim.QHAdam,
     optim.RAdam,
+    optim.SGDP,
     optim.Yogi,
 ]
 
@@ -93,12 +95,13 @@ def test_eps_validation(optimizer_class):
     optim.QHM,
     optim.RAdam,
     optim.SGDW,
+    optim.SGDP,
     optim.Shampoo,
     optim.Yogi,
 ]
 
 
-@pytest.mark.parametrize('optimizer_class', optimizers)
+@pytest.mark.parametrize('optimizer_class', weight_decay_optimizers)
 def test_weight_decay_validation(optimizer_class):
     weight_decay = -0.1
     with pytest.raises(ValueError) as ctx:
@@ -120,7 +123,7 @@ def test_weight_decay_validation(optimizer_class):
 ]
 
 
-@pytest.mark.parametrize('optimizer_class', eps_optimizers)
+@pytest.mark.parametrize('optimizer_class', betas_optimizers)
 def test_betas_validation(optimizer_class):
     betas = (-1, 0.999)
     with pytest.raises(ValueError) as ctx:
diff --git a/torch_optimizer/__init__.py b/torch_optimizer/__init__.py
@@ -32,6 +32,7 @@
 from .qhm import QHM
 from .radam import RAdam
 from .sgdw import SGDW
+from .sgdp import SGDP
 from .shampoo import Shampoo
 from .yogi import Yogi
 
@@ -53,6 +54,7 @@
     'RangerQH',
     'RangerVA',
     'SGDW',
+    'SGDP',
     'Shampoo',
     'Yogi',
     # utils
@@ -78,6 +80,7 @@
     RangerQH,
     RangerVA,
     SGDW,
+    SGDP,
     Shampoo,
     Yogi,
 ]  # type: List[Type[Optimizer]]
diff --git a/torch_optimizer/adamp.py b/torch_optimizer/adamp.py
@@ -27,6 +27,7 @@ class AdamP(Optimizer):
             invariant or not (default: 0.1)
         wd_ratio: relative weight decay applied on scale-invariant parameters
             compared to that applied on scale-variant parameters (default: 0.1)
+        nesterov: enables Nesterov momentum (default: False)
 
 
     Example:
@@ -50,7 +51,8 @@ def __init__(
             eps: float = 1e-8,
             weight_decay: float = 0,
             delta: float = 0.1,
-            wd_ratio: float = 0.1
+            wd_ratio: float = 0.1,
+            nesterov: bool = False
     ) -> None:
         if lr <= 0.0:
             raise ValueError('Invalid learning rate: {}'.format(lr))
@@ -68,6 +70,10 @@ def __init__(
             raise ValueError(
                 'Invalid weight_decay value: {}'.format(weight_decay)
             )
+        if delta < 0:
+            raise ValueError('Invalid delta value: {}'.format(delta))
+        if wd_ratio < 0:
+            raise ValueError('Invalid wd_ratio value: {}'.format(wd_ratio))
 
         defaults = dict(
             lr=lr,
@@ -76,6 +82,7 @@ def __init__(
             weight_decay=weight_decay,
             delta=delta,
             wd_ratio=wd_ratio,
+            nesterov=nesterov
         )
         super(AdamP, self).__init__(params, defaults)
 
@@ -136,6 +143,7 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
 
                 grad = p.grad.data
                 beta1, beta2 = group['betas']
+                nesterov = group['nesterov']
 
                 state = self.state[p]
 
@@ -159,7 +167,10 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                          math.sqrt(bias_correction2)).add_(group['eps'])
                 step_size = group['lr'] / bias_correction1
 
-                perturb = exp_avg / denom
+                if nesterov:
+                    perturb = (beta1 * exp_avg + (1 - beta1) * grad) / denom
+                else:
+                    perturb = exp_avg / denom
 
                 # Projection
                 wd_ratio = 1
diff --git a/torch_optimizer/sgdp.py b/torch_optimizer/sgdp.py

Original file line number	Diff line number	Diff line change
`@@ -60,7 +60,8 @@ def build_lookahead(a, *kw):`
`60`	`60`	`(optim.AccSGD, {'lr': 0.015}, 800),`
`61`	`61`	`(build_lookahead, {'lr': 1.0}, 500),`
`62`	`62`	`(optim.QHAdam, {'lr': 1.0}, 500),`
`63`		`- (optim.AdamP, {'lr': 0.01, 'betas': (0.9, 0.95), 'eps': 1e-3}, 800)`
	`63`	`+ (optim.AdamP, {'lr': 0.01, 'betas': (0.9, 0.95), 'eps': 1e-3}, 800),`
	`64`	`+ (optim.SGDP, {'lr': 0.002, 'momentum': 0.91}, 900),`
`64`	`65`	`]`
`65`	`66`
`66`	`67`