Merge pull request #96 from kozistr/refactor/optimizers

[Refactor] Cleanup the codes

Author: kozistr

pytorch_optimizer/optimizer/adabelief.py

@@ -95,6 +95,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         for group in self.param_groups:
             beta1, beta2 = group['betas']
+            weight_decay: float = group['weight_decay']
+
             if self.rectify:
                 n_sma_max: float = 2.0 / (1.0 - beta2) - 1.0
 
@@ -106,13 +108,6 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if grad.is_sparse:
                     raise NoSparseGradientError(self.__name__)
 
-                if grad.dtype in (torch.float16, torch.bfloat16):
-                    grad = grad.float()
-
-                p_fp32 = p
-                if p.dtype in (torch.float16, torch.bfloat16):
-                    p_fp32 = p_fp32.float()
-
                 state = self.state[p]
                 if len(state) == 0:
                     state['step'] = 0
@@ -122,70 +117,65 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                         state['max_exp_avg_var'] = torch.zeros_like(p)
 
                 if self.weight_decouple:
-                    decay: float = (
-                        group['lr'] * group['weight_decay'] if not self.fixed_decay else group['weight_decay']
-                    )
-                    p_fp32.mul_(1.0 - decay)
-                elif group['weight_decay'] != 0:
-                    grad.add_(p_fp32, alpha=group['weight_decay'])
-
-                exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
+                    p.mul_(1.0 - (group['lr'] * weight_decay if not self.fixed_decay else weight_decay))
+                elif weight_decay > 0.0:
+                    grad.add_(p, alpha=weight_decay)
 
                 state['step'] += 1
+                exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
 
                 bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2 = 1.0 - beta2 ** state['step']
+                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
 
                 exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                 grad_residual = grad - exp_avg
-                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2)
-                exp_avg_var.add_(group['eps'])
-                if group['amsgrad']:
-                    torch.max(state['max_exp_avg_var'], exp_avg_var, out=exp_avg_var)
+                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1.0 - beta2).add_(group['eps'])
 
-                de_nom = (exp_avg_var.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
+                if group['amsgrad']:
+                    max_exp_avg_var = state['max_exp_avg_var']
+                    torch.max(max_exp_avg_var, exp_avg_var, out=max_exp_avg_var)
+                    de_nom = max_exp_avg_var.sqrt()
+                else:
+                    de_nom = exp_avg_var.sqrt()
+                de_nom.div_(bias_correction2_sq).add_(group['eps'])
 
                 if not self.rectify:
-                    step_size = group['lr']
-                    if not group['adamd_debias_term']:
-                        step_size /= bias_correction1
-                    p_fp32.addcdiv_(exp_avg, de_nom, value=-step_size)
+                    step_size: float = group['lr'] if group['adamd_debias_term'] else group['lr'] / bias_correction1
+                    p.addcdiv_(exp_avg, de_nom, value=-step_size)
+                    continue
+
+                buffered = group['buffer'][state['step'] % 10]
+                if state['step'] == buffered[0]:
+                    n_sma, step_size = buffered[1], buffered[2]
                 else:
-                    buffered = group['buffer'][state['step'] % 10]
-                    if state['step'] == buffered[0]:
-                        n_sma, step_size = buffered[1], buffered[2]
-                    else:
-                        buffered[0] = state['step']
-                        beta2_t = beta2 ** state['step']
-                        n_sma = n_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
-                        buffered[1] = n_sma
-
-                        if n_sma >= self.n_sma_threshold:
-                            step_size = math.sqrt(
-                                (1 - beta2_t)
-                                * (n_sma - 4)
-                                / (n_sma_max - 4)
-                                * (n_sma - 2)
-                                / n_sma
-                                * n_sma_max
-                                / (n_sma_max - 2)
-                            )
-                            if not group['adamd_debias_term']:
-                                step_size /= bias_correction1
-                        elif self.degenerated_to_sgd:
-                            step_size = 1.0 / bias_correction1
-                        else:
-                            step_size = -1
-
-                        buffered[2] = step_size
+                    buffered[0] = state['step']
+                    beta2_t = beta2 ** state['step']
+                    n_sma = n_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
+                    buffered[1] = n_sma
 
                     if n_sma >= self.n_sma_threshold:
-                        de_nom = exp_avg_var.sqrt().add_(group['eps'])
-                        p_fp32.addcdiv_(exp_avg, de_nom, value=-step_size * group['lr'])
-                    elif step_size > 0:
-                        p_fp32.add_(exp_avg, alpha=-step_size * group['lr'])
+                        step_size = math.sqrt(
+                            (1 - beta2_t)
+                            * (n_sma - 4)
+                            / (n_sma_max - 4)
+                            * (n_sma - 2)
+                            / n_sma
+                            * n_sma_max
+                            / (n_sma_max - 2)
+                        )
+                        if not group['adamd_debias_term']:
+                            step_size /= bias_correction1
+                    elif self.degenerated_to_sgd:
+                        step_size = 1.0 / bias_correction1
+                    else:
+                        step_size = -1
+
+                    buffered[2] = step_size
 
-                if p.dtype in (torch.float16, torch.bfloat16):
-                    p.copy_(p_fp32)
+                if n_sma >= self.n_sma_threshold:
+                    de_nom = exp_avg_var.sqrt().add_(group['eps'])
+                    p.addcdiv_(exp_avg, de_nom, value=-step_size * group['lr'])
+                elif step_size > 0:
+                    p.add_(exp_avg, alpha=-step_size * group['lr'])
 
         return loss

pytorch_optimizer/optimizer/adabound.py

@@ -93,6 +93,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
         for group, base_lr in zip(self.param_groups, self.base_lrs):
             beta1, beta2 = group['betas']
+            weight_decay: float = group['weight_decay']
             for p in group['params']:
                 if p.grad is None:
                     continue
@@ -113,13 +114,11 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 state['step'] += 1
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
 
-                if group['weight_decay'] > 0.0:
+                if weight_decay > 0.0:
                     if self.weight_decouple:
-                        p.mul_(
-                            1.0 - (group['weight_decay'] if self.fixed_decay else group['lr'] * group['weight_decay'])
-                        )
+                        p.mul_(1.0 - (weight_decay if self.fixed_decay else group['lr'] * weight_decay))
                     else:
-                        grad.add_(p, alpha=group['weight_decay'])
+                        grad.add_(p, alpha=weight_decay)
 
                 exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)

pytorch_optimizer/optimizer/adamp.py

@@ -108,15 +108,15 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
 
                 bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2 = 1.0 - beta2 ** state['step']
+                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
 
                 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.0 - beta2)
 
-                inv_de_nom = 1.0 / (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
+                inv_de_nom = 1.0 / (exp_avg_sq.sqrt() / bias_correction2_sq).add_(group['eps'])
 
                 perturb = exp_avg.clone()
                 if group['nesterov']:
@@ -125,7 +125,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 else:
                     perturb.mul_(inv_de_nom)
 
-                wd_ratio: float = 1
+                wd_ratio: float = 1.0
                 if len(p.shape) > 1:
                     perturb, wd_ratio = projection(
                         p,
@@ -139,10 +139,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if group['weight_decay'] > 0.0:
                     p.mul_(1.0 - group['lr'] * group['weight_decay'] * wd_ratio)
 
-                step_size = group['lr']
-                if not group['adamd_debias_term']:
-                    step_size /= bias_correction1
-
+                step_size: float = group['lr'] if group['adamd_debias_term'] else group['lr'] / bias_correction1
                 p.add_(perturb, alpha=-step_size)
 
         return loss

pytorch_optimizer/optimizer/adan.py

@@ -138,7 +138,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if self.use_gc:
                     grad = centralize_gradient(grad, gc_conv_only=False)
 
-                grad_diff = grad - state['previous_grad']
+                grad_diff = -state['previous_grad']
+                grad_diff.add_(grad)
                 state['previous_grad'].copy_(grad)
 
                 update = grad + beta2 * grad_diff

pytorch_optimizer/optimizer/adapnm.py

@@ -111,25 +111,22 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 state['step'] += 1
 
                 bias_correction1 = 1 - beta1 ** state['step']
-                bias_correction2 = 1 - beta2 ** state['step']
+                bias_correction2_sq = math.sqrt(1 - beta2 ** state['step'])
 
                 exp_avg_sq = state['exp_avg_sq']
                 if state['step'] % 2 == 1:
                     exp_avg, neg_exp_avg = state['exp_avg'], state['neg_exp_avg']
                 else:
                     exp_avg, neg_exp_avg = state['neg_exp_avg'], state['exp_avg']
 
-                exp_avg.mul_(beta1 ** 2).add_(grad, alpha=1 - beta1 ** 2)  # fmt: skip
+                exp_avg.mul_(beta1 ** 2).add_(grad, alpha=1.0 - beta1 ** 2)  # fmt: skip
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                 if group['amsgrad']:
                     torch.max(state['max_exp_avg_sq'], exp_avg_sq, out=exp_avg_sq)
 
-                de_nom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
-
-                step_size = group['lr']
-                if not group['adamd_debias_term']:
-                    step_size /= bias_correction1
+                de_nom = (exp_avg_sq.sqrt() / bias_correction2_sq).add_(group['eps'])
 
+                step_size: float = group['lr'] if group['adamd_debias_term'] else group['lr'] / bias_correction1
                 pn_momentum = exp_avg.mul(1.0 + beta3).add(neg_exp_avg, alpha=-beta3).mul(1.0 / noise_norm)
                 p.addcdiv_(pn_momentum, de_nom, value=-step_size)
 

pytorch_optimizer/optimizer/diffgrad.py

@@ -85,32 +85,31 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     state['exp_avg_sq'] = torch.zeros_like(p)
                     state['previous_grad'] = torch.zeros_like(p)
 
+                state['step'] += 1
                 exp_avg, exp_avg_sq, previous_grad = state['exp_avg'], state['exp_avg_sq'], state['previous_grad']
 
-                if group['weight_decay'] != 0:
+                if group['weight_decay'] > 0.0:
                     grad.add_(p, alpha=group['weight_decay'])
 
-                state['step'] += 1
-
                 # Decay the first and second moment running average coefficient
                 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().add_(group['eps'])
 
                 bias_correction1 = 1.0 - beta1 ** state['step']
-                bias_correction2 = 1.0 - beta2 ** state['step']
+                bias_correction2_sq = math.sqrt(1.0 - beta2 ** state['step'])
 
                 # compute diffGrad coefficient (dfc)
-                diff = abs(previous_grad - grad)
-                dfc = 1.0 / (1.0 + torch.exp(-diff))
+                dfc = previous_grad.clone()
+                dfc.sub_(grad).abs_().sigmoid_().mul_(exp_avg)
                 state['previous_grad'].copy_(grad)
 
-                step_size = group['lr'] * math.sqrt(bias_correction2)
+                step_size = group['lr'] * bias_correction2_sq
                 if not group['adamd_debias_term']:
                     step_size /= bias_correction1
 
-                # update momentum with dfc (exp_avg * dfc)
-                p.addcdiv_(exp_avg * dfc, de_nom, value=-step_size)
+                # update momentum with dfc
+                p.addcdiv_(dfc, de_nom, value=-step_size)
 
         return loss

pytorch_optimizer/optimizer/diffrgrad.py

@@ -90,23 +90,12 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 if grad.is_sparse:
                     raise NoSparseGradientError(self.__name__)
 
-                if grad.dtype in (torch.float16, torch.bfloat16):
-                    grad = grad.float()
-
-                p_fp32 = p
-                if p.dtype in (torch.float16, torch.bfloat16):
-                    p_fp32 = p_fp32.float()
-
                 state = self.state[p]
                 if len(state) == 0:
                     state['step'] = 0
-                    state['exp_avg'] = torch.zeros_like(p_fp32)
-                    state['exp_avg_sq'] = torch.zeros_like(p_fp32)
-                    state['previous_grad'] = torch.zeros_like(p_fp32)
-                else:
-                    state['exp_avg'] = state['exp_avg'].type_as(p_fp32)
-                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_fp32)
-                    state['previous_grad'] = state['previous_grad'].type_as(p_fp32)
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['previous_grad'] = torch.zeros_like(p)
 
                 state['step'] += 1
                 exp_avg, exp_avg_sq, previous_grad = state['exp_avg'], state['exp_avg_sq'], state['previous_grad']
@@ -117,8 +106,8 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
 
                 # compute diffGrad coefficient (dfc)
-                diff = abs(previous_grad - grad)
-                dfc = 1.0 / (1.0 + torch.exp(-diff))
+                dfc = previous_grad.clone()
+                dfc.sub_(grad).abs_().sigmoid_().mul_(exp_avg)
                 state['previous_grad'].copy_(grad)
 
                 buffered = group['buffer'][state['step'] % 10]
@@ -149,21 +138,13 @@ def step(self, closure: CLOSURE = None) -> LOSS:
 
                     buffered[2] = step_size
 
-                if n_sma >= self.n_sma_threshold:
-                    if group['weight_decay'] != 0:
-                        p_fp32.add_(p_fp32, alpha=-group['weight_decay'] * group['lr'])
+                if group['weight_decay'] > 0.0:
+                    p.add_(p, alpha=-group['weight_decay'] * group['lr'])
 
+                if n_sma >= self.n_sma_threshold:
                     de_nom = exp_avg_sq.sqrt().add_(group['eps'])
-
-                    # update momentum with dfc
-                    p_fp32.addcdiv_(exp_avg * dfc.float(), de_nom, value=-step_size * group['lr'])
+                    p.addcdiv_(dfc, de_nom, value=-step_size * group['lr'])
                 elif step_size > 0:
-                    if group['weight_decay'] != 0:
-                        p_fp32.add_(p_fp32, alpha=-group['weight_decay'] * group['lr'])
-
-                    p_fp32.add_(exp_avg, alpha=-step_size * group['lr'])
-
-                if p.dtype in (torch.float16, torch.bfloat16):
-                    p.copy_(p_fp32)
+                    p.add_(exp_avg, alpha=-step_size * group['lr'])
 
         return loss

pytorch_optimizer/optimizer/gsam.py

@@ -140,7 +140,7 @@ def gradient_decompose(self, alpha: float = 0.0):
                 if p.grad is None:
                     continue
 
-                vertical = self.state[p]['old_g'] - cosine * old_grad_norm * p.grad.data / (
+                vertical = self.state[p]['old_g'] - cosine * old_grad_norm * p.grad / (
                     new_grad_norm + self.perturb_eps
                 )
                 p.grad.add_(vertical, alpha=-alpha)

pytorch_optimizer/optimizer/lars.py

@@ -104,13 +104,13 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                     if 'momentum_buffer' not in param_state:
                         param_state['momentum_buffer'] = grad.clone().detach()
 
-                    mu = param_state['momentum_buffer']
-                    mu.mul_(group['momentum']).add_(grad, alpha=1.0 - group['dampening'])
+                    mb = param_state['momentum_buffer']
+                    mb.mul_(group['momentum']).add_(grad, alpha=1.0 - group['dampening'])
 
                     if group['nesterov']:
-                        grad.add_(mu, alpha=group['momentum'])
+                        grad.add_(mb, alpha=group['momentum'])
                     else:
-                        grad.copy_(mu)
+                        grad.copy_(mb)
 
                 p.add_(grad, alpha=-group['lr'])