feature: support Gradient Centralization

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)