feature: disable_lr_scheduler

Author: kozistr

pytorch_optimizer/optimizer/ranger21.py

@@ -2,7 +2,7 @@
 from typing import Optional
 
 import torch
-from torch.nn import functional as f
+from torch.nn.functional import softplus
 
 from pytorch_optimizer.base.exception import NoSparseGradientError, ZeroParameterSizeError
 from pytorch_optimizer.base.optimizer import BaseOptimizer
@@ -38,6 +38,7 @@ class Ranger21(BaseOptimizer):
     :param betas: BETAS. coefficients used for computing running averages of gradient and the squared hessian trace.
     :param use_softplus: bool. use softplus to smooth.
     :param beta_softplus: float. beta.
+    :param disable_lr_scheduler: bool. whether to disable learning rate schedule.
     :param num_warm_up_iterations: Optional[int]. number of warm-up iterations. Ranger21 performs linear learning rate
         warmup.
     :param num_warm_down_iterations: Optional[int]. number of warm-down iterations. Ranger21 performs Explore-exploit
@@ -65,6 +66,7 @@ def __init__(  # pylint: disable=R0913
         betas: BETAS = (0.9, 0.999),
         use_softplus: bool = True,
         beta_softplus: float = 50.0,
+        disable_lr_scheduler: bool = False,
         num_warm_up_iterations: Optional[int] = None,
         num_warm_down_iterations: Optional[int] = None,
         warm_down_min_lr: float = 3e-5,
@@ -93,6 +95,7 @@ def __init__(  # pylint: disable=R0913
         self.min_lr = warm_down_min_lr
         self.use_softplus = use_softplus
         self.beta_softplus = beta_softplus
+        self.disable_lr_scheduler = disable_lr_scheduler
         self.agc_clipping_value = agc_clipping_value
         self.agc_eps = agc_eps
         self.centralize_gradients = centralize_gradients
@@ -245,8 +248,11 @@ def step(self, closure: CLOSURE = None) -> LOSS:
             noise_norm: float = math.sqrt((1.0 + beta2) ** 2 + beta2 ** 2)  # fmt: skip
 
             # warm up & down
-            lr: float = self.warm_up_dampening(group['lr'], group['step'])
-            lr = self.warm_down(lr, group['step'])
+            if self.disable_lr_scheduler:
+                lr: float = group['lr']
+            else:
+                lr: float = self.warm_up_dampening(group['lr'], group['step'])
+                lr = self.warm_down(lr, group['step'])
 
             for p in group['params']:
                 if p.grad is None:
@@ -279,7 +285,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 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)
+                    de_nom = softplus(de_nom, beta=self.beta_softplus)
 
                 grad = p.grad
                 centralize_gradient(grad, gc_conv_only=False)
@@ -289,7 +295,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 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)
+                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)
 
         self.lookahead_process_step()