Merge pull request #73 from Bing-su/hubconf

[Feature] support torch.hub.load

Author: kozistr

hubconf.py

@@ -0,0 +1,15 @@
+dependencies = ['torch']
+
+from functools import partial as _partial, update_wrapper as _update_wrapper
+
+from pytorch_optimizer import (
+    get_supported_optimizers as _get_supported_optimizers,
+    load_optimizer as _load_optimizer,
+)
+
+for optimizer in _get_supported_optimizers():
+    name = optimizer.__name__
+    for n in (name, name.lower()):
+        func = _partial(_load_optimizer, optimizer=n)
+        _update_wrapper(func, optimizer)
+        globals()[n] = func

pytorch_optimizer/__init__.py

@@ -1,5 +1,7 @@
 # pylint: disable=unused-import
-from typing import Callable, Dict, List
+from typing import Dict, List, Type
+
+from torch.optim import Optimizer
 
 from pytorch_optimizer.adabelief import AdaBelief
 from pytorch_optimizer.adabound import AdaBound
@@ -54,10 +56,10 @@
     SGDP,
     Shampoo,
 ]
-OPTIMIZERS: Dict[str, Callable] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
+OPTIMIZERS: Dict[str, Type[Optimizer]] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
 
 
-def load_optimizer(optimizer: str) -> Callable:
+def load_optimizer(optimizer: str) -> Type[Optimizer]:
     optimizer: str = optimizer.lower()
 
     if optimizer not in OPTIMIZERS:
@@ -66,5 +68,5 @@ def load_optimizer(optimizer: str) -> Callable:
     return OPTIMIZERS[optimizer]
 
 
-def get_supported_optimizers() -> List:
+def get_supported_optimizers() -> List[Type[Optimizer]]:
     return OPTIMIZER_LIST