Merge pull request #60 from kozistr/feature/list-optimizers

[Feature] List of supported optimizers

Author: kozistr

.github/workflows/ci.yml

@@ -15,27 +15,27 @@ jobs:
 
     steps:
     - uses: actions/checkout@v3
-    - name: set up Python ${{ matrix.python-version }}
+    - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
-    - name: cache pip
+    - name: Cache pip
       uses: actions/cache@v3
       with:
         path: ~/.cache/pip
         key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt') }}
         restore-keys: |
           ${{ runner.os }}-pip-
           ${{ runner.os }}-
-    - name: install dependencies
+    - name: Install dependencies
       run: pip install -r requirements-dev.txt
-    - name: check lint
+    - name: Check lint
       run: make check
-    - name: check test
+    - name: Check test
       run: |
         export PYTHONDONTWRITEBYTECODE=1
         make test
-    - name: check codecov
+    - name: Check codecov
       uses: codecov/codecov-action@v2
       with:
         token: ${{ secrets.CODECOV_TOKEN }}

.github/workflows/publish.yml

@@ -12,7 +12,7 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v3
-      - name: Create Release
+      - name: Create release
         id: create_release
         uses: actions/create-release@v1
         env:
@@ -32,7 +32,7 @@ jobs:
         uses: actions/setup-python@v3
         with:
           python-version: 3.9
-      - name: cache pip
+      - name: Cache pip
         uses: actions/cache@v3
         with:
           path: ~/.cache/pip

.pylintrc

@@ -96,7 +96,7 @@ score=yes
 [REFACTORING]
 
 # Maximum number of nested blocks for function / method body
-max-nested-blocks=5
+max-nested-blocks=6
 
 
 [BASIC]

README.rst

@@ -38,11 +38,11 @@ or you can use optimizer loader, simply passing a name of the optimizer.
 
 ::
 
-    from pytorch_optimizer import load_optimizers
+    from pytorch_optimizer import load_optimizer
 
     ...
     model = YourModel()
-    opt = load_optimizers(optimizer='adamp')
+    opt = load_optimizer(optimizer='adamp')
     optimizer = opt(model.parameters())
     ...
 

lint.py

@@ -14,7 +14,7 @@ def get_configuration() -> Namespace:
     parser.add_argument(
         '-t',
         '--threshold',
-        default=9.98,
+        default=10.0,
         type=float,
     )
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "1.1.0"
+version = "1.1.1"
 description = "Bunch of optimizer implementations in PyTorch with clean-code, strict types. Also, including useful optimization ideas."
 license = "Apache-2.0"
 authors = ["kozistr <[email protected]>"]

pytorch_optimizer/__init__.py

@@ -1,4 +1,8 @@
 # pylint: disable=unused-import
+from typing import Dict, List
+
+from torch.optim import Optimizer
+
 from pytorch_optimizer.adabelief import AdaBelief
 from pytorch_optimizer.adabound import AdaBound
 from pytorch_optimizer.adamp import AdamP
@@ -14,7 +18,6 @@
 from pytorch_optimizer.lookahead import Lookahead
 from pytorch_optimizer.madgrad import MADGRAD
 from pytorch_optimizer.nero import Nero
-from pytorch_optimizer.optimizers import load_optimizers
 from pytorch_optimizer.pcgrad import PCGrad
 from pytorch_optimizer.pnm import PNM
 from pytorch_optimizer.radam import RAdam
@@ -31,3 +34,37 @@
     normalize_gradient,
     unit_norm,
 )
+
+OPTIMIZER_LIST: List = [
+    AdaBelief,
+    AdaBound,
+    AdamP,
+    AdaPNM,
+    DiffGrad,
+    DiffRGrad,
+    Lamb,
+    LARS,
+    MADGRAD,
+    Nero,
+    PNM,
+    RAdam,
+    RaLamb,
+    Ranger,
+    Ranger21,
+    SGDP,
+    Shampoo,
+]
+OPTIMIZERS: Dict[str, Optimizer] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
+
+
+def load_optimizer(optimizer: str) -> OPTIMIZERS:  # pylint: disable=R0911
+    optimizer: str = optimizer.lower()
+
+    if optimizer not in OPTIMIZERS:
+        raise NotImplementedError(f'[-] not implemented optimizer : {optimizer}')
+
+    return OPTIMIZERS[optimizer]
+
+
+def get_supported_optimizers() -> List:
+    return OPTIMIZER_LIST

pytorch_optimizer/fp16.py

@@ -87,7 +87,7 @@ def decrease_loss_scale(self):
 
 
 class SafeFP16Optimizer(Optimizer):
-    def __init__(self, optimizer, aggregate_g_norms: bool = False):
+    def __init__(self, optimizer, aggregate_g_norms: bool = False):  # pylint: disable=super-init-not-called
         self.optimizer = optimizer
         self.aggregate_g_norms = aggregate_g_norms
 

pytorch_optimizer/lookahead.py

@@ -25,7 +25,7 @@ class Lookahead(Optimizer, BaseOptimizer):
           optimizer.step()
     """
 
-    def __init__(
+    def __init__(  # pylint: disable=super-init-not-called
         self,
         optimizer: Optimizer,
         k: int = 5,