Optimizer improvements, additions, cleanup

* Add MADGRAD code
* Fix Lamb (non-fused variant) to work w/ PyTorch XLA
* Tweak optimizer factory args (lr/learning_rate and opt/optimizer_name), may break compat
* Use newer fn signatures for all add,addcdiv, addcmul in optimizers
* Use upcoming PyTorch native Nadam if it's available
* Cleanup lookahead opt
* Add optimizer tests
* Remove novograd.py impl as it was messy, keep nvnovograd
* Make AdamP/SGDP work in channels_last layout
* Add rectified adablief mode (radabelief)
* Support a few more PyTorch optim, adamax, adagrad

Author: rwightman

benchmark.py

@@ -255,8 +255,8 @@ def __init__(self, model_name, device='cuda', torchscript=False, **kwargs):
 
         self.optimizer = create_optimizer_v2(
             self.model,
-            optimizer_name=kwargs.pop('opt', 'sgd'),
-            learning_rate=kwargs.pop('lr', 1e-4))
+            opt=kwargs.pop('opt', 'sgd'),
+            lr=kwargs.pop('lr', 1e-4))
 
     def _gen_target(self, batch_size):
         return torch.empty(

timm/optim/__init__.py

@@ -4,7 +4,6 @@
 from .adahessian import Adahessian
 from .lookahead import Lookahead
 from .nadam import Nadam
-from .novograd import NovoGrad
 from .nvnovograd import NvNovoGrad
 from .radam import RAdam
 from .rmsprop_tf import RMSpropTF

timm/optim/adabelief.py

@@ -18,7 +18,7 @@ class AdaBelief(Optimizer):
         amsgrad (boolean, optional): whether to use the AMSGrad variant of this
             algorithm from the paper `On the Convergence of Adam and Beyond`_
             (default: False)
-        weight_decouple (boolean, optional): ( default: True) If set as True, then
+        decoupled_decay (boolean, optional): ( default: True) If set as True, then
             the optimizer uses decoupled weight decay as in AdamW
         fixed_decay (boolean, optional): (default: False) This is used when weight_decouple
             is set as True.
@@ -39,9 +39,9 @@ class AdaBelief(Optimizer):
       - link to args.yaml: https://gist.github.com/juntang-zhuang/517ce3c27022b908bb93f78e4f786dc3
     """
 
-    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
-                 weight_decay=0, amsgrad=False, weight_decouple=True, fixed_decay=False, rectify=True,
-                 degenerated_to_sgd=True):
+    def __init__(
+            self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16, weight_decay=0, amsgrad=False,
+            decoupled_decay=True, fixed_decay=False, rectify=True, degenerated_to_sgd=True):
 
         if not 0.0 <= lr:
             raise ValueError("Invalid learning rate: {}".format(lr))
@@ -52,21 +52,17 @@ def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-16,
         if not 0.0 <= betas[1] < 1.0:
             raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
 
-        self.degenerated_to_sgd = degenerated_to_sgd
         if isinstance(params, (list, tuple)) and len(params) > 0 and isinstance(params[0], dict):
             for param in params:
                 if 'betas' in param and (param['betas'][0] != betas[0] or param['betas'][1] != betas[1]):
                     param['buffer'] = [[None, None, None] for _ in range(10)]
 
-        defaults = dict(lr=lr, betas=betas, eps=eps,
-                        weight_decay=weight_decay, amsgrad=amsgrad, buffer=[[None, None, None] for _ in range(10)])
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad,
+            degenerated_to_sgd=degenerated_to_sgd, decoupled_decay=decoupled_decay, rectify=rectify,
+            fixed_decay=fixed_decay, buffer=[[None, None, None] for _ in range(10)])
         super(AdaBelief, self).__init__(params, defaults)
 
-        self.degenerated_to_sgd = degenerated_to_sgd
-        self.weight_decouple = weight_decouple
-        self.rectify = rectify
-        self.fixed_decay = fixed_decay
-
     def __setstate__(self, state):
         super(AdaBelief, self).__setstate__(state)
         for group in self.param_groups:
@@ -133,8 +129,8 @@ def step(self, closure=None):
                         state['max_exp_avg_var'] = torch.zeros_like(p.data)
 
                 # perform weight decay, check if decoupled weight decay
-                if self.weight_decouple:
-                    if not self.fixed_decay:
+                if group['decoupled_decay']:
+                    if not group['fixed_decay']:
                         p.data.mul_(1.0 - group['lr'] * group['weight_decay'])
                     else:
                         p.data.mul_(1.0 - group['weight_decay'])
@@ -152,7 +148,7 @@ def step(self, closure=None):
                 # Update first and second moment running average
                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                 grad_residual = grad - exp_avg
-                exp_avg_var.mul_(beta2).addcmul_( grad_residual, grad_residual, value=1 - beta2)
+                exp_avg_var.mul_(beta2).addcmul_(grad_residual, grad_residual, value=1 - beta2)
 
                 if amsgrad:
                     max_exp_avg_var = state['max_exp_avg_var']
@@ -165,34 +161,36 @@ def step(self, closure=None):
                     denom = (exp_avg_var.add_(group['eps']).sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
 
                 # update
-                if not self.rectify:
+                if not group['rectify']:
                     # Default update
                     step_size = group['lr'] / bias_correction1
-                    p.data.addcdiv_( exp_avg, denom, value=-step_size)
-
-                else:  # Rectified update, forked from RAdam
+                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
+                else:
+                    # Rectified update, forked from RAdam
                     buffered = group['buffer'][int(state['step'] % 10)]
                     if state['step'] == buffered[0]:
-                        N_sma, step_size = buffered[1], buffered[2]
+                        num_sma, step_size = buffered[1], buffered[2]
                     else:
                         buffered[0] = state['step']
                         beta2_t = beta2 ** state['step']
-                        N_sma_max = 2 / (1 - beta2) - 1
-                        N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
-                        buffered[1] = N_sma
+                        num_sma_max = 2 / (1 - beta2) - 1
+                        num_sma = num_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
+                        buffered[1] = num_sma
 
                         # more conservative since it's an approximated value
-                        if N_sma >= 5:
+                        if num_sma >= 5:
                             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)) / (1 - beta1 ** state['step'])
-                        elif self.degenerated_to_sgd:
+                                (1 - beta2_t) *
+                                (num_sma - 4) / (num_sma_max - 4) *
+                                (num_sma - 2) / num_sma *
+                                num_sma_max / (num_sma_max - 2)) / (1 - beta1 ** state['step'])
+                        elif group['degenerated_to_sgd']:
                             step_size = 1.0 / (1 - beta1 ** state['step'])
                         else:
                             step_size = -1
                         buffered[2] = step_size
 
-                    if N_sma >= 5:
+                    if num_sma >= 5:
                         denom = exp_avg_var.sqrt().add_(group['eps'])
                         p.data.addcdiv_(exp_avg, denom, value=-step_size * group['lr'])
                     elif step_size > 0:

timm/optim/adafactor.py

@@ -34,15 +34,13 @@ class Adafactor(torch.optim.Optimizer):
         beta1 (float): coefficient used for computing running averages of gradient (default: None)
         weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
         scale_parameter (bool): if True, learning rate is scaled by root mean square of parameter (default: True)
-        relative_step (bool): if True, time-dependent learning rate is computed
-            instead of external learning rate (default: True)
         warmup_init (bool): time-dependent learning rate computation depends on
             whether warm-up initialization is being used (default: False)
     """
 
     def __init__(self, params, lr=None, eps=1e-30, eps_scale=1e-3, clip_threshold=1.0,
                  decay_rate=-0.8, betas=None, weight_decay=0.0, scale_parameter=True, warmup_init=False):
-        relative_step = lr is None
+        relative_step = not lr
         if warmup_init and not relative_step:
             raise ValueError('warmup_init requires relative_step=True')
 
@@ -138,37 +136,32 @@ def step(self, closure=None):
                     exp_avg_sq_row = state['exp_avg_sq_row']
                     exp_avg_sq_col = state['exp_avg_sq_col']
 
-                    exp_avg_sq_row.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-1))
-                    exp_avg_sq_col.mul_(beta2t).add_(1.0 - beta2t, update.mean(dim=-2))
-                    #exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)  # pytorch 1.6+
-                    #exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
 
                     # Approximation of exponential moving average of square of gradient
                     update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
                     update.mul_(grad)
                 else:
                     exp_avg_sq = state['exp_avg_sq']
 
-                    exp_avg_sq.mul_(beta2t).add_(1.0 - beta2t, update)
-                    #exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)  # pytorch 1.6+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)
                     update = exp_avg_sq.rsqrt().mul_(grad)
 
                 update.div_((self._rms(update) / group['clip_threshold']).clamp_(min=1.0))
                 update.mul_(lr_t)
 
                 if use_first_moment:
                     exp_avg = state['exp_avg']
-                    exp_avg.mul_(group["beta1"]).add_(1 - group["beta1"], update)
-                    #exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])  # pytorch 1.6+
+                    exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])
                     update = exp_avg
 
                 if group['weight_decay'] != 0:
-                    p_data_fp32.add_(-group["weight_decay"] * lr_t, p_data_fp32)
-                    #p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr_t)  # pytorch 1.6+
+                    p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr_t)
 
                 p_data_fp32.add_(-update)
 
                 if p.data.dtype in {torch.float16, torch.bfloat16}:
                     p.data.copy_(p_data_fp32)
 
-        return loss
+        return loss

timm/optim/adamp.py

@@ -9,48 +9,43 @@
 """
 
 import torch
-import torch.nn as nn
-from torch.optim.optimizer import Optimizer, required
+import torch.nn.functional as F
+from torch.optim.optimizer import Optimizer
 import math
 
-class AdamP(Optimizer):
-    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
-                 weight_decay=0, delta=0.1, wd_ratio=0.1, nesterov=False):
-        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
-                        delta=delta, wd_ratio=wd_ratio, nesterov=nesterov)
-        super(AdamP, self).__init__(params, defaults)
-
-    def _channel_view(self, x):
-        return x.view(x.size(0), -1)
 
-    def _layer_view(self, x):
-        return x.view(1, -1)
+def _channel_view(x) -> torch.Tensor:
+    return x.reshape(x.size(0), -1)
 
-    def _cosine_similarity(self, x, y, eps, view_func):
-        x = view_func(x)
-        y = view_func(y)
 
-        x_norm = x.norm(dim=1).add_(eps)
-        y_norm = y.norm(dim=1).add_(eps)
-        dot = (x * y).sum(dim=1)
+def _layer_view(x) -> torch.Tensor:
+    return x.reshape(1, -1)
 
-        return dot.abs() / x_norm / y_norm
 
-    def _projection(self, p, grad, perturb, delta, wd_ratio, eps):
-        wd = 1
-        expand_size = [-1] + [1] * (len(p.shape) - 1)
-        for view_func in [self._channel_view, self._layer_view]:
+def projection(p, grad, perturb, delta: float, wd_ratio: float, eps: float):
+    wd = 1.
+    expand_size = (-1,) + (1,) * (len(p.shape) - 1)
+    for view_func in [_channel_view, _layer_view]:
+        param_view = view_func(p.data)
+        grad_view = view_func(grad)
+        cosine_sim = F.cosine_similarity(grad_view, param_view, dim=1, eps=eps).abs_()
 
-            cosine_sim = self._cosine_similarity(grad, p.data, eps, view_func)
+        if cosine_sim.max() < delta / math.sqrt(param_view.size(1)):
+            p_n = p.data / param_view.norm(p=2, dim=1).add_(eps).reshape(expand_size)
+            perturb -= p_n * view_func(p_n * perturb).sum(dim=1).reshape(expand_size)
+            wd = wd_ratio
+            return perturb, wd
 
-            if cosine_sim.max() < delta / math.sqrt(view_func(p.data).size(1)):
-                p_n = p.data / view_func(p.data).norm(dim=1).view(expand_size).add_(eps)
-                perturb -= p_n * view_func(p_n * perturb).sum(dim=1).view(expand_size)
-                wd = wd_ratio
+    return perturb, wd
 
-                return perturb, wd
 
-        return perturb, wd
+class AdamP(Optimizer):
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0, delta=0.1, wd_ratio=0.1, nesterov=False):
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
+            delta=delta, wd_ratio=wd_ratio, nesterov=nesterov)
+        super(AdamP, self).__init__(params, defaults)
 
     def step(self, closure=None):
         loss = None
@@ -81,8 +76,8 @@ def step(self, closure=None):
                 bias_correction1 = 1 - beta1 ** state['step']
                 bias_correction2 = 1 - beta2 ** state['step']
 
-                exp_avg.mul_(beta1).add_(1 - beta1, grad)
-                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
 
                 denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
                 step_size = group['lr'] / bias_correction1
@@ -93,15 +88,15 @@ def step(self, closure=None):
                     perturb = exp_avg / denom
 
                 # Projection
-                wd_ratio = 1
+                wd_ratio = 1.
                 if len(p.shape) > 1:
-                    perturb, wd_ratio = self._projection(p, grad, perturb, group['delta'], group['wd_ratio'], group['eps'])
+                    perturb, wd_ratio = projection(p, grad, perturb, group['delta'], group['wd_ratio'], group['eps'])
 
                 # Weight decay
                 if group['weight_decay'] > 0:
-                    p.data.mul_(1 - group['lr'] * group['weight_decay'] * wd_ratio)
+                    p.data.mul_(1. - group['lr'] * group['weight_decay'] * wd_ratio)
 
                 # Step
-                p.data.add_(-step_size, perturb)
+                p.data.add_(perturb, alpha=-step_size)
 
         return loss

timm/optim/adamw.py

@@ -1,5 +1,8 @@
 """ AdamW Optimizer
 Impl copied from PyTorch master
+
+NOTE: Builtin optim.AdamW is used by the factory, this impl only serves as a Python based reference, will be removed
+someday
 """
 import math
 import torch
@@ -100,8 +103,8 @@ def step(self, closure=None):
                 bias_correction2 = 1 - beta2 ** state['step']
 
                 # Decay the first and second moment running average coefficient
-                exp_avg.mul_(beta1).add_(1 - beta1, grad)
-                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                 if amsgrad:
                     # Maintains the maximum of all 2nd moment running avg. till now
                     torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
@@ -112,6 +115,6 @@ def step(self, closure=None):
 
                 step_size = group['lr'] / bias_correction1
 
-                p.data.addcdiv_(-step_size, exp_avg, denom)
+                p.data.addcdiv_(exp_avg, denom, value=-step_size)
 
         return loss