Remove some changes to existing behaviour and small correctness issues with optimizer. tensor step/LR support.

Author: rwightman

timm/optim/_helpers.py

@@ -1,6 +1,6 @@
 """Small optimizer helpers shared by timm optimizer implementations."""
 
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Union
 
 import torch
 from torch import Tensor
@@ -10,12 +10,22 @@ def _get_scalar_dtype() -> torch.dtype:
     return torch.float64 if torch.get_default_dtype() == torch.float64 else torch.float32
 
 
-def _init_scalar(value: float = 0.0, device=None) -> Tensor:
-    return torch.tensor(value, dtype=_get_scalar_dtype(), device=device)
+def _init_scalar(
+        value: Union[float, Tensor] = 0.0,
+        device=None,
+        dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    if isinstance(value, Tensor):
+        dtype = dtype or value.dtype
+        device = value.device if device is None else torch.device(device)
+        if value.device == device and value.dtype == dtype:
+            return value
+        return value.to(device=device, dtype=dtype)
+    return torch.tensor(value, dtype=dtype or _get_scalar_dtype(), device=device)
 
 
-def _zeros_scalar(device=None) -> Tensor:
-    return torch.zeros((), dtype=_get_scalar_dtype(), device=device)
+def _zeros_scalar(device=None, dtype: Optional[torch.dtype] = None) -> Tensor:
+    return torch.zeros((), dtype=dtype or _get_scalar_dtype(), device=device)
 
 
 def _is_compiling() -> bool:

timm/optim/adafactor_bv.py

@@ -16,7 +16,7 @@
 from torch import Tensor
 from torch.optim import Optimizer
 
-from ._helpers import _get_scalar_dtype, _get_value, _init_scalar, _validate_scalar
+from ._helpers import _get_value, _init_scalar, _validate_scalar
 from ._types import ParamsT
 
 
@@ -111,7 +111,7 @@ def __setstate__(self, state):
             for p in group['params']:
                 p_state = self.state.get(p, {})
                 if len(p_state) != 0 and 'step' in p_state:
-                    p_state['step'] = _init_scalar(float(p_state['step']), device='cpu')
+                    p_state['step'] = _init_scalar(p_state['step'], device='cpu', dtype=torch.float64)
 
                 if 'exp_avg' in p_state and torch.is_tensor(p_state['exp_avg']):
                     # FIXME this is a bit of a hack, optimizer.load_state_dict appears to upcast
@@ -148,7 +148,7 @@ def step(self, closure=None):
                 state = self.state[p]
 
                 if len(state) == 0:
-                    state['step'] = _init_scalar(device='cpu')
+                    state['step'] = _init_scalar(device='cpu', dtype=torch.float64)
 
                     shape = p.grad.shape
                     factored_dims = _factored_dims(

timm/optim/adamw.py

@@ -95,7 +95,7 @@ def __setstate__(self, state):
                 p_state = self.state.get(p, {})
                 if p_state and 'step' in p_state:
                     p_state['step'] = _init_scalar(
-                        float(p_state['step']),
+                        p_state['step'],
                         device=p.device if group['capturable'] else 'cpu',
                     )
 

timm/optim/lamb.py

@@ -129,7 +129,7 @@ def __setstate__(self, state):
             group.setdefault('decoupled_decay', False)
             group.setdefault('corrected_weight_decay', False)
             if 'step' in group:
-                group['step'] = _init_scalar(float(group['step']), device='cpu')
+                group['step'] = _init_scalar(group['step'], device='cpu')
 
     def _get_clip_grad_norm(self):
         max_grad_norm = self.defaults['max_grad_norm']
@@ -169,14 +169,15 @@ def step(self, closure=None):
             grad_averaging = 1 if group['grad_averaging'] else 0
             beta3 = 1 - beta1 if grad_averaging else 1.0
 
-            if not any(p.grad is not None for p in group['params']):
-                continue
-
             # assume same step across group now to simplify things
             # per parameter step can be easily support by making it tensor, or pass list into kernel
             if 'step' not in group:
                 group['step'] = _init_scalar(device='cpu')
             group['step'].add_(1)
+
+            if not any(p.grad is not None for p in group['params']):
+                continue
+
             step = _get_value(group['step'])
 
             if bias_correction:

timm/optim/laprop.py

@@ -66,19 +66,15 @@ def __setstate__(self, state):
                 if not p_state:
                     continue
                 if 'step' in p_state:
-                    p_state['step'] = _init_scalar(float(p_state['step']), device='cpu')
-                if (
-                        'exp_avg_lr_1' in p_state
-                        and torch.is_tensor(group['lr'])
-                        and not torch.is_tensor(p_state['exp_avg_lr_1'])
-                ):
-                    p_state['exp_avg_lr_1'] = torch.tensor(
-                        float(p_state['exp_avg_lr_1']),
+                    p_state['step'] = _init_scalar(p_state['step'], device='cpu')
+                if 'exp_avg_lr_1' in p_state and torch.is_tensor(group['lr']):
+                    p_state['exp_avg_lr_1'] = _init_scalar(
+                        p_state['exp_avg_lr_1'],
                         dtype=group['lr'].dtype,
                         device=group['lr'].device,
                     )
                 if 'exp_avg_lr_2' in p_state:
-                    p_state['exp_avg_lr_2'] = _init_scalar(float(p_state['exp_avg_lr_2']), device='cpu')
+                    p_state['exp_avg_lr_2'] = _init_scalar(p_state['exp_avg_lr_2'], device='cpu')
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -129,9 +125,10 @@ def step(self, closure=None):
 
                 # 1 - beta1 ** state['step']
                 if torch.is_tensor(group['lr']):
+                    lr_safe = torch.where(group['lr'] != 0., group['lr'], torch.ones_like(group['lr']))
                     bias_correction1 = torch.where(
                         group['lr'] != 0.,
-                        state['exp_avg_lr_1'] / group['lr'],
+                        state['exp_avg_lr_1'] / lr_safe,
                         torch.ones_like(group['lr']),
                     )
                 else:

timm/optim/nadamw.py

@@ -87,7 +87,7 @@ def __setstate__(self, state):
                 p_state = self.state.get(p, {})
                 if p_state and 'step' in p_state:
                     p_state['step'] = _init_scalar(
-                        float(p_state['step']),
+                        p_state['step'],
                         device=p.device if group['capturable'] else 'cpu',
                     )
 

timm/optim/rmsprop_tf.py

@@ -94,7 +94,7 @@ def __setstate__(self, state):
             for p in group['params']:
                 p_state = self.state.get(p, {})
                 if p_state and 'step' in p_state:
-                    p_state['step'] = _init_scalar(float(p_state['step']), device='cpu')
+                    p_state['step'] = _init_scalar(p_state['step'], device='cpu')
 
     @torch.no_grad()
     def step(self, closure=None):