Update typing in other scheduler classes. Fix spacing in cosine typing. Add some basic scheduler unit tests

Author: rwightman

tests/test_scheduler.py

@@ -35,16 +35,16 @@ def __init__(
             cycle_mul: float = 1.,
             cycle_decay: float = 1.,
             cycle_limit: int = 1,
-            warmup_t: int=0,
-            warmup_lr_init: float=0.,
-            warmup_prefix: bool=False,
-            t_in_epochs: bool=True,
-            noise_range_t: Union[List[int], Tuple[int, int], int, None]=None,
-            noise_pct: float=0.67,
-            noise_std: float=1.0,
-            noise_seed: int=42,
-            k_decay: float=1.0,
-            initialize: bool=True,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            warmup_prefix: bool = False,
+            t_in_epochs: bool = True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: int = 42,
+            k_decay: float = 1.0,
+            initialize: bool = True,
     ) -> None:
         super().__init__(
             optimizer,

timm/scheduler/cosine_lr.py

@@ -5,7 +5,7 @@
 import torch
 import bisect
 from timm.scheduler.scheduler import Scheduler
-from typing import List
+from typing import List, Tuple, Union
 
 class MultiStepLRScheduler(Scheduler):
     """
@@ -16,15 +16,15 @@ def __init__(
             optimizer: torch.optim.Optimizer,
             decay_t: List[int],
             decay_rate: float = 1.,
-            warmup_t=0,
-            warmup_lr_init=0,
-            warmup_prefix=True,
-            t_in_epochs=True,
-            noise_range_t=None,
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=42,
-            initialize=True,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            warmup_prefix: bool = True,
+            t_in_epochs: bool = True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: int = 42,
+            initialize: bool = True,
     ) -> None:
         super().__init__(
             optimizer,

timm/scheduler/multistep_lr.py

@@ -5,7 +5,7 @@
 Hacked together by / Copyright 2020 Ross Wightman
 """
 import torch
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from .scheduler import Scheduler
 
@@ -15,23 +15,22 @@ class PlateauLRScheduler(Scheduler):
 
     def __init__(
             self,
-            optimizer,
-            decay_rate=0.1,
-            patience_t=10,
-            verbose=True,
-            threshold=1e-4,
-            cooldown_t=0,
-            warmup_t=0,
-            warmup_lr_init=0,
-            lr_min=0,
-            mode='max',
-            noise_range_t=None,
-            noise_type='normal',
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=None,
-            initialize=True,
-    ):
+            optimizer: torch.optim.Optimizer,
+            decay_rate: float = 0.1,
+            patience_t: int = 10,
+            threshold: float = 1e-4,
+            cooldown_t: int = 0,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            lr_min: float = 0.,
+            mode: str = 'max',
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_type: str = 'normal',
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: Optional[int] = None,
+            initialize: bool = True,
+    ) -> None:
         super().__init__(
             optimizer,
             'lr',
@@ -47,11 +46,10 @@ def __init__(
             self.optimizer,
             patience=patience_t,
             factor=decay_rate,
-            verbose=verbose,
             threshold=threshold,
             cooldown=cooldown_t,
             mode=mode,
-            min_lr=lr_min
+            min_lr=lr_min,
         )
 
         self.warmup_t = warmup_t
@@ -63,19 +61,19 @@ def __init__(
             self.warmup_steps = [1 for _ in self.base_values]
         self.restore_lr = None
 
-    def state_dict(self):
+    def state_dict(self) -> Dict[str, Any]:
         return {
             'best': self.lr_scheduler.best,
             'last_epoch': self.lr_scheduler.last_epoch,
         }
 
-    def load_state_dict(self, state_dict):
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         self.lr_scheduler.best = state_dict['best']
         if 'last_epoch' in state_dict:
             self.lr_scheduler.last_epoch = state_dict['last_epoch']
 
     # override the base class step fn completely
-    def step(self, epoch, metric=None):
+    def step(self, epoch: int, metric: Optional[float] = None) -> None:
         if epoch <= self.warmup_t:
             lrs = [self.warmup_lr_init + epoch * s for s in self.warmup_steps]
             super().update_groups(lrs)
@@ -88,15 +86,15 @@ def step(self, epoch, metric=None):
 
             # step the base scheduler if metric given
             if metric is not None:
-                self.lr_scheduler.step(metric, epoch)
+                self.lr_scheduler.step(metric)
 
             if self._is_apply_noise(epoch):
                 self._apply_noise(epoch)
 
     def step_update(self, num_updates: int, metric: Optional[float] = None):
         return None
 
-    def _apply_noise(self, epoch):
+    def _apply_noise(self, epoch: int) -> None:
         noise = self._calculate_noise(epoch)
 
         # apply the noise on top of previous LR, cache the old value so we can restore for normal

timm/scheduler/plateau_lr.py

@@ -6,7 +6,7 @@
 """
 import math
 import logging
-from typing import List
+from typing import List, Tuple, Union
 
 import torch
 
@@ -31,16 +31,16 @@ def __init__(
             cycle_mul: float = 1.,
             cycle_decay: float = 1.,
             cycle_limit: int = 1,
-            warmup_t=0,
-            warmup_lr_init=0,
-            warmup_prefix=False,
-            t_in_epochs=True,
-            noise_range_t=None,
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=42,
-            k_decay=1.0,
-            initialize=True,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            warmup_prefix: bool = False,
+            t_in_epochs: bool = True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: int = 42,
+            k_decay: float = 1.0,
+            initialize: bool = True,
     ) -> None:
         super().__init__(
             optimizer,

timm/scheduler/poly_lr.py

@@ -1,6 +1,6 @@
 import abc
 from abc import ABC
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 
@@ -29,11 +29,11 @@ def __init__(
             optimizer: torch.optim.Optimizer,
             param_group_field: str,
             t_in_epochs: bool = True,
-            noise_range_t=None,
-            noise_type='normal',
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=None,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_type: str = 'normal',
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: Optional[int] = None,
             initialize: bool = True,
     ) -> None:
         self.optimizer = optimizer
@@ -81,14 +81,14 @@ def step(self, epoch: int, metric: Optional[float] = None) -> None:
             values = self._add_noise(values, epoch)
             self.update_groups(values)
 
-    def step_update(self, num_updates: int, metric: Optional[float] = None):
+    def step_update(self, num_updates: int, metric: Optional[float] = None) -> None:
         self.metric = metric
         values = self._get_values(num_updates, on_epoch=False)
         if values is not None:
             values = self._add_noise(values, num_updates)
             self.update_groups(values)
 
-    def update_groups(self, values):
+    def update_groups(self, values: Union[float, List[float]]) -> None:
         if not isinstance(values, (list, tuple)):
             values = [values] * len(self.optimizer.param_groups)
         for param_group, value in zip(self.optimizer.param_groups, values):
@@ -97,13 +97,13 @@ def update_groups(self, values):
             else:
                 param_group[self.param_group_field] = value
 
-    def _add_noise(self, lrs, t):
+    def _add_noise(self, lrs: List[float], t: int) -> List[float]:
         if self._is_apply_noise(t):
             noise = self._calculate_noise(t)
             lrs = [v + v * noise for v in lrs]
         return lrs
 
-    def _is_apply_noise(self, t) -> bool:
+    def _is_apply_noise(self, t: int) -> bool:
         """Return True if scheduler in noise range."""
         apply_noise = False
         if self.noise_range_t is not None:

timm/scheduler/scheduler.py

@@ -69,7 +69,7 @@ def create_scheduler_v2(
         cooldown_epochs: int = 0,
         patience_epochs: int = 10,
         decay_rate: float = 0.1,
-        min_lr: float = 0,
+        min_lr: float = 0.,
         warmup_lr: float = 1e-5,
         warmup_epochs: int = 0,
         warmup_prefix: bool = False,

timm/scheduler/scheduler_factory.py

@@ -6,7 +6,7 @@
 """
 import math
 import torch
-from typing import List
+from typing import List, Tuple, Union
 
 
 from .scheduler import Scheduler
@@ -21,15 +21,15 @@ def __init__(
             optimizer: torch.optim.Optimizer,
             decay_t: float,
             decay_rate: float = 1.,
-            warmup_t=0,
-            warmup_lr_init=0,
-            warmup_prefix=True,
-            t_in_epochs=True,
-            noise_range_t=None,
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=42,
-            initialize=True,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            warmup_prefix: bool = True,
+            t_in_epochs: bool = True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: int = 42,
+            initialize: bool = True,
     ) -> None:
         super().__init__(
             optimizer,

timm/scheduler/step_lr.py

@@ -8,7 +8,7 @@
 import math
 import numpy as np
 import torch
-from typing import List
+from typing import List, Tuple, Union
 
 from .scheduler import Scheduler
 
@@ -32,15 +32,15 @@ def __init__(
             cycle_mul: float = 1.,
             cycle_decay: float = 1.,
             cycle_limit: int = 1,
-            warmup_t=0,
-            warmup_lr_init=0,
-            warmup_prefix=False,
-            t_in_epochs=True,
-            noise_range_t=None,
-            noise_pct=0.67,
-            noise_std=1.0,
-            noise_seed=42,
-            initialize=True,
+            warmup_t: int = 0,
+            warmup_lr_init: float = 0.,
+            warmup_prefix: bool = False,
+            t_in_epochs: bool = True,
+            noise_range_t: Union[List[int], Tuple[int, int], int, None] = None,
+            noise_pct: float = 0.67,
+            noise_std: float = 1.0,
+            noise_seed: int = 42,
+            initialize: bool = True,
     ) -> None:
         super().__init__(
             optimizer,