pytorch-optimizer 2.11.1__tar.gz → 2.11.2__tar.gz

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
-Name: pytorch-optimizer
-Version: 2.11.1
+Name: pytorch_optimizer
+Version: 2.11.2
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 Home-page: https://github.com/kozistr/pytorch_optimizer
 License: Apache-2.0
-Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,AdamP,AdamS,Adan,AggMo,AliG,Amos,Apollo,AvaGrad,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,GSAM,LARS,Lamb,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,Tiger,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,AdamP,AdamS,Adan,AggMo,AliG,Amos,Apollo,AvaGrad,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,Fromage,Gravity,GSAM,LARS,Lamb,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,Tiger,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -54,7 +54,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, 59 optimizers, 10 lr schedulers, and 13 loss functions are supported!
+| Currently, **60 optimizers**, **10 lr schedulers**, and **13 loss functions** are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -69,20 +69,20 @@ So, please double-check the license before using it at your work.
 Installation
 ~~~~~~~~~~~~
 
-::
+.. code-block:: bash
 
     $ pip3 install -U pytorch-optimizer
 
 If there's a version issue when installing the package, try with `--no-deps` option.
 
-::
+.. code-block:: bash
 
     $ pip3 install -U --no-deps pytorch-optimizer
 
 Simple Usage
 ~~~~~~~~~~~~
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import AdamP
 
@@ -99,7 +99,7 @@ Simple Usage
 
 Also, you can load the optimizer via `torch.hub`
 
-::
+.. code-block:: python
 
     import torch
 
@@ -109,7 +109,7 @@ Also, you can load the optimizer via `torch.hub`
 
 If you want to build the optimizer with parameters & configs, there's `create_optimizer()` API.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import create_optimizer
 
@@ -127,7 +127,7 @@ Supported Optimizers
 
 You can check the supported optimizers with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_optimizers
 
@@ -268,7 +268,7 @@ Supported LR Scheduler
 
 You can check the supported learning rate schedulers with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_lr_schedulers
 
@@ -287,7 +287,7 @@ Supported Loss Function
 
 You can check the supported loss functions with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_loss_functions
 
@@ -471,9 +471,9 @@ Hyeongchan Kim / `@kozistr <http://kozistr.tech/about>`__
    :target: https://pytorch-optimizers.readthedocs.io/en/latest/?badge=latest
 .. |PyPI version| image:: https://badge.fury.io/py/pytorch-optimizer.svg
    :target: https://badge.fury.io/py/pytorch-optimizer
-.. |PyPi download| image:: https://pepy.tech/badge/pytorch-optimizer
+.. |PyPi download| image:: https://static.pepy.tech/badge/pytorch-optimizer
    :target: https://pepy.tech/project/pytorch-optimizer
-.. |PyPi month download| image:: https://pepy.tech/badge/pytorch-optimizer/month
+.. |PyPi month download| image:: https://static.pepy.tech/badge/pytorch-optimizer/month
    :target: https://pepy.tech/project/pytorch-optimizer
 .. |PyPI pyversions| image:: https://img.shields.io/pypi/pyversions/pytorch-optimizer.svg
    :target: https://pypi.python.org/pypi/pytorch-optimizer/

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/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, 59 optimizers, 10 lr schedulers, and 13 loss functions are supported!
+| Currently, **60 optimizers**, **10 lr schedulers**, and **13 loss functions** are supported!
 |
 | Highly inspired by `pytorch-optimizer <https://github.com/jettify/pytorch-optimizer>`__.
 
@@ -31,20 +31,20 @@ So, please double-check the license before using it at your work.
 Installation
 ~~~~~~~~~~~~
 
-::
+.. code-block:: bash
 
     $ pip3 install -U pytorch-optimizer
 
 If there's a version issue when installing the package, try with `--no-deps` option.
 
-::
+.. code-block:: bash
 
     $ pip3 install -U --no-deps pytorch-optimizer
 
 Simple Usage
 ~~~~~~~~~~~~
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import AdamP
 
@@ -61,7 +61,7 @@ Simple Usage
 
 Also, you can load the optimizer via `torch.hub`
 
-::
+.. code-block:: python
 
     import torch
 
@@ -71,7 +71,7 @@ Also, you can load the optimizer via `torch.hub`
 
 If you want to build the optimizer with parameters & configs, there's `create_optimizer()` API.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import create_optimizer
 
@@ -89,7 +89,7 @@ Supported Optimizers
 
 You can check the supported optimizers with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_optimizers
 
@@ -230,7 +230,7 @@ Supported LR Scheduler
 
 You can check the supported learning rate schedulers with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_lr_schedulers
 
@@ -249,7 +249,7 @@ Supported Loss Function
 
 You can check the supported loss functions with below code.
 
-::
+.. code-block:: python
 
     from pytorch_optimizer import get_supported_loss_functions
 
@@ -433,9 +433,9 @@ Hyeongchan Kim / `@kozistr <http://kozistr.tech/about>`__
    :target: https://pytorch-optimizers.readthedocs.io/en/latest/?badge=latest
 .. |PyPI version| image:: https://badge.fury.io/py/pytorch-optimizer.svg
    :target: https://badge.fury.io/py/pytorch-optimizer
-.. |PyPi download| image:: https://pepy.tech/badge/pytorch-optimizer
+.. |PyPi download| image:: https://static.pepy.tech/badge/pytorch-optimizer
    :target: https://pepy.tech/project/pytorch-optimizer
-.. |PyPi month download| image:: https://pepy.tech/badge/pytorch-optimizer/month
+.. |PyPi month download| image:: https://static.pepy.tech/badge/pytorch-optimizer/month
    :target: https://pepy.tech/project/pytorch-optimizer
 .. |PyPI pyversions| image:: https://img.shields.io/pypi/pyversions/pytorch-optimizer.svg
    :target: https://pypi.python.org/pypi/pytorch-optimizer/

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.11.1"
+version = "2.11.2"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -13,10 +13,10 @@ keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
     "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "AdamP",
     "AdamS", "Adan", "AggMo", "AliG", "Amos", "Apollo", "AvaGrad", "CAME", "DAdaptAdaGrad", "DAdaptAdam", "DAdaptAdan",
-    "DAdaptSGD", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead", "MADGRAD",
-    "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
-    "Ranger21", "RotoGrad", "SAM", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM",
-    "SWATS", "Tiger", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard",
+    "DAdaptSGD", "DAdaptLion", "DiffGrad", "Fromage", "Gravity", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead",
+    "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam",
+    "Ranger", "Ranger21", "RotoGrad", "SAM", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH",
+    "SRMM", "SWATS", "Tiger", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard",
     "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge",
 ]
 classifiers = [
@@ -60,7 +60,7 @@ black = [
     { version = "==23.3.0", python = ">=3.7,<3.8" },
     { version = "^23.7.0", python = ">=3.8"}
 ]
-ruff = "^0.0.278"
+ruff = "^0.0.286"
 pytest = "^7.4.0"
 pytest-cov = "^4.1.0"
 
@@ -122,7 +122,6 @@ testpaths = "tests"
 [tool.coverage.run]
 omit = [
     "./pytorch_optimizer/optimizer/gsam.py",
-    "./pytorch_optimizer/optimizer/fp16.py",
     "./pytorch_optimizer/optimizer/rotograd.py",
 ]
 

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/__init__.py

@@ -48,7 +48,7 @@ from pytorch_optimizer.optimizer.amos import Amos
 from pytorch_optimizer.optimizer.apollo import Apollo
 from pytorch_optimizer.optimizer.avagrad import AvaGrad
 from pytorch_optimizer.optimizer.came import CAME
-from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptSGD
+from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptLion, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.fromage import Fromage
@@ -171,6 +171,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     LOMO,
     Tiger,
     CAME,
+    DAdaptLion,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/adabelief.py

@@ -35,7 +35,7 @@ class AdaBelief(Optimizer, BaseOptimizer):
         weight_decay: float = 0.0,
         weight_decouple: bool = True,
         fixed_decay: bool = False,
-        rectify: bool = True,
+        rectify: bool = False,
         n_sma_threshold: int = 5,
         degenerated_to_sgd: bool = True,
         ams_bound: bool = False,

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/dadapt.py

@@ -699,3 +699,132 @@ class DAdaptAdan(Optimizer, BaseOptimizer):
             group['k'] += 1
 
         return loss
+
+
+class DAdaptLion(Optimizer, BaseOptimizer):
+    r"""Lion with D-Adaptation. Leave LR set to 1 unless you encounter instability. This implementation is based on V3.
+
+    :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 d0: float. initial D estimate for D-adaptation (default 1e-6). Rarely needs changing.
+    :param weight_decay: float. weight decay (L2 penalty).
+    :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.
+    :param fixed_decay: bool. fix weight decay.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1.0,
+        betas: BETAS = (0.9, 0.999),
+        d0: float = 1e-6,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = False,
+        fixed_decay: bool = False,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'd': d0,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'step': 0,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'DAdaptLion'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                state = self.state[p]
+
+                state['exp_avg'] = torch.zeros_like(p)
+                state['s'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        group = self.param_groups[0]
+        device = group['params'][0].device
+
+        if 'numerator_weighted' not in group:
+            group['numerator_weighted'] = torch.tensor([0.0], device=device)
+        numerator_weighted = group['numerator_weighted']
+
+        sk_l1 = torch.tensor([0.0], device=device)
+        numerator_accumulator = torch.tensor([0.0], device=device)
+
+        beta1, beta2 = group['betas']
+        beta2_sq = math.sqrt(beta2)
+
+        d, lr = group['d'], group['lr']
+        d_lr: float = d * lr
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise NoSparseGradientError(str(self))
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['s'] = torch.zeros_like(p)
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=d_lr,
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                exp_avg, s = state['exp_avg'], state['s']
+
+                update = exp_avg.clone().mul_(beta1).add_(grad, alpha=1.0 - beta1).sign_()
+                p.add_(update, alpha=-d_lr)
+
+                exp_avg.mul_(beta2).add_(grad, alpha=(1.0 - beta2) * d_lr)
+
+                numerator_accumulator.add_(torch.dot(update.flatten(), s.flatten()), alpha=d_lr)
+                s.mul_(beta2_sq).add_(update, alpha=(1.0 - beta2_sq) * d_lr)
+
+                sk_l1.add_(s.abs().sum())
+
+        numerator_weighted.mul_(beta2_sq).add_(numerator_accumulator, alpha=1.0 - beta2_sq)
+
+        if sk_l1 == 0:
+            return loss
+
+        if lr > 0.0:
+            d_hat: float = (numerator_weighted / ((1.0 - beta2_sq) * sk_l1)).item()
+            d = max(d, d_hat)
+
+        for group in self.param_groups:
+            group['step'] += 1
+
+            group['numerator_weighted'] = numerator_weighted
+            group['d'] = d
+
+        return loss

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/fp16.py

@@ -90,7 +90,7 @@ class DynamicLossScaler:
             self.loss_scale = max(self.loss_scale, self.threshold)
 
 
-class SafeFP16Optimizer(Optimizer):
+class SafeFP16Optimizer(Optimizer):  # pragma: no cover
     r"""Safe FP16 Optimizer.
 
     :param optimizer: OPTIMIZER.

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/lomo.py

@@ -46,7 +46,7 @@ class LOMO(BaseOptimizer, Optimizer):
         self.grad_norms: List[torch.Tensor] = []
         self.clip_coef: Optional[float] = None
 
-        p0: torch.Tensor = list(self.model.parameters())[0]
+        p0: torch.Tensor = next(iter(self.model.parameters()))
 
         self.grad_func: Callable[[Any], Any] = (
             self.fuse_update_zero3() if hasattr(p0, 'ds_tensor') else self.fuse_update()

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/lookahead.py

@@ -1,13 +1,14 @@
 from collections import defaultdict
-from typing import Dict
+from typing import Callable, Dict
 
 import torch
+from torch.optim import Optimizer
 
 from pytorch_optimizer.base.optimizer import BaseOptimizer
-from pytorch_optimizer.base.types import CLOSURE, LOSS, OPTIMIZER, STATE
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER, STATE
 
 
-class Lookahead(BaseOptimizer):
+class Lookahead(Optimizer, BaseOptimizer):
     r"""k steps forward, 1 step back.
 
     :param optimizer: OPTIMIZER. base optimizer.
@@ -27,12 +28,16 @@ class Lookahead(BaseOptimizer):
         self.validate_range(alpha, 'alpha', 0.0, 1.0)
         self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
 
+        self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
+        self._optimizer_step_post_hooks: Dict[int, Callable] = {}
+
         self.alpha = alpha
         self.k = k
         self.pullback_momentum = pullback_momentum
 
         self.optimizer = optimizer
         self.param_groups = self.optimizer.param_groups
+
         self.state: STATE = defaultdict(dict)
 
         for group in self.param_groups:
@@ -46,6 +51,13 @@ class Lookahead(BaseOptimizer):
                 if self.pullback_momentum == 'pullback':
                     state['slow_momentum'] = torch.zeros_like(p)
 
+        self.defaults: DEFAULTS = {
+            'lookahead_alpha': alpha,
+            'lookahead_k': k,
+            'lookahead_pullback_momentum': pullback_momentum,
+            **optimizer.defaults,
+        }
+
     def __getstate__(self):
         return {
             'state': self.state,

{pytorch_optimizer-2.11.1 → pytorch_optimizer-2.11.2}/pytorch_optimizer/optimizer/rotograd.py

@@ -219,7 +219,7 @@ class RotateOnly(nn.Module):
         self.backbone.to(*args, **kwargs)
         for head in self.heads:
             head.to(*args, **kwargs)
-        return super(RotateOnly, self).to(*args, **kwargs)
+        return super().to(*args, **kwargs)
 
     def train(self, mode: bool = True) -> nn.Module:
         super().train(mode)
@@ -284,7 +284,7 @@ class RotateOnly(nn.Module):
         if not self.training:
             raise AssertionError('Backward should only be called when training')
 
-        if self.iteration_counter == 0 or self.iteration_counter == self.burn_in_period:
+        if self.iteration_counter in (0, self.burn_in_period):
             for i, loss in enumerate(losses):
                 self.initial_losses[i] = loss.item()