pytorch-optimizer 3.0.0__tar.gz → 3.0.1__tar.gz

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: pytorch_optimizer
-Version: 3.0.0
+Version: 3.0.1
 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,Adalite,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,Fromage,GaLore,Gravity,GSAM,LARS,Lamb,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,Tiger,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes
+Keywords: pytorch,deep-learning,optimizer,lr scheduler,A2Grad,ASGD,AccSGD,AdaBelief,AdaBound,AdaDelta,AdaFactor,AdaMax,AdaMod,AdaNorm,AdaPNM,AdaSmooth,AdaHessian,Adai,Adalite,AdamP,AdamS,Adan,AggMo,Aida,AliG,Amos,Apollo,AvaGrad,bSAM,CAME,DAdaptAdaGrad,DAdaptAdam,DAdaptAdan,DAdaptSGD,DAdaptLion,DiffGrad,FAdam,Fromage,GaLore,Gravity,GrokFast,GSAM,LARS,Lamb,Lion,LOMO,Lookahead,MADGRAD,MSVAG,Nero,NovoGrad,PAdam,PCGrad,PID,PNM,Prodigy,QHAdam,QHM,RAdam,Ranger,Ranger21,RotoGrad,SAM,ScheduleFreeSGD,ScheduleFreeAdamW,SGDP,Shampoo,ScalableShampoo,SGDW,SignSGD,SM3,SopihaH,SRMM,SWATS,Tiger,WSAM,Yogi,BCE,BCEFocal,Focal,FocalCosine,SoftF1,Dice,LDAM,Jaccard,Bi-Tempered,Tversky,FocalTversky,LovaszHinge,bitsandbytes
 Author: kozistr
 Author-email: kozistr@gmail.com
 Maintainer: kozistr
@@ -49,7 +49,7 @@ Description-Content-Type: text/markdown
 
 **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, **67 optimizers (+ `bitsandbytes`)**, **11 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **69 optimizers (+ `bitsandbytes`)**, **11 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -203,6 +203,8 @@ supported_optimizers = get_supported_optimizers()
 | Adalite       | *Adalite optimizer*                                                                               | [github](https://github.com/VatsaDev/adalite)                                                                  | <https://github.com/VatsaDev/adalite>                                                      | [cite](https://github.com/VatsaDev/adalite)                                                                       |
 | bSAM          | *SAM as an Optimal Relaxation of Bayes*                                                           | [github](https://github.com/team-approx-bayes/bayesian-sam)                                                    | <https://arxiv.org/abs/2210.01620>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv221001620M/exportcitation)                                      |
 | Schedule-Free | *Schedule-Free Optimizers*                                                                        | [github](https://github.com/facebookresearch/schedule_free)                                                    | <https://github.com/facebookresearch/schedule_free>                                        | [cite](https://github.com/facebookresearch/schedule_free)                                                         |
+| FAdam         | *Adam is a natural gradient optimizer using diagonal empirical Fisher information*                | [github](https://github.com/lessw2020/fadam_pytorch)                                                           | <https://arxiv.org/abs/2405.12807>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240512807H/exportcitation)                                      |
+| Grokfast      | *Accelerated Grokking by Amplifying Slow Gradients*                                               | [github](https://github.com/ironjr/grokfast)                                                                   | <https://arxiv.org/abs/2405.20233>                                                         | [cite](https://github.com/ironjr/grokfast?tab=readme-ov-file#citation)                                            |
 
 ## Supported LR Scheduler
 
@@ -363,7 +365,7 @@ If you use this software, please cite it below. Or you can get it from "cite thi
         month = jan,
         title = {{pytorch_optimizer: optimizer & lr scheduler & loss function collections in PyTorch}},
         url = {https://github.com/kozistr/pytorch_optimizer},
-        version = {2.12.0},
+        version = {3.0.1},
         year = {2021}
     }
 

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/README.md

@@ -10,7 +10,7 @@
 
 **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, **67 optimizers (+ `bitsandbytes`)**, **11 lr schedulers**, and **13 loss functions** are supported!  
+Currently, **69 optimizers (+ `bitsandbytes`)**, **11 lr schedulers**, and **13 loss functions** are supported!  
 
 Highly inspired by [pytorch-optimizer](https://github.com/jettify/pytorch-optimizer).
 
@@ -164,6 +164,8 @@ supported_optimizers = get_supported_optimizers()
 | Adalite       | *Adalite optimizer*                                                                               | [github](https://github.com/VatsaDev/adalite)                                                                  | <https://github.com/VatsaDev/adalite>                                                      | [cite](https://github.com/VatsaDev/adalite)                                                                       |
 | bSAM          | *SAM as an Optimal Relaxation of Bayes*                                                           | [github](https://github.com/team-approx-bayes/bayesian-sam)                                                    | <https://arxiv.org/abs/2210.01620>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2022arXiv221001620M/exportcitation)                                      |
 | Schedule-Free | *Schedule-Free Optimizers*                                                                        | [github](https://github.com/facebookresearch/schedule_free)                                                    | <https://github.com/facebookresearch/schedule_free>                                        | [cite](https://github.com/facebookresearch/schedule_free)                                                         |
+| FAdam         | *Adam is a natural gradient optimizer using diagonal empirical Fisher information*                | [github](https://github.com/lessw2020/fadam_pytorch)                                                           | <https://arxiv.org/abs/2405.12807>                                                         | [cite](https://ui.adsabs.harvard.edu/abs/2024arXiv240512807H/exportcitation)                                      |
+| Grokfast      | *Accelerated Grokking by Amplifying Slow Gradients*                                               | [github](https://github.com/ironjr/grokfast)                                                                   | <https://arxiv.org/abs/2405.20233>                                                         | [cite](https://github.com/ironjr/grokfast?tab=readme-ov-file#citation)                                            |
 
 ## Supported LR Scheduler
 
@@ -324,7 +326,7 @@ If you use this software, please cite it below. Or you can get it from "cite thi
         month = jan,
         title = {{pytorch_optimizer: optimizer & lr scheduler & loss function collections in PyTorch}},
         url = {https://github.com/kozistr/pytorch_optimizer},
-        version = {2.12.0},
+        version = {3.0.1},
         year = {2021}
     }
 

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "3.0.0"
+version = "3.0.1"
 description = "optimizer & lr scheduler & objective function collections in PyTorch"
 license = "Apache-2.0"
 authors = ["kozistr <kozistr@gmail.com>"]
@@ -13,12 +13,12 @@ keywords = [
     "pytorch", "deep-learning", "optimizer", "lr scheduler", "A2Grad", "ASGD", "AccSGD", "AdaBelief", "AdaBound",
     "AdaDelta", "AdaFactor", "AdaMax", "AdaMod", "AdaNorm", "AdaPNM", "AdaSmooth", "AdaHessian", "Adai", "Adalite",
     "AdamP", "AdamS", "Adan", "AggMo", "Aida", "AliG", "Amos", "Apollo", "AvaGrad", "bSAM", "CAME", "DAdaptAdaGrad",
-    "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad", "Fromage", "GaLore", "Gravity", "GSAM", "LARS",
-    "Lamb", "Lion", "LOMO", "Lookahead", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PAdam", "PCGrad", "PID", "PNM",
-    "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM", "ScheduleFreeSGD",
-    "ScheduleFreeAdamW", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH", "SRMM", "SWATS",
-    "Tiger", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM", "Jaccard",
-    "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes",
+    "DAdaptAdam", "DAdaptAdan", "DAdaptSGD", "DAdaptLion", "DiffGrad", "FAdam", "Fromage", "GaLore", "Gravity",
+    "GrokFast", "GSAM", "LARS", "Lamb", "Lion", "LOMO", "Lookahead", "MADGRAD", "MSVAG", "Nero", "NovoGrad", "PAdam",
+    "PCGrad", "PID", "PNM", "Prodigy", "QHAdam", "QHM", "RAdam", "Ranger", "Ranger21", "RotoGrad", "SAM",
+    "ScheduleFreeSGD", "ScheduleFreeAdamW", "SGDP", "Shampoo", "ScalableShampoo", "SGDW", "SignSGD", "SM3", "SopihaH",
+    "SRMM", "SWATS", "Tiger", "WSAM", "Yogi", "BCE", "BCEFocal", "Focal", "FocalCosine", "SoftF1", "Dice", "LDAM",
+    "Jaccard", "Bi-Tempered", "Tversky", "FocalTversky", "LovaszHinge", "bitsandbytes",
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/__init__.py

@@ -54,11 +54,13 @@ from pytorch_optimizer.optimizer.avagrad import AvaGrad
 from pytorch_optimizer.optimizer.came import CAME
 from pytorch_optimizer.optimizer.dadapt import DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptLion, DAdaptSGD
 from pytorch_optimizer.optimizer.diffgrad import DiffGrad
+from pytorch_optimizer.optimizer.fadam import FAdam
 from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
 from pytorch_optimizer.optimizer.fromage import Fromage
 from pytorch_optimizer.optimizer.galore import GaLore, GaLoreProjector
 from pytorch_optimizer.optimizer.gc import centralize_gradient
 from pytorch_optimizer.optimizer.gravity import Gravity
+from pytorch_optimizer.optimizer.grokfast import GrokFastAdamW, gradfilter_ema, gradfilter_ma
 from pytorch_optimizer.optimizer.lamb import Lamb
 from pytorch_optimizer.optimizer.lars import LARS
 from pytorch_optimizer.optimizer.lion import Lion
@@ -190,6 +192,8 @@ OPTIMIZER_LIST: List[OPTIMIZER] = [
     BSAM,
     ScheduleFreeSGD,
     ScheduleFreeAdamW,
+    FAdam,
+    GrokFastAdamW,
 ]
 OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/base/optimizer.py

@@ -215,7 +215,7 @@ class BaseOptimizer(ABC):
         return grad * exp_grad_norm / grad_norm if exp_grad_norm > grad_norm else grad
 
     @staticmethod
-    def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)'):
+    def validate_range(x: float, name: str, low: float, high: float, range_type: str = '[)') -> None:
         if range_type == '[)' and not low <= x < high:
             raise ValueError(f'[-] {name} must be in the range [{low}, {high})')
         if range_type == '[]' and not low <= x <= high:
@@ -226,40 +226,42 @@ class BaseOptimizer(ABC):
             raise ValueError(f'[-] {name} must be in the range ({low}, {high})')
 
     @staticmethod
-    def validate_non_negative(x: Optional[float], name: str):
+    def validate_non_negative(x: Optional[float], name: str) -> None:
         if x is not None and x < 0.0:
             raise ValueError(f'[-] {name} must be non-negative')
 
     @staticmethod
-    def validate_positive(x: Union[float, int], name: str):
+    def validate_positive(x: Union[float, int], name: str) -> None:
         if x <= 0:
             raise ValueError(f'[-] {name} must be positive')
 
     @staticmethod
-    def validate_boundary(constant: float, boundary: float, bound_type: str = 'upper'):
+    def validate_boundary(constant: float, boundary: float, bound_type: str = 'upper') -> None:
         if bound_type == 'upper' and constant > boundary:
             raise ValueError(f'[-] constant {constant} must be in a range of (-inf, {boundary}]')
         if bound_type == 'lower' and constant < boundary:
             raise ValueError(f'[-] constant {constant} must be in a range of [{boundary}, inf)')
 
     @staticmethod
-    def validate_step(step: int, step_type: str):
+    def validate_step(step: int, step_type: str) -> None:
         if step < 1:
             raise NegativeStepError(step, step_type=step_type)
 
     @staticmethod
-    def validate_options(x: str, name: str, options: List[str]):
+    def validate_options(x: str, name: str, options: List[str]) -> None:
         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_learning_rate(learning_rate: Optional[float]):
+    def validate_learning_rate(learning_rate: Optional[float]) -> None:
         if learning_rate is not None and learning_rate < 0.0:
             raise NegativeLRError(learning_rate)
 
-    def validate_betas(self, betas: BETAS):
-        self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')
+    def validate_betas(self, betas: BETAS) -> None:
+        if betas[0] is not None:
+            self.validate_range(betas[0], 'beta1', 0.0, 1.0, range_type='[]')
+
         self.validate_range(betas[1], 'beta2', 0.0, 1.0, range_type='[]')
 
         if len(betas) < 3:
@@ -268,7 +270,7 @@ class BaseOptimizer(ABC):
         if betas[2] is not None:
             self.validate_range(betas[2], 'beta3', 0.0, 1.0, range_type='[]')
 
-    def validate_nus(self, nus: Union[float, Tuple[float, float]]):
+    def validate_nus(self, nus: Union[float, Tuple[float, float]]) -> None:
         if isinstance(nus, float):
             self.validate_range(nus, 'nu', 0.0, 1.0, range_type='[]')
         else:

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/base/types.py

@@ -6,7 +6,7 @@ from torch.optim.lr_scheduler import _LRScheduler
 
 CLOSURE = Optional[Callable[[], float]]
 LOSS = Optional[float]
-BETAS = Union[Tuple[float, float], Tuple[float, float, float]]
+BETAS = Union[Tuple[float, float], Tuple[float, float, float], Tuple[None, float]]
 DEFAULTS = Dict
 PARAMETERS = Optional[Union[Iterable[Dict], Iterable[torch.Tensor]]]
 STATE = Dict

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/optimizer/adafactor.py

@@ -10,11 +10,12 @@ from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMET
 
 
 class AdaFactor(Optimizer, BaseOptimizer):
-    r"""Adaptive Learning Rates with Sublinear Memory Cost.
+    r"""Adaptive Learning Rates with Sublinear Memory Cost with some tweaks.
 
     :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 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.
@@ -27,6 +28,9 @@ class AdaFactor(Optimizer, BaseOptimizer):
         is being used.
     :param eps1: float. term added to the denominator to improve numerical stability.
     :param eps2: float. term added to the denominator to improve numerical stability.
+    :param momentum_dtype: torch.dtype. type of momentum variable. In VIT paper observed that storing momentum in
+        half-precision (bfloat16 type) does not affect training dynamics and has no effect on the outcome while
+        reducing optimize overhead from 2-fold to 1.5-fold.
     """
 
     def __init__(
@@ -45,6 +49,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
         warmup_init: bool = False,
         eps1: float = 1e-30,
         eps2: float = 1e-3,
+        momentum_dtype: torch.dtype = torch.bfloat16,
     ):
         self.validate_learning_rate(lr)
         self.validate_betas(betas)
@@ -56,6 +61,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
         self.clip_threshold = clip_threshold
         self.eps1 = eps1
         self.eps2 = eps2
+        self.momentum_dtype = momentum_dtype
 
         defaults: DEFAULTS = {
             'lr': lr,
@@ -87,7 +93,8 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 grad_shape: Tuple[int, ...] = grad.shape
                 factored: bool = self.get_options(grad_shape)
 
-                state['exp_avg'] = torch.zeros_like(p)
+                if group['betas'][0] is not None:
+                    state['exp_avg'] = torch.zeros_like(p, dtype=self.momentum_dtype)
 
                 if factored:
                     state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype, device=grad.device)
@@ -149,7 +156,7 @@ class AdaFactor(Optimizer, BaseOptimizer):
             else:
                 group['step'] = 1
 
-            beta1, _ = group['betas']
+            beta1, beta2 = group['betas']
 
             beta2_t: float = 1.0 - math.pow(group['step'], self.decay_rate)
 
@@ -167,7 +174,8 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 factored: bool = self.get_options(grad_shape)
 
                 if len(state) == 0:
-                    state['exp_avg'] = torch.zeros_like(p)
+                    if beta1 is not None:
+                        state['exp_avg'] = torch.zeros_like(p, dtype=self.momentum_dtype)
 
                     if factored:
                         state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1], dtype=grad.dtype, device=grad.device)
@@ -205,6 +213,8 @@ class AdaFactor(Optimizer, BaseOptimizer):
                 else:
                     exp_avg_sq = state['exp_avg_sq']
                     exp_avg_sq.mul_(beta2_t).add_(update, alpha=1.0 - beta2_t)
+                    exp_avg_sq.clamp_(max=beta2)
+
                     torch.rsqrt(exp_avg_sq, out=update)
 
                 if group['ams_bound']:
@@ -216,8 +226,11 @@ class AdaFactor(Optimizer, BaseOptimizer):
 
                 update.div_((self.get_rms(update) / self.clip_threshold).clamp_(min=1.0)).mul_(lr)
 
-                exp_avg = state['exp_avg']
-                exp_avg.mul_(beta1).add_(update, alpha=1.0 - beta1)
+                if beta1 is not None:
+                    exp_avg = state['exp_avg']
+                    exp_avg.mul_(beta1).add_(update, alpha=1.0 - beta1)
+
+                    update = exp_avg
 
                 self.apply_weight_decay(
                     p=p,
@@ -228,6 +241,6 @@ class AdaFactor(Optimizer, BaseOptimizer):
                     fixed_decay=group['fixed_decay'],
                 )
 
-                p.add_(-exp_avg)
+                p.add_(-update)
 
         return loss

pytorch_optimizer-3.0.1/pytorch_optimizer/optimizer/fadam.py

@@ -0,0 +1,124 @@
+import torch
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+
+class FAdam(Optimizer, BaseOptimizer):
+    r"""Adam is a natural gradient optimizer using diagonal empirical Fisher information.
+
+    :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 weight_decay: float. weight decay (L2 penalty).
+    :param clip: float. maximum norm of the gradient.
+    :param p: float. momentum factor.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    :param momentum_dtype: torch.dtype. type of momentum.
+    :param fim_dtype: torch.dtype. type of fim.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-3,
+        betas: BETAS = (0.9, 0.999),
+        weight_decay: float = 0.1,
+        clip: float = 1.0,
+        p: float = 0.5,
+        eps: float = 1e-8,
+        momentum_dtype: torch.dtype = torch.float32,
+        fim_dtype: torch.dtype = torch.float32,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_positive(clip, 'clip')
+        self.validate_positive(p, 'p')
+        self.validate_non_negative(eps, 'eps')
+
+        self.momentum_dtype = momentum_dtype
+        self.fim_dtype = fim_dtype
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'clip': clip,
+            'p': p,
+            'eps': eps,
+        }
+
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'FAdam'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['momentum'] = torch.zeros_like(p, dtype=self.momentum_dtype)
+                state['fim'] = torch.zeros_like(p, dtype=self.fim_dtype)
+
+    @torch.no_grad()
+    def step(self, closure: CLOSURE = None) -> LOSS:
+        loss: LOSS = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            curr_beta2: float = beta2 * (1 - beta2 ** (group['step'] - 1)) / (1 - beta2 ** group['step'])
+
+            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['momentum'] = torch.zeros_like(p, dtype=self.momentum_dtype)
+                    state['fim'] = torch.zeros_like(p, dtype=self.fim_dtype)
+
+                momentum, fim = state['momentum'], state['fim']
+                fim.mul_(curr_beta2).addcmul_(grad, grad, value=1.0 - curr_beta2)
+
+                rms_grad = grad.pow(2).mean().sqrt_()
+                curr_eps = min(rms_grad, 1) * group['eps']
+
+                fim_base = fim.pow(group['p']).add_(curr_eps)
+                grad_nat = grad / fim_base
+
+                rms = grad_nat.pow(2).mean().sqrt_()
+                divisor = max(1, rms) / group['clip']
+                grad_nat.div_(divisor)
+
+                momentum.mul_(beta1).add_(grad_nat, alpha=1.0 - beta1)
+
+                grad_weights = p / fim_base
+
+                rms = torch.pow(grad_weights, 2).mean().sqrt_()
+                divisor = max(1, rms) / group['clip']
+                grad_weights.div_(divisor)
+
+                grad_weights.mul_(group['weight_decay']).add_(momentum)
+
+                p.add_(grad_weights, alpha=-group['lr'])
+
+        return loss

pytorch_optimizer-3.0.1/pytorch_optimizer/optimizer/grokfast.py

@@ -0,0 +1,232 @@
+import math
+from collections import deque
+from typing import Dict, Literal, Optional
+
+import torch
+from torch import nn
+from torch.optim.optimizer import Optimizer
+
+from pytorch_optimizer.base.exception import NoSparseGradientError
+from pytorch_optimizer.base.optimizer import BaseOptimizer
+from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
+
+FILTER_TYPE = Literal['mean', 'sum']
+
+
+@torch.no_grad()
+def gradfilter_ma(
+    model: nn.Module,
+    grads: Optional[Dict[str, deque]] = None,
+    window_size: int = 100,
+    lamb: float = 5.0,
+    filter_type: FILTER_TYPE = 'mean',
+    warmup: bool = True,
+) -> Dict[str, deque]:
+    r"""Grokfast-MA.
+
+    Example:
+    -------
+        Here's an example::
+
+            loss.backwards()  # Calculate the gradients.
+
+            grads = gradfilter_ma(model, grads=grads, window_size=window_size, lamb=lamb)
+
+            optimizer.step()  # Call the optimizer.
+
+    :param model: nn.Module. model that contains every trainable parameters.
+    :param grads: Optional[Dict[str, deque]]. running memory (Queue for windowed moving average). initialize by setting
+        it to None. feed the output of the method recursively after on.
+    :param window_size: int. the width of the filter window. additional memory requirements increases linearly with
+        respect to the windows size.
+    :param lamb: float. amplifying factor hyperparameter of the filter.
+    :param filter_type: FILTER_TYPE. aggregation method for the running queue.
+    :param warmup: bool. if true, filter is not applied until the queue is filled.
+    """
+    if grads is None:
+        grads = {n: deque(maxlen=window_size) for n, p in model.named_parameters() if p.requires_grad}
+
+    for n, p in model.named_parameters():
+        if p.requires_grad:
+            grads[n].append(p.grad)
+
+            if not warmup or len(grads[n]) == window_size:
+                if filter_type == 'mean':
+                    avg = sum(grads[n]) / len(grads[n])
+                elif filter_type == 'sum':
+                    avg = sum(grads[n])
+                else:
+                    raise ValueError(f'not supported filter_type {filter_type}')
+
+                p.grad.add_(avg, alpha=lamb)
+
+    return grads
+
+
+@torch.no_grad()
+def gradfilter_ema(
+    model: nn.Module,
+    grads: Optional[Dict[str, torch.Tensor]] = None,
+    alpha: float = 0.98,
+    lamb: float = 2.0,
+) -> Dict[str, torch.Tensor]:
+    r"""Grokfast.
+
+    Example:
+    -------
+        Here's an example::
+
+            loss.backwards()  # Calculate the gradients.
+
+            grads = gradfilter_ema(model, grads=grads, alpha=alpha, lamb=lamb)
+
+            optimizer.step()  # Call the optimizer.
+
+    :param model: nn.Module. model that contains every trainable parameters.
+    :param grads: Optional[Dict[str, deque]]. running memory (EMA). Initialize by setting it to None. Feed the output
+        of the method recursively after on.
+    :param alpha: int. momentum hyperparameter of the EMA.
+    :param lamb: float. amplifying factor hyperparameter of the filter.
+    """
+    if grads is None:
+        grads = {n: p.grad for n, p in model.named_parameters() if p.requires_grad}
+
+    for n, p in model.named_parameters():
+        if p.requires_grad:
+            grads[n].mul_(alpha).add_(p.grad, alpha=1.0 - alpha)
+            p.grad.add_(grads[n], alpha=lamb)
+
+    return grads
+
+
+class GrokFastAdamW(Optimizer, BaseOptimizer):
+    r"""Accelerated Grokking by Amplifying Slow Gradients with AdamW.
+
+    :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 grokfast: bool. whether to use grokfast.
+    :param grokfast_alpha: float. momentum hyperparameter of the EMA.
+    :param grokfast_lamb: float. amplifying factor hyperparameter of the filter..
+    :param grokfast_after_step: int. warmup step for grokfast.
+    :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.
+    :param eps: float. term added to the denominator to improve numerical stability.
+    """
+
+    def __init__(
+        self,
+        params: PARAMETERS,
+        lr: float = 1e-4,
+        betas: BETAS = (0.9, 0.99),
+        grokfast: bool = True,
+        grokfast_alpha: float = 0.98,
+        grokfast_lamb: float = 2.0,
+        grokfast_after_step: int = 0,
+        weight_decay: float = 0.0,
+        weight_decouple: bool = True,
+        fixed_decay: bool = False,
+        normalize_lr: bool = True,
+        eps: float = 1e-8,
+    ):
+        self.validate_learning_rate(lr)
+        self.validate_betas(betas)
+        self.validate_non_negative(weight_decay, 'weight_decay')
+        self.validate_range(grokfast_alpha, 'grokfast_alpha', 0.0, 1.0)
+        self.validate_non_negative(eps, 'eps')
+
+        if grokfast and normalize_lr:
+            lr /= 1.0 + grokfast_lamb
+
+        defaults: DEFAULTS = {
+            'lr': lr,
+            'betas': betas,
+            'weight_decay': weight_decay,
+            'weight_decouple': weight_decouple,
+            'fixed_decay': fixed_decay,
+            'grokfast': grokfast,
+            'grokfast_alpha': grokfast_alpha,
+            'grokfast_lamb': grokfast_lamb,
+            'grokfast_after_step': grokfast_after_step,
+            'eps': eps,
+        }
+        super().__init__(params, defaults)
+
+    def __str__(self) -> str:
+        return 'GrokFastAdamW'
+
+    @torch.no_grad()
+    def reset(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                state = self.state[p]
+
+                state['exp_avg'] = torch.zeros_like(p)
+                state['exp_avg_sq'] = 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()
+
+        for group in self.param_groups:
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            beta1, beta2 = group['betas']
+
+            bias_correction1: float = 1.0 - beta1 ** group['step']
+            bias_correction2_sq: float = math.sqrt(1.0 - beta2 ** group['step'])
+
+            should_grokfast: bool = (
+                group['grokfast'] and group['step'] > group['grokfast_after_step'] and group['grokfast_lamb'] > 0
+            )
+
+            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['exp_avg_sq'] = torch.zeros_like(p)
+                    if should_grokfast:
+                        state['grok_exp_avg'] = grad.clone()
+
+                self.apply_weight_decay(
+                    p=p,
+                    grad=grad,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    weight_decouple=group['weight_decouple'],
+                    fixed_decay=group['fixed_decay'],
+                )
+
+                if should_grokfast:
+                    grok_exp_avg = state['grok_exp_avg']
+                    grok_exp_avg.lerp_(grad, weight=1.0 - group['grokfast_alpha'])
+
+                    grad.add_(grok_exp_avg, alpha=group['grokfast_lamb'])
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+
+                de_nom = exp_avg_sq.sqrt().div_(bias_correction2_sq).clamp_(min=group['eps'])
+
+                update = exp_avg.div(bias_correction1).div_(de_nom)
+
+                p.add_(update, alpha=-group['lr'])
+
+        return loss

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/optimizer/lookahead.py

@@ -36,7 +36,6 @@ class Lookahead(Optimizer, BaseOptimizer):
         self.pullback_momentum = pullback_momentum
 
         self.optimizer = optimizer
-        self.param_groups = self.optimizer.param_groups
 
         self.state: STATE = defaultdict(dict)
 
@@ -58,6 +57,10 @@ class Lookahead(Optimizer, BaseOptimizer):
             **optimizer.defaults,
         }
 
+    @property
+    def param_groups(self):
+        return self.optimizer.param_groups
+
     def __getstate__(self):
         return {
             'state': self.state,

{pytorch_optimizer-3.0.0 → pytorch_optimizer-3.0.1}/pytorch_optimizer/optimizer/utils.py

@@ -282,7 +282,7 @@ def reduce_max_except_dim(x: torch.Tensor, dim: int) -> torch.Tensor:
 
 def reg_noise(
     network1: nn.Module, network2: nn.Module, num_data: int, lr: float, eta: float = 8e-3, temperature: float = 1e-4
-) -> torch.Tensor | float:
+) -> Union[torch.Tensor, float]:
     r"""Entropy-MCMC: Sampling from flat basins with ease.
 
     usage: https://github.com/lblaoke/EMCMC/blob/master/exp/cifar10_emcmc.py