pytorch-optimizer 2.9.0__tar.gz → 2.9.1__tar.gz

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.1
 Name: pytorch-optimizer
-Version: 2.9.0
-Summary: optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas.
+Version: 2.9.1
+Summary: optimizer & lr scheduler 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,Adai,AdamP,AdamS,Adan,AggMo,AliG,Apollo,AvaGrad,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DiffGrad,Fromage,Gravity,LARS,Lamb,Lion,MADGRAD,MSVAG,Nero,NovoGrad,PID,PNM,QHAdam,QHM,RAdam,Ranger,Ranger21,SGDP,SGDW,SM3,SRMM,SWATS,ScalableShampoo,Shampoo,Yogi,SAM,GSAM,PCGrad,RotoGrad
@@ -172,9 +172,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | Adai         | *Disentangling the Effects of Adaptive Learning Rate and Momentum*                                | `github <https://github.com/zeke-xie/adaptive-inertia-adai>`__                    | `https://arxiv.org/abs/2006.15815 <https://arxiv.org/abs/2006.15815>`__                       | `cite <https://github.com/zeke-xie/adaptive-inertia-adai#citing>`__                                                  |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                                        |
+| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                                        |
+| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | 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->`__     | `cite <https://github.com/juntang-zhuang/GSAM#citation>`__                                                           |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/README.rst

@@ -129,9 +129,9 @@ You can check the supported optimizers & lr schedulers.
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | Adai         | *Disentangling the Effects of Adaptive Learning Rate and Momentum*                                | `github <https://github.com/zeke-xie/adaptive-inertia-adai>`__                    | `https://arxiv.org/abs/2006.15815 <https://arxiv.org/abs/2006.15815>`__                       | `cite <https://github.com/zeke-xie/adaptive-inertia-adai#citing>`__                                                  |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                                        |
+| SAM          | *Sharpness-Aware Minimization*                                                                    | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2010.01412 <https://arxiv.org/abs/2010.01412>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2020arXiv201001412F/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
-| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__     | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                                        |
+| ASAM         | *Adaptive Sharpness-Aware Minimization*                                                           | `github <https://github.com/davda54/sam>`__                                       | `https://arxiv.org/abs/2102.11600 <https://arxiv.org/abs/2102.11600>`__                       | `cite <https://ui.adsabs.harvard.edu/abs/2021arXiv210211600K/exportcitation>`__                                      |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+
 | 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->`__     | `cite <https://github.com/juntang-zhuang/GSAM#citation>`__                                                           |
 +--------------+---------------------------------------------------------------------------------------------------+-----------------------------------------------------------------------------------+-----------------------------------------------------------------------------------------------+----------------------------------------------------------------------------------------------------------------------+

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "2.9.0"
-description = "optimizer & lr scheduler implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
+version = "2.9.1"
+description = "optimizer & lr scheduler collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
 maintainers = ["kozistr <kozistr@gmail.com>"]

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/base/optimizer.py

@@ -1,6 +1,6 @@
 import math
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 
@@ -90,31 +90,6 @@ class BaseOptimizer(ABC):
 
         return grad * exp_grad_norm / grad_norm if exp_grad_norm > grad_norm else grad
 
-    @staticmethod
-    def validate_learning_rate(learning_rate: Optional[float]):
-        if learning_rate is not None and learning_rate < 0.0:
-            raise NegativeLRError(learning_rate)
-
-    @staticmethod
-    def validate_betas(betas: BETAS):
-        if not 0.0 <= betas[0] <= 1.0:
-            raise ValueError(f'[-] beta1 {betas[0]} must be in the range [0, 1]')
-        if not 0.0 <= betas[1] <= 1.0:
-            raise ValueError(f'[-] beta2 {betas[1]} must be in the range [0, 1]')
-
-        if len(betas) < 3:
-            return
-
-        if not 0.0 <= betas[2] <= 1.0:
-            raise ValueError(f'[-] beta3 {betas[2]} must be in the range [0, 1]')
-
-    @staticmethod
-    def validate_weight_decay_type(weight_decay_type: str):
-        if weight_decay_type not in ('l2', 'decoupled', 'stable'):
-            raise ValueError(
-                f'[-] weight_decay_type {weight_decay_type} must be one of (\'l2\', \'decoupled\', \'stable\')'
-            )
-
     @staticmethod
     def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)'):
         if range_type == '[)' and not low <= x < high:
@@ -149,37 +124,31 @@ class BaseOptimizer(ABC):
             raise NegativeStepError(step, step_type=step_type)
 
     @staticmethod
-    def validate_pullback_momentum(pullback_momentum: str):
-        if pullback_momentum not in ('none', 'reset', 'pullback'):
-            raise ValueError(
-                f'[-] pullback_momentum {pullback_momentum} must be one of (\'none\' or \'reset\' or \'pullback\')'
-            )
+    def validate_options(x: str, name: str, options: List[str]):
+        if x not in options:
+            opts: str = ' or '.join([f'\'{option}\'' for option in options]).strip()
+            raise ValueError(f'[-] {name} {x} must be one of ({opts})')
 
     @staticmethod
-    def validate_reduction(reduction: str):
-        if reduction not in ('mean', 'sum'):
-            raise ValueError(f'[-] reduction {reduction} must be one of (\'mean\' or \'sum\')')
+    def validate_learning_rate(learning_rate: Optional[float]):
+        if learning_rate is not None and learning_rate < 0.0:
+            raise NegativeLRError(learning_rate)
 
-    @staticmethod
-    def validate_rebound(rebound: str):
-        if rebound not in ('constant', 'belief'):
-            raise ValueError(f'[-] rebound {rebound} must be one of (\'constant\' or \'belief\')')
+    def validate_betas(self, betas: BETAS):
+        self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')
+        self.validate_range(betas[1], 'beta2', 0.0, 1.0, range_type='[]')
 
-    @staticmethod
-    def validate_a2grad_variant(variant: str):
-        if variant not in ('uni', 'inc', 'exp'):
-            raise ValueError(f'[-] A2Grad variant {variant} must be one of (\'uni\' or \'inc\' or \'exp\')')
+        if len(betas) < 3:
+            return
 
-    @staticmethod
-    def validate_nus(nus: Union[float, Tuple[float, float]]):
+        self.validate_range(betas[2], 'beta3', 0.0, 1.0, range_type='[]')
+
+    def validate_nus(self, nus: Union[float, Tuple[float, float]]):
         if isinstance(nus, float):
-            if not 0.0 <= nus <= 1.0:
-                raise ValueError(f'[-] nus {nus} must be in the range [0, 1]')
+            self.validate_range(nus, 'nu', 0.0, 1.0, range_type='[]')
         else:
-            if not 0.0 <= nus[0] <= 1.0:
-                raise ValueError(f'[-] nus1 {nus[0]} must be in the range [0, 1]')
-            if not 0.0 <= nus[1] <= 1.0:
-                raise ValueError(f'[-] nus2 {nus[1]} must be in the range [0, 1]')
+            self.validate_range(nus[0], 'nu1', 0.0, 1.0, range_type='[]')
+            self.validate_range(nus[1], 'nu2', 0.0, 1.0, range_type='[]')
 
     @abstractmethod
     def reset(self):

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/a2grad.py

@@ -32,7 +32,7 @@ class A2Grad(Optimizer, BaseOptimizer):
         self.validate_learning_rate(lr)
         self.validate_non_negative(lips, 'lips')
         self.validate_non_negative(rho, 'rho')
-        self.validate_a2grad_variant(variant)
+        self.validate_options(variant, 'variant', ['uni', 'inc', 'exp'])
 
         self.variant = variant
 

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/adai.py

@@ -103,7 +103,7 @@ class Adai(Optimizer, BaseOptimizer):
                 state['step'] += 1
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 bias_correction2: float = 1.0 - beta2 ** state['step']
 

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/adamp.py

@@ -122,7 +122,7 @@ class AdamP(Optimizer, BaseOptimizer):
                         state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 s_grad = self.get_adanorm_gradient(
                     grad=grad,

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/adan.py

@@ -130,7 +130,7 @@ class Adan(Optimizer, BaseOptimizer):
                 grad.mul_(clip_global_grad_norm)
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 grad_diff = state['previous_grad']
                 grad_diff.add_(grad)

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/agc.py

@@ -6,7 +6,7 @@ from pytorch_optimizer.optimizer.utils import unit_norm
 def agc(p: torch.Tensor, grad: torch.Tensor, agc_eps: float, agc_clip_val: float, eps: float = 1e-6) -> torch.Tensor:
     r"""Clip gradient values in excess of the unit wise norm.
 
-    :param p: Parameter. parameter.
+    :param p: torch.Tensor. parameter.
     :param grad: torch.Tensor, gradient.
     :param agc_eps: float. agc epsilon to clip the norm of parameter.
     :param agc_clip_val: float. norm clip.

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/apollo.py

@@ -37,9 +37,9 @@ class Apollo(Optimizer, BaseOptimizer):
     ):
         self.validate_learning_rate(lr)
         self.validate_range(beta, 'beta', 0.0, 1.0, range_type='[]')
-        self.validate_rebound(rebound)
+        self.validate_options(rebound, 'rebound', ['constant', 'belief'])
         self.validate_non_negative(weight_decay, 'weight_decay')
-        self.validate_weight_decay_type(weight_decay_type)
+        self.validate_options(weight_decay_type, 'weight_decay_type', ['l2', 'decoupled', 'stable'])
         self.validate_non_negative(eps, 'eps')
 
         self.lr = lr

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/gc.py

@@ -1,14 +1,12 @@
 import torch
 
 
-def centralize_gradient(x: torch.Tensor, gc_conv_only: bool = False) -> torch.Tensor:
+def centralize_gradient(x: torch.Tensor, gc_conv_only: bool = False):
     r"""Gradient Centralization (GC).
 
     :param x: torch.Tensor. gradient.
     :param gc_conv_only: bool. 'False' for both conv & fc layers.
-    :return: torch.Tensor. centralized gradient.
     """
     size: int = x.dim()
     if (gc_conv_only and size > 3) or (not gc_conv_only and size > 1):
         x.add_(-x.mean(dim=tuple(range(1, size)), keepdim=True))
-    return x

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/lion.py

@@ -90,7 +90,7 @@ class Lion(Optimizer, BaseOptimizer):
                         state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
 
                 if self.use_gc:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 self.apply_weight_decay(
                     p=p,

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/lookahead.py

@@ -25,7 +25,7 @@ class Lookahead(BaseOptimizer):
     ):
         self.validate_positive(k, 'k')
         self.validate_range(alpha, 'alpha', 0.0, 1.0)
-        self.validate_pullback_momentum(pullback_momentum)
+        self.validate_options(pullback_momentum, 'pullback_momentum', ['none', 'reset', 'pullback'])
 
         self.alpha = alpha
         self.k = k

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/pcgrad.py

@@ -18,7 +18,7 @@ class PCGrad(BaseOptimizer):
     """
 
     def __init__(self, optimizer: OPTIMIZER, reduction: str = 'mean'):
-        self.validate_reduction(reduction)
+        self.validate_options(reduction, 'reduction', ['mean', 'sum'])
 
         self.optimizer = optimizer
         self.reduction = reduction

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/ranger.py

@@ -140,7 +140,7 @@ class Ranger(Optimizer, BaseOptimizer):
                         state['exp_grad_norm'] = torch.zeros((1,), dtype=grad.dtype, device=grad.device)
 
                 if self.use_gc and grad.dim() > self.gc_gradient_threshold:
-                    grad = centralize_gradient(grad, gc_conv_only=False)
+                    centralize_gradient(grad, gc_conv_only=False)
 
                 self.apply_weight_decay(
                     p=p,

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/ranger21.py

@@ -96,12 +96,9 @@ class Ranger21(Optimizer, BaseOptimizer):
         self.lookahead_blending_alpha = lookahead_blending_alpha
         self.norm_loss_factor = norm_loss_factor
 
-        # lookahead
         self.lookahead_step: int = 0
-
-        # learning rate
-        self.starting_lr = lr
-        self.current_lr = lr
+        self.starting_lr: float = lr
+        self.current_lr: float = lr
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -114,7 +111,6 @@ class Ranger21(Optimizer, BaseOptimizer):
         }
         super().__init__(params, defaults)
 
-        # warmup iterations
         self.num_warm_up_iterations: int = (
             self.build_warm_up_iterations(num_iterations, betas[1])
             if num_warm_up_iterations is None
@@ -140,8 +136,7 @@ class Ranger21(Optimizer, BaseOptimizer):
 
                 state['grad_ma'] = torch.zeros_like(p)
                 state['variance_ma'] = torch.zeros_like(p)
-                state['lookahead_params'] = torch.empty_like(p)
-                state['lookahead_params'].copy_(p)
+                state['lookahead_params'] = p.clone()
                 state['neg_grad_ma'] = torch.zeros_like(p)
                 state['max_variance_ma'] = torch.zeros_like(p)
 
@@ -162,28 +157,21 @@ class Ranger21(Optimizer, BaseOptimizer):
 
         warm_up_current_pct: float = min(1.0, (step / self.num_warm_up_iterations))
 
-        new_lr: float = lr * warm_up_current_pct
-        self.current_lr = new_lr
+        self.current_lr = lr * warm_up_current_pct
 
-        return new_lr
+        return self.current_lr
 
     def warm_down(self, lr: float, iteration: int) -> float:
         if iteration < self.start_warm_down:
             return lr
 
         # start iteration from 1, not 0
-        warm_down_iteration: int = (iteration + 1) - self.start_warm_down
-        warm_down_iteration = max(warm_down_iteration, 1)
-
-        warm_down_pct: float = warm_down_iteration / (self.num_warm_down_iterations + 1)
-        warm_down_pct = min(warm_down_pct, 1.0)
-
-        new_lr: float = self.starting_lr - self.warm_down_lr_delta * warm_down_pct
-        new_lr = max(new_lr, self.min_lr)
+        warm_down_iteration: int = max((iteration + 1) - self.start_warm_down, 1)
+        warm_down_pct: float = min(warm_down_iteration / (self.num_warm_down_iterations + 1), 1.0)
 
-        self.current_lr = new_lr
+        self.current_lr = max(self.starting_lr - self.warm_down_lr_delta * warm_down_pct, self.min_lr)
 
-        return new_lr
+        return self.current_lr
 
     @torch.no_grad()
     def step(self, closure: CLOSURE = None) -> LOSS:
@@ -220,8 +208,7 @@ class Ranger21(Optimizer, BaseOptimizer):
                 if len(state) == 0:
                     state['grad_ma'] = torch.zeros_like(p)
                     state['variance_ma'] = torch.zeros_like(p)
-                    state['lookahead_params'] = torch.empty_like(p)
-                    state['lookahead_params'].copy_(p)
+                    state['lookahead_params'] = p.clone()
                     state['neg_grad_ma'] = torch.zeros_like(p)
                     state['max_variance_ma'] = torch.zeros_like(p)
 
@@ -229,8 +216,8 @@ class Ranger21(Optimizer, BaseOptimizer):
                 grad.copy_(agc(p, grad, self.agc_eps, self.agc_clipping_value))
 
                 # Apply gradient centralization & normalization
-                grad = centralize_gradient(grad, gc_conv_only=False)
-                grad = normalize_gradient(grad)
+                centralize_gradient(grad, gc_conv_only=False)
+                normalize_gradient(grad)
 
                 # second moment estimation
                 # using positive-negative momentum and bias correction
@@ -245,7 +232,6 @@ class Ranger21(Optimizer, BaseOptimizer):
 
         # Phase 2 - Apply weight decay and step
         for group in self.param_groups:
-            lr: float = group['lr']
             beta1, beta2 = group['betas']
 
             bias_correction1: float = 1.0 - beta1 ** group['step']  # fmt: skip
@@ -254,7 +240,7 @@ class Ranger21(Optimizer, BaseOptimizer):
             noise_norm: float = math.sqrt((1.0 + beta2) ** 2 + beta2 ** 2)  # fmt: skip
 
             # warm up & down
-            lr = self.warm_up_dampening(lr, group['step'])
+            lr: float = self.warm_up_dampening(group['lr'], group['step'])
             lr = self.warm_down(lr, group['step'])
 
             for p in group['params']:
@@ -265,7 +251,7 @@ class Ranger21(Optimizer, BaseOptimizer):
                 self.apply_weight_decay(
                     p=p,
                     grad=None,
-                    lr=group['lr'],
+                    lr=lr,
                     weight_decay=group['weight_decay'],
                     weight_decouple=group['weight_decouple'],
                     fixed_decay=group['fixed_decay'],
@@ -287,16 +273,16 @@ class Ranger21(Optimizer, BaseOptimizer):
 
                 de_nom = (variance_ma.sqrt() / bias_correction2_sq).add_(group['eps'])
 
+                if self.use_softplus:
+                    de_nom = f.softplus(de_nom, beta=self.beta_softplus)
+
                 grad = p.grad
-                grad = centralize_gradient(grad, gc_conv_only=False)
-                grad = normalize_gradient(grad)
+                centralize_gradient(grad, gc_conv_only=False)
+                normalize_gradient(grad)
 
                 grad_ma.mul_(beta1 ** 2).add_(grad, alpha=1.0 - beta1 ** 2)  # fmt: skip
 
-                step_size: float = lr if group['adam_debias'] else lr / bias_correction1
-
-                if self.use_softplus:
-                    de_nom = f.softplus(de_nom, beta=self.beta_softplus)
+                step_size: float = self.apply_adam_debias(group['adam_debias'], lr, bias_correction1)
 
                 pn_momentum = grad_ma.mul(1.0 + 1.0).add(neg_grad_ma, alpha=-1.0).mul(1.0 / noise_norm)
                 p.addcdiv_(pn_momentum, de_nom, value=-step_size)

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/shampoo.py

@@ -45,7 +45,6 @@ class Shampoo(Optimizer, BaseOptimizer):
         self.validate_non_negative(matrix_eps, 'matrix_eps')
 
         self.preconditioning_compute_steps = preconditioning_compute_steps
-        self.matrix_eps = matrix_eps
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -53,6 +52,7 @@ class Shampoo(Optimizer, BaseOptimizer):
             'weight_decay': weight_decay,
             'weight_decouple': weight_decouple,
             'fixed_decay': fixed_decay,
+            'matrix_eps': matrix_eps,
         }
         super().__init__(params, defaults)
 
@@ -93,7 +93,7 @@ class Shampoo(Optimizer, BaseOptimizer):
                         state['momentum_buffer'] = grad.clone()
 
                     for dim_id, dim in enumerate(grad.size()):
-                        state[f'pre_cond_{dim_id}'] = self.matrix_eps * torch.eye(dim, out=grad.new(dim, dim))
+                        state[f'pre_cond_{dim_id}'] = group['matrix_eps'] * torch.eye(dim, out=grad.new(dim, dim))
                         state[f'inv_pre_cond_{dim_id}'] = grad.new(dim, dim).zero_()
 
                 if momentum > 0.0:

{pytorch_optimizer-2.9.0 → pytorch_optimizer-2.9.1}/pytorch_optimizer/optimizer/utils.py

@@ -27,13 +27,12 @@ def to_real(x: torch.Tensor) -> torch.Tensor:
     return x.real if torch.is_complex(x) else x
 
 
-def normalize_gradient(x: torch.Tensor, use_channels: bool = False, epsilon: float = 1e-8) -> torch.Tensor:
+def normalize_gradient(x: torch.Tensor, use_channels: bool = False, epsilon: float = 1e-8):
     r"""Normalize gradient with stddev.
 
     :param x: torch.Tensor. gradient.
     :param use_channels: bool. channel-wise normalization.
     :param epsilon: float. eps.
-    :return: torch.Tensor. normalized gradient.
     """
     size: int = x.dim()
     if size > 1 and use_channels:
@@ -42,7 +41,6 @@ def normalize_gradient(x: torch.Tensor, use_channels: bool = False, epsilon: flo
     elif torch.numel(x) > 2:
         s = x.std().add_(epsilon)
         x.div_(s)
-    return x
 
 
 def flatten_grad(grads: List[torch.Tensor]) -> torch.Tensor: