pytorch-optimizer 3.3.3__tar.gz → 3.3.4__tar.gz

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: pytorch_optimizer
-Version: 3.3.3
+Version: 3.3.4
 Summary: optimizer & lr scheduler & objective function collections in PyTorch
 License: Apache-2.0
-Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,ADOPT,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,APOLLO,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DeMo,DiffGrad,FAdam,Fromage,FTRL,GaLore,Grams,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,OrthoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SRMM,StableAdamW,SWATS,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdamG,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdEMAMix,ADOPT,AdaHessian,Adai,Adalite,AdaLomo,AdamMini,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,APOLLO,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DeMo,DiffGrad,FAdam,Fromage,FTRL,GaLore,Grams,Gravity,GrokFast,GSAM,Kate,Lamb,LaProp,LARS,Lion,LOMO,Lookahead,MADGRAD,MARS,MSVAG,Muno,Nero,NovoGrad,OrthoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,ScheduleFreeRAdam,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SOAP,SopihaH,SPAM,SRMM,StableAdamW,SWATS,TAM,Tiger,TRAC,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes,WSD,QGaLore
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -51,8 +51,8 @@ Description-Content-Type: text/markdown
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-* Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
+* Wide range of supported optimizers. Currently, **92 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
 * Somewhat a bit more optimized compared to the original implementation
@@ -239,6 +239,8 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
 | OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
 | Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
+| SPAM          | *Spike-Aware Adam with Momentum Reset for Stable LLM Training*                                    | [github](https://github.com/TianjinYellow/SPAM-Optimizer)                                                      | <https://arxiv.org/abs/2501.06842>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250106842H/exportcitation)                                                        |
+| TAM           | *Torque-Aware Momentum*                                                                           |                                                                                                                | <https://arxiv.org/abs/2412.18790>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241218790M/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/README.md

@@ -10,8 +10,8 @@
 
 ## The reasons why you use `pytorch-optimizer`.
 
-* Wide range of supported optimizers. Currently, **89 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
-* Including many variants such as `Cautious`, `AdamD`, `Gradient Centrailiaztion`
+* Wide range of supported optimizers. Currently, **92 optimizers (+ `bitsandbytes`, `qgalore`, `torchao`)**, **16 lr schedulers**, and **13 loss functions** are supported!
+* Including many variants such as `ADOPT`, `Cautious`, `AdamD`, `StableAdamW`, and `Gradient Centrailiaztion`
 * Easy to use, clean, and tested codes
 * Active maintenance
 * Somewhat a bit more optimized compared to the original implementation
@@ -198,6 +198,8 @@ get_supported_optimizers(['adam*', 'ranger*'])
 | Grams         | *Gradient Descent with Adaptive Momentum Scaling*                                                 |                                                                                                                | <https://arxiv.org/abs/2412.17107>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241217107C/exportcitation)                                                        |
 | OrthoGrad     | *Grokking at the Edge of Numerical Stability*                                                     | [github](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability)                         | <https://arxiv.org/abs/2501.04697>                                                          | [cite](https://github.com/LucasPrietoAl/grokking-at-the-edge-of-numerical-stability?tab=readme-ov-file#citation)                    |
 | Adam-ATAN2    | *Scaling Exponents Across Parameterizations and Optimizers*                                       |                                                                                                                | <https://arxiv.org/abs/2407.05872>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240705872E/exportcitation)                                                        |
+| SPAM          | *Spike-Aware Adam with Momentum Reset for Stable LLM Training*                                    | [github](https://github.com/TianjinYellow/SPAM-Optimizer)                                                      | <https://arxiv.org/abs/2501.06842>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2025arXiv250106842H/exportcitation)                                                        |
+| TAM           | *Torque-Aware Momentum*                                                                           |                                                                                                                | <https://arxiv.org/abs/2412.18790>                                                          | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv241218790M/exportcitation)                                                        |
 
 ## Supported LR Scheduler
 

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.3.3"
+version = "3.3.4"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -18,9 +18,9 @@ keywords = [
     "Kate", "Lamb", "LaProp", "LARS", "Lion", "LOMO", "Lookahead", "MADGRAD", "MARS", "MSVAG", "Muno", "Nero",
     "NovoGrad", "OrthoGrad", "PAdam", "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger",
     "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD", "ScheduleFreeAdamW", "ScheduleFreeRAdam", "SGDP", "Shampoo",
-    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SRMM", "StableAdamW", "SWATS", "Tiger", "TRAC",
-    "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard", "Bi-Tempered",
-    "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
+    "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SOAP", "SopihaH", "SPAM", "SRMM", "StableAdamW", "SWATS", "TAM",
+    "Tiger", "TRAC", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard",
+    "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes", "WSD", "QGaLore",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/__init__.py

@@ -61,8 +61,10 @@ from pytorch_optimizer.optimizer import (
     SGDW,
     SM3,
     SOAP,
+    SPAM,
     SRMM,
     SWATS,
+    TAM,
     TRAC,
     WSAM,
     A2Grad,
@@ -87,6 +89,7 @@ from pytorch_optimizer.optimizer import (
     AdaPNM,
     AdaShift,
     AdaSmooth,
+    AdaTAM,
     AdEMAMix,
     AggMo,
     Aida,

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/base/scheduler.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
 from typing import List
 
+from torch.optim import Optimizer
+
 from pytorch_optimizer.base.exception import NegativeLRError, NegativeStepError
-from pytorch_optimizer.base.types import OPTIMIZER
 
 
 class BaseLinearWarmupScheduler(ABC):
@@ -10,7 +11,7 @@ class BaseLinearWarmupScheduler(ABC):
 
         The LR Scheduler class based on this class has linear warmup strategy.
 
-    :param optimizer: Optimizer. OPTIMIZER. It will set learning rate to all trainable parameters in optimizer.
+    :param optimizer: Optimizer. It will set learning rate to all trainable parameters in optimizer.
     :param t_max: int. total steps to train.
     :param max_lr: float. maximum lr.
     :param min_lr: float. minimum lr.
@@ -20,7 +21,7 @@ class BaseLinearWarmupScheduler(ABC):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         t_max: int,
         max_lr: float,
         min_lr: float = 0.0,

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/base/types.py

@@ -11,6 +11,7 @@ DEFAULTS = Dict
 PARAMETERS = Optional[Union[Iterable[Dict], Iterable[torch.Tensor]]]
 STATE = Dict
 OPTIMIZER = Type[Optimizer]
+OPTIMIZER_INSTANCE_OR_CLASS = Union[OPTIMIZER, Optimizer]
 SCHEDULER = Type[LRScheduler]
 
 HUTCHINSON_G = Literal['gaussian', 'rademacher']

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/lr_scheduler/cosine_anealing.py

@@ -1,10 +1,9 @@
 import math
 from typing import List, Optional
 
+from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from pytorch_optimizer.base.types import OPTIMIZER
-
 
 class CosineAnnealingWarmupRestarts(LRScheduler):
     r"""CosineAnnealingWarmupRestarts.
@@ -21,7 +20,7 @@ class CosineAnnealingWarmupRestarts(LRScheduler):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         first_cycle_steps: int,
         cycle_mult: float = 1.0,
         max_lr: float = 1e-4,
@@ -53,7 +52,7 @@ class CosineAnnealingWarmupRestarts(LRScheduler):
 
         self.init_lr()
 
-    def init_lr(self):
+    def init_lr(self) -> None:
         self.base_lrs = []
         for param_group in self.optimizer.param_groups:
             param_group['lr'] = self.min_lr

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/lr_scheduler/rex.py

@@ -1,9 +1,8 @@
 from typing import List, Optional
 
+from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from pytorch_optimizer.base.types import OPTIMIZER
-
 
 class REXScheduler(LRScheduler):
     r"""Revisiting Budgeted Training with an Improved Schedule.
@@ -16,7 +15,7 @@ class REXScheduler(LRScheduler):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         total_steps: int,
         max_lr: float = 1.0,
         min_lr: float = 0.0,
@@ -35,7 +34,7 @@ class REXScheduler(LRScheduler):
 
         self.init_lr()
 
-    def init_lr(self):
+    def init_lr(self) -> None:
         self.base_lrs = []
         for param_group in self.optimizer.param_groups:
             param_group['lr'] = self.min_lr

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/__init__.py

@@ -1,10 +1,10 @@
 import fnmatch
 from importlib.util import find_spec
-from typing import Dict, List, Optional, Sequence, Set, Union
+from typing import Dict, List, Optional, Sequence, Set, Type, Union
 
 import torch
 from torch import nn
-from torch.optim import AdamW
+from torch.optim import AdamW, Optimizer
 
 from pytorch_optimizer.base.types import OPTIMIZER, PARAMETERS
 from pytorch_optimizer.optimizer.a2grad import A2Grad
@@ -83,8 +83,10 @@ from pytorch_optimizer.optimizer.shampoo import ScalableShampoo, Shampoo
 from pytorch_optimizer.optimizer.sm3 import SM3
 from pytorch_optimizer.optimizer.soap import SOAP
 from pytorch_optimizer.optimizer.sophia import SophiaH
+from pytorch_optimizer.optimizer.spam import SPAM
 from pytorch_optimizer.optimizer.srmm import SRMM
 from pytorch_optimizer.optimizer.swats import SWATS
+from pytorch_optimizer.optimizer.tam import TAM, AdaTAM
 from pytorch_optimizer.optimizer.tiger import Tiger
 from pytorch_optimizer.optimizer.trac import TRAC
 from pytorch_optimizer.optimizer.yogi import Yogi
@@ -251,6 +253,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     SRMM,
     AvaGrad,
     AdaShift,
+    TAM,
+    AdaTAM,
     AdaDelta,
     Amos,
     AdaHessian,
@@ -286,6 +290,7 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     MARS,
     SGDSaI,
     Grams,
+    SPAM,
     Ranger25,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
@@ -298,8 +303,9 @@ def create_optimizer(
     weight_decay: float = 0.0,
     wd_ban_list: List[str] = ('bias', 'LayerNorm.bias', 'LayerNorm.weight'),
     use_lookahead: bool = False,
+    use_orthograd: bool = False,
     **kwargs,
-):
+) -> Optimizer:
     r"""Build optimizer.
 
     :param model: nn.Module. model.
@@ -307,7 +313,8 @@ def create_optimizer(
     :param lr: float. learning rate.
     :param weight_decay: float. weight decay.
     :param wd_ban_list: List[str]. weight decay ban list by layer.
-    :param use_lookahead: bool. use lookahead.
+    :param use_lookahead: bool. use Lookahead.
+    :param use_orthograd: bool. use OrthoGrad.
     """
     optimizer_name = optimizer_name.lower()
 
@@ -315,14 +322,17 @@ def create_optimizer(
         get_optimizer_parameters(model, weight_decay, wd_ban_list) if weight_decay > 0.0 else model.parameters()
     )
 
-    optimizer = load_optimizer(optimizer_name)
+    optimizer_class: OPTIMIZER = load_optimizer(optimizer_name)
 
     if optimizer_name == 'alig':
-        optimizer = optimizer(parameters, max_lr=lr, **kwargs)
+        optimizer = optimizer_class(parameters, max_lr=lr, **kwargs)
     elif optimizer_name in {'lomo', 'adalomo', 'adammini'}:
-        optimizer = optimizer(model, lr=lr, **kwargs)
+        optimizer = optimizer_class(model, lr=lr, **kwargs)
     else:
-        optimizer = optimizer(parameters, lr=lr, **kwargs)
+        optimizer = optimizer_class(parameters, lr=lr, **kwargs)
+
+    if use_orthograd:
+        optimizer = OrthoGrad(optimizer, **kwargs)
 
     if use_lookahead:
         optimizer = Lookahead(

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/adafactor.py

@@ -13,8 +13,8 @@ class AdaFactor(BaseOptimizer):
 
     :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. if beta1 is None, first momentum will be skipped.
+    :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
+        if beta1 is None, first momentum will be skipped.
     :param decay_rate: float. coefficient used to compute running averages of square gradient.
     :param weight_decay: float. weight decay (L2 penalty).
     :param weight_decouple: bool. the optimizer uses decoupled weight decay as in AdamW.

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/adamp.py

@@ -141,6 +141,10 @@ class AdamP(BaseOptimizer):
                 inv_de_nom = exp_avg_sq.rsqrt().add_(group['eps']).mul_(bias_correction2_sq)
 
                 perturb = exp_avg.clone()
+
+                if self.cautious:
+                    self.apply_cautious(perturb, grad)
+
                 if group['nesterov']:
                     perturb.mul_(beta1).addcmul_(grad, inv_de_nom, value=1.0 - beta1)
                 else:
@@ -173,9 +177,6 @@ class AdamP(BaseOptimizer):
                     bias_correction1=bias_correction1,
                 )
 
-                if self.cautious:
-                    self.apply_cautious(perturb, grad)
-
                 p.add_(perturb, alpha=-step_size)
 
         return loss

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/adamw.py

@@ -107,8 +107,7 @@ class StableAdamW(BaseOptimizer):
                 exp_avg.lerp_(grad, weight=beta1_comp)
                 exp_avg_sq.mul_(beta2_hat).addcmul_(grad, grad, value=1.0 - beta2_hat)
 
-                rms = self.get_stable_adamw_rms(grad, exp_avg_sq, eps=eps_p2)
-                lr = group['lr'] / rms
+                lr: float = group['lr'] / self.get_stable_adamw_rms(grad, exp_avg_sq, eps=eps_p2)
 
                 self.apply_weight_decay(
                     p,

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/experimental/ranger25.py

@@ -11,6 +11,8 @@ from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMET
 class Ranger25(BaseOptimizer):
     r"""Mixin' every fancy optimizer hacks.
 
+    ADOPT + AdEMAMix + Cautious + StableAdamW + Adam-Atan2
+
     :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.
@@ -19,10 +21,10 @@ class Ranger25(BaseOptimizer):
     :param fixed_decay: bool. fix weight decay.
     :param alpha: float. usually between 4 and 10 would work well.
     :param t_alpha_beta3: Optional[float]. total number of iterations is preferred when needed.
-    :param n_sma_threshold: number of SMA threshold (recommended is 5).
     :param cautious: bool. whether to use the Cautious variant.
     :param stable_adamw: bool. whether to use stable AdamW variant.
-    :param eps: float. term added to the denominator to improve numerical stability.
+    :param eps: Optional[float]. term added to the denominator to improve numerical stability. when eps is None and
+        stable_adamw is False, adam-atan2 feature will be used.
     """
 
     def __init__(
@@ -35,10 +37,9 @@ class Ranger25(BaseOptimizer):
         fixed_decay: bool = False,
         alpha: float = 5.0,
         t_alpha_beta3: Optional[float] = None,
-        n_sma_threshold: int = 5,
         cautious: bool = True,
         stable_adamw: bool = True,
-        eps: float = 1e-8,
+        eps: Optional[float] = 1e-8,
         **kwargs,
     ):
         self.validate_learning_rate(lr)
@@ -48,9 +49,8 @@ class Ranger25(BaseOptimizer):
         self.validate_non_negative(weight_decay, 'weight_decay')
         self.validate_non_negative(eps, 'eps')
 
-        self.n_sma_threshold = n_sma_threshold
         self.cautious = cautious
-        self.stable_adamw = stable_adamw
+        self.stable_adamw: bool = stable_adamw if isinstance(eps, float) else False
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -60,7 +60,7 @@ class Ranger25(BaseOptimizer):
             'fixed_decay': fixed_decay,
             'alpha': alpha,
             't_alpha_beta3': t_alpha_beta3,
-            'eps': eps,
+            'eps': eps if (eps is not None) or (eps is None and not stable_adamw) else 1e-8,
         }
 
         super().__init__(params, defaults)
@@ -147,16 +147,14 @@ class Ranger25(BaseOptimizer):
 
                 exp_avg, exp_avg_sq, exp_avg_slow = state['exp_avg'], state['exp_avg_sq'], state['exp_avg_slow']
 
-                de_nom = exp_avg_sq.sqrt().clamp_(min=group['eps'])
-
-                normed_grad = grad.div(de_nom).clamp_(-clip, clip)
+                normed_grad = grad.div(
+                    exp_avg_sq.sqrt().clamp_(min=group['eps'] if group['eps'] is not None else 1e-8)
+                ).clamp_(-clip, clip)
 
                 exp_avg.mul_(beta1).add_(normed_grad, alpha=1.0 - beta1)
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                 exp_avg_slow.mul_(beta3_t).add_(normed_grad, alpha=1.0 - beta3_t)
 
-                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).add_(group['eps'])
-
                 update = exp_avg.clone()
                 if self.cautious:
                     self.apply_cautious(update, grad)
@@ -164,21 +162,14 @@ class Ranger25(BaseOptimizer):
                 if self.stable_adamw:
                     step_size /= self.get_stable_adamw_rms(grad, exp_avg_sq)
 
-                step_size, n_sma = self.get_rectify_step_size(
-                    is_rectify=True,
-                    step=group['step'],
-                    lr=step_size,
-                    beta2=beta2,
-                    n_sma_threshold=self.n_sma_threshold,
-                    degenerated_to_sgd=False,
-                )
+                update.add_(exp_avg_slow, alpha=alpha_t)
+
+                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq)
 
-                update.add_(exp_avg_slow, alpha=alpha_t).div_(de_nom)
+                if group['eps'] is not None:
+                    p.addcdiv_(update, de_nom.add_(group['eps']), value=-step_size)
+                    continue
 
-                if n_sma >= self.n_sma_threshold:
-                    de_nom = exp_avg_sq.sqrt().add_(group['eps'])
-                    p.addcdiv_(update, de_nom, value=-step_size)
-                else:
-                    p.add_(update, alpha=-step_size)
+                p.add_(update.atan2_(de_nom), alpha=-step_size)
 
         return loss

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/fp16.py

@@ -4,7 +4,7 @@ import torch
 from torch import nn
 from torch.optim import Optimizer
 
-from pytorch_optimizer.base.types import CLOSURE, OPTIMIZER, PARAMETERS
+from pytorch_optimizer.base.types import CLOSURE, PARAMETERS
 from pytorch_optimizer.optimizer.utils import clip_grad_norm, has_overflow
 
 
@@ -100,7 +100,7 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: Optimizer,
         aggregate_g_norms: bool = False,
         min_loss_scale: float = 2 ** -5,
     ):  # fmt: skip
@@ -113,10 +113,9 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
 
         # we want the optimizer to be tracking the fp32 parameters
         if len(optimizer.param_groups) != 1:
-            # future implementers: this should hopefully be a matter of just
-            # iterating through the param groups and keeping track of the pointer
-            # through the fp32_params
-            raise NotImplementedError('[-] Need to implement the parameter group transfer.')
+            # future implementers: this should hopefully be a matter of just iterating through the param groups and
+            # keeping track of the pointer through the fp32_params
+            raise NotImplementedError('Need to implement the parameter group transfer.')
 
         optimizer.param_groups[0]['params'] = self.fp32_params
 
@@ -124,17 +123,16 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
         self.needs_sync: bool = True
 
     @classmethod
-    def get_parameters(cls, optimizer: OPTIMIZER):
+    def get_parameters(cls, optimizer: Optimizer) -> List:
         params: List = []
-        for pg in optimizer.param_groups:
-            params += list(pg['params'])
+        for group in optimizer.param_groups:
+            params += list(group['params'])
         return params
 
     @classmethod
     def build_fp32_params(
         cls, parameters: PARAMETERS, flatten: bool = True
     ) -> Union[torch.Tensor, List[torch.Tensor]]:
-        # create FP32 copy of parameters and grads
         if flatten:
             total_param_size: int = sum(p.numel() for p in parameters)
             fp32_params = torch.zeros(total_param_size, dtype=torch.float, device=parameters[0].device)
@@ -142,7 +140,7 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
             offset: int = 0
             for p in parameters:
                 p_num_el = p.numel()
-                fp32_params[offset : offset + p_num_el].copy_(p.view(-1))
+                fp32_params[offset:offset + p_num_el].copy_(p.view(-1))  # fmt: skip
                 offset += p_num_el
 
             fp32_params = nn.Parameter(fp32_params)
@@ -168,9 +166,9 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
     def load_state_dict(self, state_dict: Dict):
         r"""Load an optimizer state dict.
 
-            In general, we should prefer the configuration of the existing optimizer instance
-            (e.g., learning rate) over that found in the state_dict. This allows us to
-            resume training from a checkpoint using a new set of optimizer args.
+            In general, we should prefer the configuration of the existing optimizer instance (e.g., learning rate)
+            over that found in the state_dict. This allows us to resume training from a checkpoint using a new set of
+            optimizer args.
 
         :param state_dict: Dict. state_dict.
         """
@@ -196,35 +194,34 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
         if update_main_grads:
             self.update_main_grads()
 
-    def sync_fp16_grads_to_fp32(self, multiply_grads: float = 1.0):
+    def sync_fp16_grads_to_fp32(self, multiply_grads: float = 1.0) -> None:
         r"""Sync fp16 to fp32 gradients."""
         if self.needs_sync:
             if self.scaler is not None:
-                # correct for dynamic loss scaler
                 multiply_grads /= self.scaler.loss_scale
 
-            # copy FP16 grads to FP32
-            for p, p32 in zip(self.fp16_params, self.fp32_params):
-                if not p.requires_grad:
+            for p16, p32 in zip(self.fp16_params, self.fp32_params):
+                if not p16.requires_grad:
                     continue
 
-                if p.grad is not None:
-                    p32.grad.copy_(p.grad)
+                if p16.grad is not None:
+                    p32.grad.copy_(p16.grad)
                     p32.grad.mul_(multiply_grads)
                 else:
-                    p32.grad = torch.zeros_like(p, dtype=torch.float)
+                    p32.grad = torch.zeros_like(p16, dtype=torch.float)
 
             self.needs_sync = False
 
-    def multiply_grads(self, c: float):
+    def multiply_grads(self, c: float) -> None:
         r"""Multiply grads by a constant c."""
         if self.needs_sync:
             self.sync_fp16_grads_to_fp32(c)
-        else:
-            for p32 in self.fp32_params:
-                p32.grad.mul_(c)
+            return
 
-    def update_main_grads(self):
+        for p32 in self.fp32_params:
+            p32.grad.mul_(c)
+
+    def update_main_grads(self) -> None:
         self.sync_fp16_grads_to_fp32()
 
     def clip_main_grads(self, max_norm: float):
@@ -249,8 +246,7 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
 
                     raise FloatingPointError(
                         f'Minimum loss scale reached ({self.min_loss_scale}). Your loss is probably exploding. '
-                        'Try lowering the learning rate, using gradient clipping or '
-                        'increasing the batch size.\n'
+                        'Try lowering the learning rate, using gradient clipping or increasing the batch size.\n'
                         f'Overflow: setting loss scale to {self.scaler.loss_scale}'
                     )
 
@@ -261,16 +257,15 @@ class SafeFP16Optimizer(Optimizer):  # pragma: no cover
         self.sync_fp16_grads_to_fp32()
         self.optimizer.step(closure)
 
-        # copy FP32 params back into FP16 model
-        for p, p32 in zip(self.fp16_params, self.fp32_params):
-            if not p.requires_grad:
+        for p16, p32 in zip(self.fp16_params, self.fp32_params):
+            if not p16.requires_grad:
                 continue
-            p.data.copy_(p32)
+            p16.data.copy_(p32)
 
-    def zero_grad(self):
+    def zero_grad(self) -> None:
         r"""Clear the gradients of all optimized parameters."""
-        for p in self.fp16_params:
-            p.grad = None
+        for p16 in self.fp16_params:
+            p16.grad = None
         for p32 in self.fp32_params:
             p32.grad.zero_()
         self.needs_sync = False

{pytorch_optimizer-3.3.3 → pytorch_optimizer-3.3.4}/pytorch_optimizer/optimizer/lookahead.py

@@ -2,15 +2,16 @@ from collections import defaultdict
 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, DEFAULTS, LOSS, OPTIMIZER, STATE
+from pytorch_optimizer.base.types import CLOSURE, DEFAULTS, LOSS, OPTIMIZER_INSTANCE_OR_CLASS, STATE
 
 
 class Lookahead(BaseOptimizer):
     r"""k steps forward, 1 step back.
 
-    :param optimizer: OPTIMIZER. base optimizer.
+    :param optimizer: OPTIMIZER_INSTANCE_OR_CLASS. base optimizer.
     :param k: int. number of lookahead steps.
     :param alpha: float. linear interpolation factor.
     :param pullback_momentum: str. change to inner optimizer momentum on interpolation update.
@@ -18,7 +19,7 @@ class Lookahead(BaseOptimizer):
 
     def __init__(
         self,
-        optimizer: OPTIMIZER,
+        optimizer: OPTIMIZER_INSTANCE_OR_CLASS,
         k: int = 5,
         alpha: float = 0.5,
         pullback_momentum: str = 'none',
@@ -28,10 +29,17 @@ class Lookahead(BaseOptimizer):
         self.validate_range(alpha, 'alpha', 0.0, 1.0)
         self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
 
+        if isinstance(optimizer, Optimizer):
+            self.optimizer = optimizer
+        elif 'params' in kwargs:
+            params = kwargs.pop('params')
+            self.optimizer = optimizer(params, **kwargs)
+        else:
+            raise ValueError('Need to pass `params` when you pass the torch.optim.Optimizer instance.')
+
         self._optimizer_step_pre_hooks: Dict[int, Callable] = {}
         self._optimizer_step_post_hooks: Dict[int, Callable] = {}
 
-        self.optimizer = optimizer
         self.alpha = alpha
         self.k = k
         self.pullback_momentum = pullback_momentum
@@ -53,7 +61,7 @@ class Lookahead(BaseOptimizer):
             'lookahead_alpha': alpha,
             'lookahead_k': k,
             'lookahead_pullback_momentum': pullback_momentum,
-            **optimizer.defaults,
+            **self.optimizer.defaults,
         }
 
     @property