Merge pull request #155 from kozistr/feature/avagrad-optimizer

[Feature] Implement AvaGrad optimizer

Author: kozistr

Makefile

@@ -13,7 +13,7 @@ test:
 
 check:
 	black -S -l 119 --check pytorch_optimizer tests hubconf.py
-	ruff pytorch_optimizer tests hubconf.py
+	ruff check pytorch_optimizer tests hubconf.py
 
 requirements:
 	python -m poetry export -f requirements.txt --output requirements.txt --without-hashes

README.rst

@@ -16,7 +16,7 @@ pytorch-optimizer
 
 | **pytorch-optimizer** is optimizer & lr scheduler collections in PyTorch.
 | I just re-implemented (speed & memory tweaks, plug-ins) the algorithm while based on the original paper. Also, It includes useful and practical optimization ideas.
-| Currently, 47 optimizers, 6 lr schedulers are supported!
+| Currently, 48 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -179,6 +179,10 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | SRMM         | *Stochastic regularized majorization-minimization with weakly convex and multi-convex surrogates* | `github <https://github.com/HanbaekLyu/SRMM>`__                                   | `https://arxiv.org/abs/2201.01652 <https://arxiv.org/abs/2201.01652>`__                       |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| AvaGrad      | *Domain-independent Dominance of Adaptive Methods*                                                | `github <https://github.com/lolemacs/avagrad>`__                                  | `https://arxiv.org/abs/1912.01823 <https://arxiv.org/abs/1912.01823>`__                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| PCGrad       | *Gradient Surgery for Multi-Task Learning*                                                        | `github <https://github.com/tianheyu927/PCGrad>`__                                | `https://arxiv.org/abs/2001.06782 <https://arxiv.org/abs/2001.06782>`__                       |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -197,7 +201,7 @@ Also, most of the captures are taken from ``Ranger21`` paper.
 +------------------------------------------+---------------------------------------------+--------------------------------------------+
 | `Lookahead`_                             | `Chebyshev learning rate schedule`_         | `(Adaptive) Sharpness-Aware Minimization`_ |
 +------------------------------------------+---------------------------------------------+--------------------------------------------+
-| `On the Convergence of Adam and Beyond`_ | `Gradient Surgery for Multi-Task Learning`_ |                                            |
+| `On the Convergence of Adam and Beyond`_ | `Improved bias-correction in Adam`_         | `Adaptive Gradient Norm Correction`_       |
 +------------------------------------------+---------------------------------------------+--------------------------------------------+
 
 Adaptive Gradient Clipping
@@ -291,7 +295,7 @@ Lookahead
 Chebyshev learning rate schedule
 --------------------------------
 
-Acceleration via Fractal Learning Rate Schedules
+Acceleration via Fractal Learning Rate Schedules.
 
 -  paper : `arXiv <https://arxiv.org/abs/2103.01338v1>`__
 
@@ -310,10 +314,16 @@ On the Convergence of Adam and Beyond
 
 - paper : `paper <https://openreview.net/forum?id=ryQu7f-RZ>`__
 
-Gradient Surgery for Multi-Task Learning
-----------------------------------------
+Improved bias-correction in Adam
+--------------------------------
+
+| With the default bias-correction, Adam may actually make larger than requested gradient updates early in training.
+
+- paper : `arXiv <https://arxiv.org/abs/2110.10828>`_
+
+Adaptive Gradient Norm Correction
+---------------------------------
 
-- paper : `paper <https://arxiv.org/abs/2001.06782>`__
 
 Citations
 ---------
@@ -358,7 +368,7 @@ Citations
 
 `On the Convergence of Adam and Beyond <https://ui.adsabs.harvard.edu/abs/2019arXiv190409237R/exportcitation>`__
 
-`Gradient surgery for multi-task learning <https://ui.adsabs.harvard.edu/abs/2020arXiv200106782Y/exportcitation>`__
+`Gradient surgery for multi-task learning <https://github.com/tianheyu927/PCGrad#reference>`__
 
 `AdamD <https://ui.adsabs.harvard.edu/abs/2021arXiv211010828S/exportcitation>`__
 
@@ -420,6 +430,8 @@ Citations
 
 `SRMM <https://ui.adsabs.harvard.edu/abs/2022arXiv220101652L/exportcitation>`__
 
+`AvaGrad <https://ui.adsabs.harvard.edu/abs/2019arXiv191201823S/exportcitation>`__
+
 Citation
 --------
 

docs/conf.py

@@ -9,22 +9,19 @@
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
-#
 import os
 import sys
 
 sys.path.insert(0, os.path.abspath('../'))
 
-import sphinx_rtd_theme
-
 # -- Project information -----------------------------------------------------
 
 project = 'pytorch-optimizers'
 copyright = '2021, kozistr'
 author = 'kozistr'
 
 # The full version, including alpha/beta/rc tags
-release = '2.2.1'
+release = '2.8.0'
 
 
 # -- General configuration ---------------------------------------------------
@@ -33,16 +30,16 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'sphinx_rtd_theme',
+    # 'sphinx_rtd_theme',
     'sphinx.ext.autodoc',
     'sphinx.ext.napoleon',
-    'sphinx.ext.doctest',
     'sphinx.ext.todo',
     'sphinx.ext.coverage',
     'sphinx.ext.mathjax',
-    'sphinx.ext.ifconfig',
     'sphinx.ext.viewcode',
-    'sphinx.ext.intersphinx',
+    # 'sphinx.ext.autosummary',
+    # 'sphinx.ext.ifconfig',
+    # 'sphinx.ext.intersphinx',
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -59,7 +56,7 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
+html_theme = 'alabaster'
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,

docs/index.rst

@@ -3,7 +3,7 @@
    You can adapt this file completely to your liking, but it should at least
    contain the root `toctree` directive.
 
-Welcome to pytorch-optimizers's documentation!
+Welcome to pytorch-optimizers documentation!
 ==============================================
 
 .. include:: ../README.rst

docs/optimizer_api.rst

@@ -432,3 +432,11 @@ SRMM
 
 .. autoclass:: pytorch_optimizer.SRMM
     :members:
+
+.. _AvaGrad:
+
+AvaGrad
+-------
+
+.. autoclass:: pytorch_optimizer.AvaGrad
+    :members: