Rework get optimizer function (#97)

Author: jettify

requirements-dev.txt

@@ -7,6 +7,7 @@ flake8==3.7.9
 ipdb==0.13.2
 ipython==7.13.0
 mypy==0.770
+numpy==1.18.2
 pyroma==2.6
 pytest-cov==2.8.1
 pytest==5.4.1

tests/test_get.py

@@ -80,7 +80,7 @@ def build_lookahead(*a, **kw):
     build_lookahead,
     optim.Ranger,
     optim.RangerQH,
-    optim.RangerVA
+    optim.RangerVA,
 ]
 
 

tests/test_optimizer.py

@@ -64,9 +64,9 @@ def build_lookahead(*a, **kw):
     (build_lookahead, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
     (optim.QHM, {'lr': 0.1, 'weight_decay': 1e-5, 'momentum': 0.2}, 200),
     (optim.QHAdam, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
-    (optim.Ranger, {'lr': .1, 'weight_decay': 1e-3}, 200),
-    (optim.RangerQH, {'lr': .01, 'weight_decay': 1e-3}, 200),
-    (optim.RangerVA, {'lr': .01, 'weight_decay': 1e-3}, 200),
+    (optim.Ranger, {'lr': 0.1, 'weight_decay': 1e-3}, 200),
+    (optim.RangerQH, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
+    (optim.RangerVA, {'lr': 0.01, 'weight_decay': 1e-3}, 200),
 ]
 
 

tests/test_optimizer_with_nn.py

@@ -1,4 +1,8 @@
-from typing import Optional, Type
+from typing import Optional, Type, List, Dict
+
+from pytorch_ranger import Ranger, RangerQH, RangerVA
+from torch.optim import Optimizer
+
 from .accsgd import AccSGD
 from .adabound import AdaBound
 from .adamod import AdaMod
@@ -12,8 +16,6 @@
 from .radam import RAdam
 from .sgdw import SGDW
 from .yogi import Yogi
-from pytorch_ranger import Ranger, RangerQH, RangerVA
-from torch import optim
 
 
 __all__ = (
@@ -32,12 +34,14 @@
     'Yogi',
     'Ranger',
     'RangerQH',
-    'RangerVA'
+    'RangerVA',
+    # utils
+    'get',
 )
 __version__ = '0.0.1a11'
 
 
-package_opts = [
+_package_opts = [
     AccSGD,
     AdaBound,
     AdaMod,
@@ -54,38 +58,18 @@
     Ranger,
     RangerQH,
     RangerVA,
-]
+]  # type: List[Optimizer]
 
-builtin_opts = [
-    optim.Adadelta,
-    optim.Adagrad,
-    optim.Adam,
-    optim.AdamW,
-    optim.SparseAdam,
-    optim.Adamax,
-    optim.ASGD,
-    optim.SGD,
-    optim.Rprop,
-    optim.RMSprop,
-    optim.LBFGS
-]
 
-NAME_OPTIM_MAP = {
-    opt.__name__.lower(): opt for opt in package_opts + builtin_opts
-}
+_NAME_OPTIM_MAP = {
+    opt.__name__.lower(): opt for opt in _package_opts
+}  # type: Dict[str, Optimizer]
 
 
-def get(name: str,) -> Optional[Type[optim.Optimizer]]:
+def get(name: str) -> Optional[Type[Optimizer]]:
     r"""Returns an optimizer class from its name. Case insensitive.
 
     Args:
         name: the optimizer name.
     """
-    if isinstance(name, str):
-        cls = NAME_OPTIM_MAP.get(name.lower())
-        if cls is None:
-            raise ValueError('Could not interpret optimizer name: ' +
-                             str(name))
-        return cls
-    raise ValueError('Could not interpret optimizer name: ' +
-                     str(name))
+    return _NAME_OPTIM_MAP.get(name.lower())