Merge pull request #101 from kozistr/feature/d-adaptation

[Feature] Implements D-Adaptation optimizers

Author: kozistr

README.rst

@@ -114,6 +114,8 @@ Supported Optimizers
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 | GSAM         | *Surrogate Gap Guided Sharpness-Aware Minimization*                                    | `github <https://github.com/juntang-zhuang/GSAM>`__                               | `https://openreview.net/pdf?id=edONMAnhLu- <https://openreview.net/pdf?id=edONMAnhLu->`__     |
 +--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
+| D-Adaptation | *Learning-Rate-Free Learning by D-Adaptation*                                          | `github <https://github.com/facebookresearch/dadaptation>`__                      | `https://arxiv.org/abs/2301.07733 <https://arxiv.org/abs/2301.07733>`__                       |
++--------------+----------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+
 
 Useful Resources
 ----------------
@@ -307,6 +309,8 @@ Citations
 
 `GSAM <https://github.com/juntang-zhuang/GSAM#citation>`__
 
+`D-Adaptation <https://ui.adsabs.harvard.edu/abs/2023arXiv230107733D/exportcitation>`__
+
 Citation
 --------
 

docs/optimizer_api.rst

@@ -200,3 +200,27 @@ GSAM
 
 .. autoclass:: pytorch_optimizer.GSAM
     :members:
+
+.. _DAdaptAdaGrad:
+
+DAdaptAdaGrad
+-------------
+
+.. autoclass:: pytorch_optimizer.DAdaptAdaGrad
+    :members:
+
+.. _DAdaptAdam:
+
+DAdaptAdam
+----------
+
+.. autoclass:: pytorch_optimizer.DAdaptAdam
+    :members:
+
+.. _DAdaptSGD:
+
+DAdaptSGD
+---------
+
+.. autoclass:: pytorch_optimizer.DAdaptSGD
+    :members:

docs/util_api.rst

@@ -65,3 +65,108 @@ disable_running_stats
 
 .. autoclass:: pytorch_optimizer.disable_running_stats
     :members:
+
+.. _LayerWiseGrafting:
+
+LayerWiseGrafting
+-----------------
+
+.. autoclass:: pytorch_optimizer.LayerWiseGrafting
+    :members:
+
+.. _Graft:
+
+Graft
+-----
+
+.. autoclass:: pytorch_optimizer.Graft
+    :members:
+
+.. _SGDGraft:
+
+SGDGraft
+--------
+
+.. autoclass:: pytorch_optimizer.SGDGraft
+    :members:
+
+.. _SQRTNGraft:
+
+SQRTNGraft
+----------
+
+.. autoclass:: pytorch_optimizer.SQRTNGraft
+    :members:
+
+.. _AdaGradGraft:
+
+AdaGradGraft
+------------
+
+.. autoclass:: pytorch_optimizer.AdaGradGraft
+    :members:
+
+.. _RMSPropGraft:
+
+RMSPropGraft
+------------
+
+.. autoclass:: pytorch_optimizer.RMSPropGraft
+    :members:
+
+.. _BlockPartitioner:
+
+BlockPartitioner
+----------------
+
+.. autoclass:: pytorch_optimizer.BlockPartitioner
+    :members:
+
+.. _PreConditionerType:
+
+PreConditionerType
+------------------
+
+.. autoclass:: pytorch_optimizer.PreConditionerType
+    :members:
+
+.. _PreConditioner:
+
+PreConditioner
+--------------
+
+.. autoclass:: pytorch_optimizer.PreConditioner
+    :members:
+
+.. _power_iter:
+
+power_iter
+----------
+
+.. autoclass:: pytorch_optimizer.power_iter
+    :members:
+
+.. _matrix_power:
+
+matrix_power
+------------
+
+.. autoclass:: pytorch_optimizer.matrix_power
+    :members:
+
+.. _compute_power:
+
+compute_power
+-------------
+
+.. autoclass:: pytorch_optimizer.compute_power
+    :members:
+
+.. _merge_small_dims:
+
+merge_small_dims
+----------------
+
+.. autoclass:: pytorch_optimizer.merge_small_dims
+    :members:
+re

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.3.1"
+version = "2.4.0"
 description = "optimizer & lr scheduler 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

@@ -21,6 +21,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.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
 from pytorch_optimizer.optimizer.diffrgrad import DiffRGrad
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
@@ -40,6 +41,21 @@
 from pytorch_optimizer.optimizer.sam import SAM
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import Shampoo
+from pytorch_optimizer.optimizer.shampoo_utils import (
+    AdaGradGraft,
+    BlockPartitioner,
+    Graft,
+    LayerWiseGrafting,
+    PreConditioner,
+    PreConditionerType,
+    RMSPropGraft,
+    SGDGraft,
+    SQRTNGraft,
+    compute_power,
+    matrix_power,
+    merge_small_dims,
+    power_iter,
+)
 from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     disable_running_stats,
@@ -69,6 +85,9 @@
     Ranger21,
     SGDP,
     Shampoo,
+    DAdaptAdaGrad,
+    DAdaptAdam,
+    DAdaptSGD,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}