More optimizer cleanup. Change all to no longer use .data. Improve (b)float16 use with adabelief. Add XLA compatible Lars.

Author: rwightman

tests/test_optim.py

@@ -490,6 +490,33 @@ def test_lamb(optimizer):
     _test_model(optimizer, dict(lr=1e-3))
 
 
+@pytest.mark.parametrize('optimizer',  ['lars', 'larc', 'nlars', 'nlarc'])
+def test_lars(optimizer):
+    _test_basic_cases(
+        lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
+    )
+    _test_basic_cases(
+        lambda weight, bias: create_optimizer_v2(
+            _build_params_dict(weight, bias, lr=1e-3),
+            optimizer,
+            lr=1e-1)
+    )
+    _test_basic_cases(
+        lambda weight, bias: create_optimizer_v2(
+            _build_params_dict_single(weight, bias, lr=1e-3),
+            optimizer,
+            lr=1e-3)
+    )
+    _test_basic_cases(
+        lambda weight, bias: create_optimizer_v2(
+            _build_params_dict_single(weight, bias, lr=1e-3), optimizer)
+    )
+    _test_rosenbrock(
+        lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
+    )
+    _test_model(optimizer, dict(lr=1e-3))
+
+
 @pytest.mark.parametrize('optimizer',  ['madgrad', 'madgradw'])
 def test_madgrad(optimizer):
     _test_basic_cases(

timm/optim/adabelief.py

@@ -68,6 +68,7 @@ def __setstate__(self, state):
         for group in self.param_groups:
             group.setdefault('amsgrad', False)
 
+    @torch.no_grad()
     def reset(self):
         for group in self.param_groups:
             for p in group['params']:
@@ -77,14 +78,15 @@ def reset(self):
                 # State initialization
                 state['step'] = 0
                 # Exponential moving average of gradient values
-                state['exp_avg'] = torch.zeros_like(p.data)
+                state['exp_avg'] = torch.zeros_like(p)
 
                 # Exponential moving average of squared gradient values
-                state['exp_avg_var'] = torch.zeros_like(p.data)
+                state['exp_avg_var'] = torch.zeros_like(p)
                 if amsgrad:
                     # Maintains max of all exp. moving avg. of sq. grad. values
-                    state['max_exp_avg_var'] = torch.zeros_like(p.data)
+                    state['max_exp_avg_var'] = torch.zeros_like(p)
 
+    @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
         Arguments:
@@ -93,50 +95,47 @@ def step(self, closure=None):
         """
         loss = None
         if closure is not None:
-            loss = closure()
+            with torch.enable_grad():
+                loss = closure()
 
         for group in self.param_groups:
             for p in group['params']:
                 if p.grad is None:
                     continue
-                
-                # cast data type
-                half_precision = False
-                if p.data.dtype == torch.float16:
-                    half_precision = True
-                    p.data = p.data.float()
-                    p.grad = p.grad.float()
-
-                grad = p.grad.data
+                grad = p.grad
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
                 if grad.is_sparse:
                     raise RuntimeError(
                         'AdaBelief does not support sparse gradients, please consider SparseAdam instead')
-                amsgrad = group['amsgrad']
 
-                state = self.state[p]
+                p_fp32 = p
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p_fp32 = p_fp32.float()
 
+                amsgrad = group['amsgrad']
                 beta1, beta2 = group['betas']
-
+                state = self.state[p]
                 # State initialization
                 if len(state) == 0:
                     state['step'] = 0
                     # Exponential moving average of gradient values
-                    state['exp_avg'] = torch.zeros_like(p.data)
+                    state['exp_avg'] = torch.zeros_like(p_fp32)
                     # Exponential moving average of squared gradient values
-                    state['exp_avg_var'] = torch.zeros_like(p.data)
+                    state['exp_avg_var'] = torch.zeros_like(p_fp32)
                     if amsgrad:
                         # Maintains max of all exp. moving avg. of sq. grad. values
-                        state['max_exp_avg_var'] = torch.zeros_like(p.data)
+                        state['max_exp_avg_var'] = torch.zeros_like(p_fp32)
 
                 # perform weight decay, check if decoupled weight decay
                 if group['decoupled_decay']:
                     if not group['fixed_decay']:
-                        p.data.mul_(1.0 - group['lr'] * group['weight_decay'])
+                        p_fp32.mul_(1.0 - group['lr'] * group['weight_decay'])
                     else:
-                        p.data.mul_(1.0 - group['weight_decay'])
+                        p_fp32.mul_(1.0 - group['weight_decay'])
                 else:
                     if group['weight_decay'] != 0:
-                        grad.add_(p.data, alpha=group['weight_decay'])
+                        grad.add_(p_fp32, alpha=group['weight_decay'])
 
                 # get current state variable
                 exp_avg, exp_avg_var = state['exp_avg'], state['exp_avg_var']
@@ -164,7 +163,7 @@ def step(self, closure=None):
                 if not group['rectify']:
                     # Default update
                     step_size = group['lr'] / bias_correction1
-                    p.data.addcdiv_(exp_avg, denom, value=-step_size)
+                    p_fp32.addcdiv_(exp_avg, denom, value=-step_size)
                 else:
                     # Rectified update, forked from RAdam
                     buffered = group['buffer'][int(state['step'] % 10)]
@@ -192,12 +191,11 @@ def step(self, closure=None):
 
                     if num_sma >= 5:
                         denom = exp_avg_var.sqrt().add_(group['eps'])
-                        p.data.addcdiv_(exp_avg, denom, value=-step_size * group['lr'])
+                        p_fp32.addcdiv_(exp_avg, denom, value=-step_size * group['lr'])
                     elif step_size > 0:
-                        p.data.add_(exp_avg, alpha=-step_size * group['lr'])
+                        p_fp32.add_(exp_avg, alpha=-step_size * group['lr'])
 
-                if half_precision:
-                    p.data = p.data.half()
-                    p.grad = p.grad.half() 
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p.copy_(p_fp32)
 
         return loss

timm/optim/adafactor.py

@@ -76,29 +76,30 @@ def _approx_sq_grad(self, exp_avg_sq_row, exp_avg_sq_col):
         c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
         return torch.mul(r_factor, c_factor)
 
+    @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
         Arguments:
             closure (callable, optional): A closure that reevaluates the model and returns the loss.
         """
         loss = None
         if closure is not None:
-            loss = closure()
+            with torch.enable_grad():
+                loss = closure()
 
         for group in self.param_groups:
             for p in group['params']:
                 if p.grad is None:
                     continue
-                grad = p.grad.data
+                grad = p.grad
                 if grad.dtype in {torch.float16, torch.bfloat16}:
                     grad = grad.float()
                 if grad.is_sparse:
                     raise RuntimeError('Adafactor does not support sparse gradients.')
 
                 state = self.state[p]
-                grad_shape = grad.shape
 
-                factored, use_first_moment = self._get_options(group, grad_shape)
+                factored, use_first_moment = self._get_options(group, grad.shape)
                 # State Initialization
                 if len(state) == 0:
                     state['step'] = 0
@@ -107,8 +108,8 @@ def step(self, closure=None):
                         # Exponential moving average of gradient values
                         state['exp_avg'] = torch.zeros_like(grad)
                     if factored:
-                        state['exp_avg_sq_row'] = torch.zeros(grad_shape[:-1]).to(grad)
-                        state['exp_avg_sq_col'] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                        state['exp_avg_sq_row'] = torch.zeros(grad.shape[:-1]).to(grad)
+                        state['exp_avg_sq_col'] = torch.zeros(grad.shape[:-2] + grad.shape[-1:]).to(grad)
                     else:
                         state['exp_avg_sq'] = torch.zeros_like(grad)
 
@@ -122,12 +123,12 @@ def step(self, closure=None):
                     else:
                         state['exp_avg_sq'] = state['exp_avg_sq'].to(grad)
 
-                p_data_fp32 = p.data
-                if p.data.dtype in {torch.float16, torch.bfloat16}:
-                    p_data_fp32 = p_data_fp32.float()
+                p_fp32 = p
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p_fp32 = p_fp32.float()
 
                 state['step'] += 1
-                state['RMS'] = self._rms(p_data_fp32)
+                state['RMS'] = self._rms(p_fp32)
                 lr_t = self._get_lr(group, state)
 
                 beta2t = 1.0 - math.pow(state['step'], group['decay_rate'])
@@ -157,11 +158,10 @@ def step(self, closure=None):
                     update = exp_avg
 
                 if group['weight_decay'] != 0:
-                    p_data_fp32.add_(p_data_fp32, alpha=-group['weight_decay'] * lr_t)
+                    p_fp32.add_(p_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)
+                p_fp32.add_(-update)
+                if p.dtype in {torch.float16, torch.bfloat16}:
+                    p.copy_(p_fp32)
 
         return loss

timm/optim/adamp.py

@@ -26,12 +26,13 @@ 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)
+        param_view = view_func(p)
         grad_view = view_func(grad)
         cosine_sim = F.cosine_similarity(grad_view, param_view, dim=1, eps=eps).abs_()
 
+        # FIXME this is a problem for PyTorch XLA
         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)
+            p_n = p / 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
@@ -47,17 +48,19 @@ def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
             delta=delta, wd_ratio=wd_ratio, nesterov=nesterov)
         super(AdamP, self).__init__(params, defaults)
 
+    @torch.no_grad()
     def step(self, closure=None):
         loss = None
         if closure is not None:
-            loss = closure()
+            with torch.enable_grad():
+                loss = closure()
 
         for group in self.param_groups:
             for p in group['params']:
                 if p.grad is None:
                     continue
 
-                grad = p.grad.data
+                grad = p.grad
                 beta1, beta2 = group['betas']
                 nesterov = group['nesterov']
 
@@ -66,8 +69,8 @@ def step(self, closure=None):
                 # State initialization
                 if len(state) == 0:
                     state['step'] = 0
-                    state['exp_avg'] = torch.zeros_like(p.data)
-                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
 
                 # Adam
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
@@ -94,9 +97,9 @@ def step(self, closure=None):
 
                 # Weight decay
                 if group['weight_decay'] > 0:
-                    p.data.mul_(1. - group['lr'] * group['weight_decay'] * wd_ratio)
+                    p.mul_(1. - group['lr'] * group['weight_decay'] * wd_ratio)
 
                 # Step
-                p.data.add_(perturb, alpha=-step_size)
+                p.add_(perturb, alpha=-step_size)
 
         return loss

timm/optim/adamw.py

@@ -55,6 +55,7 @@ def __setstate__(self, state):
         for group in self.param_groups:
             group.setdefault('amsgrad', False)
 
+    @torch.no_grad()
     def step(self, closure=None):
         """Performs a single optimization step.
 
@@ -64,7 +65,8 @@ def step(self, closure=None):
         """
         loss = None
         if closure is not None:
-            loss = closure()
+            with torch.enable_grad():
+                loss = closure()
 
         for group in self.param_groups:
             for p in group['params']:
@@ -75,7 +77,7 @@ def step(self, closure=None):
                 p.data.mul_(1 - group['lr'] * group['weight_decay'])
 
                 # Perform optimization step
-                grad = p.grad.data
+                grad = p.grad
                 if grad.is_sparse:
                     raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                 amsgrad = group['amsgrad']
@@ -86,12 +88,12 @@ def step(self, closure=None):
                 if len(state) == 0:
                     state['step'] = 0
                     # Exponential moving average of gradient values
-                    state['exp_avg'] = torch.zeros_like(p.data)
+                    state['exp_avg'] = torch.zeros_like(p)
                     # Exponential moving average of squared gradient values
-                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
                     if amsgrad:
                         # Maintains max of all exp. moving avg. of sq. grad. values
-                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+                        state['max_exp_avg_sq'] = torch.zeros_like(p)
 
                 exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                 if amsgrad:
@@ -115,6 +117,6 @@ def step(self, closure=None):
 
                 step_size = group['lr'] / bias_correction1
 
-                p.data.addcdiv_(exp_avg, denom, value=-step_size)
+                p.addcdiv_(exp_avg, denom, value=-step_size)
 
         return loss