Merge pull request #30 from kozistr/feature/adamp-gc

[Feature] Support Gradient Centralization on `AdamP` optimizer

Author: kozistr

pytorch_optimizer/adamp.py

@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.gc import centralize_gradient
 from pytorch_optimizer.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
@@ -32,6 +33,7 @@ def __init__(
         weight_decay: float = 0.0,
         delta: float = 0.1,
         wd_ratio: float = 0.1,
+        use_gc: bool = False,
         nesterov: bool = False,
         eps: float = 1e-8,
     ):
@@ -43,13 +45,15 @@ def __init__(
         :param delta: float. threshold that determines whether a set of parameters is scale invariant or not
         :param wd_ratio: float. relative weight decay applied on scale-invariant parameters compared to that applied
             on scale-variant parameters
+        :param use_gc: bool. use gradient centralization
         :param nesterov: bool. enables Nesterov momentum
         :param eps: float. term added to the denominator to improve numerical stability
         """
         self.lr = lr
         self.betas = betas
         self.weight_decay = weight_decay
         self.wd_ratio = wd_ratio
+        self.use_gc = use_gc
         self.eps = eps
 
         self.check_valid_parameters()
@@ -146,6 +150,9 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                 grad = p.grad.data
 
+                if self.use_gc:
+                    grad = centralize_gradient(grad, gc_conv_only=False)
+
                 exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
 

pytorch_optimizer/ranger.py

@@ -4,6 +4,7 @@
 import torch
 from torch.optim.optimizer import Optimizer
 
+from pytorch_optimizer.gc import centralize_gradient
 from pytorch_optimizer.types import BETAS, BUFFER, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 
 
@@ -117,8 +118,8 @@ def step(self, _: CLOSURE = None) -> LOSS:
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                 beta1, beta2 = group['betas']
 
-                if grad.dim() > self.gc_gradient_threshold:
-                    grad.add_(-grad.mean(dim=tuple(range(1, grad.dim())), keepdim=True))
+                if self.use_gc and grad.dim() > self.gc_gradient_threshold:
+                    grad = centralize_gradient(grad, gc_conv_only=False)
 
                 state['step'] += 1