Merge pull request #157 from kozistr/feature/adashift-optimizer

[Feature] Implement AdaShift optimizer

Author: kozistr

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, 48 optimizers, 6 lr schedulers are supported!
+| Currently, 49 optimizers, 6 lr schedulers are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -207,7 +207,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | Chebyshev LR | *Acceleration via Fractal Learning Rate Schedules*                                                |                                                                                   | `https://arxiv.org/abs/2103.01338 <https://arxiv.org/abs/2103.01338>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210301338A/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| Untuned WU   | *On the adequacy of untuned warmup for adaptive optimization*                                     |                                                                                   | `https://arxiv.org/abs/1910.04209 <https://arxiv.org/abs/1910.04209>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2019arXiv191004209M/exportcitation>`__                                      |
+| Un-tuned WU  | *On the adequacy of untuned warmup for adaptive optimization*                                     |                                                                                   | `https://arxiv.org/abs/1910.04209 <https://arxiv.org/abs/1910.04209>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2019arXiv191004209M/exportcitation>`__                                      |
++--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
+| AdaShift     | *Decorrelation and Convergence of Adaptive Learning Rate Methods*                                 | `github <https://github.com/MichaelKonobeev/adashift>`__                          | `https://arxiv.org/abs/1810.00143v4 <https://arxiv.org/abs/1810.00143v4>`__                   | `cite <https://ui.adsabs.harvard.edu/abs/2018arXiv181000143Z/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 
 Useful Resources

docs/optimizer_api.rst

@@ -440,3 +440,11 @@ AvaGrad
 
 .. autoclass:: pytorch_optimizer.AvaGrad
     :members:
+
+.. _AdaShift:
+
+AdaShift
+--------
+
+.. autoclass:: pytorch_optimizer.AdaShift
+    :members:

pytorch_optimizer/__init__.py

@@ -28,6 +28,7 @@
 from pytorch_optimizer.optimizer.adan import Adan
 from pytorch_optimizer.optimizer.adanorm import AdaNorm
 from pytorch_optimizer.optimizer.adapnm import AdaPNM
+from pytorch_optimizer.optimizer.adashift import AdaShift
 from pytorch_optimizer.optimizer.adasmooth import AdaSmooth
 from pytorch_optimizer.optimizer.agc import agc
 from pytorch_optimizer.optimizer.aggmo import AggMo
@@ -141,6 +142,7 @@
     AdaSmooth,
     SRMM,
     AvaGrad,
+    AdaShift,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

pytorch_optimizer/base/optimizer.py

@@ -127,7 +127,7 @@ def validate_range(x: float, name: str, low: float, high: float, range_type: str
             raise ValueError(f'[-] {name} must be in the range ({low}, {high})')
 
     @staticmethod
-    def validate_negative(x: float, name: str):
+    def validate_non_negative(x: float, name: str):
         if x < 0.0:
             raise ValueError(f'[-] {name} must be non-negative')
 

pytorch_optimizer/optimizer/a2grad.py

@@ -30,8 +30,8 @@ def __init__(
         variant: str = 'uni',
     ):
         self.validate_learning_rate(lr)
-        self.validate_negative(lips, 'lips')
-        self.validate_negative(rho, 'rho')
+        self.validate_non_negative(lips, 'lips')
+        self.validate_non_negative(rho, 'rho')
         self.validate_a2grad_variant(variant)
 
         self.variant = variant

pytorch_optimizer/optimizer/adabelief.py

@@ -46,8 +46,8 @@ def __init__(
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
-        self.validate_negative(weight_decay, 'weight_decay')
-        self.validate_negative(eps, 'eps')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
 
         self.n_sma_threshold = n_sma_threshold
         self.degenerated_to_sgd = degenerated_to_sgd

pytorch_optimizer/optimizer/adabound.py

@@ -41,8 +41,8 @@ def __init__(
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
-        self.validate_negative(weight_decay, 'weight_decay')
-        self.validate_negative(eps, 'eps')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
 
         defaults: DEFAULTS = {
             'lr': lr,

pytorch_optimizer/optimizer/adafactor.py

@@ -48,9 +48,9 @@ def __init__(
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
-        self.validate_negative(weight_decay, 'weight_decay')
-        self.validate_negative(eps1, 'eps1')
-        self.validate_negative(eps2, 'eps2')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps1, 'eps1')
+        self.validate_non_negative(eps2, 'eps2')
 
         self.decay_rate = decay_rate
         self.clip_threshold = clip_threshold

pytorch_optimizer/optimizer/adai.py

@@ -39,8 +39,8 @@ def __init__(
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
-        self.validate_negative(weight_decay, 'weight_decay')
-        self.validate_negative(eps, 'eps')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
 
         self.use_gc = use_gc
 

pytorch_optimizer/optimizer/adamax.py

@@ -36,8 +36,8 @@ def __init__(
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
-        self.validate_negative(weight_decay, 'weight_decay')
-        self.validate_negative(eps, 'eps')
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_non_negative(eps, 'eps')
 
         defaults: DEFAULTS = {
             'lr': lr,