Merge pull request #79 from kozistr/feature/lr_scheduler

[Feature] Supports LR Schedulers

Author: kozistr

README.rst

@@ -29,7 +29,7 @@ Install
 
 ::
 
-    $ pip3 install pytorch-optimizer
+    $ pip3 install -U pytorch-optimizer
 
 Simple Usage
 ~~~~~~~~~~~~

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_optimizer"
-version = "1.3.2"
+version = "2.0.0"
 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]>"]
@@ -9,7 +9,7 @@ readme = "README.rst"
 homepage = "https://github.com/kozistr/pytorch_optimizer"
 repository = "https://github.com/kozistr/pytorch_optimizer"
 documentation = "https://pytorch-optimizers.readthedocs.io/en/latest"
-keywords = ["pytorch", "deep-learning", "optimizer"]
+keywords = ["pytorch", "deep-learning", "optimizer", "lr scheduler"]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Development Status :: 5 - Production/Stable",

pytorch_optimizer/__init__.py

@@ -1,42 +1,42 @@
 # pylint: disable=unused-import
-from typing import Dict, List, Type
+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
-from pytorch_optimizer.adan import Adan
-from pytorch_optimizer.adapnm import AdaPNM
-from pytorch_optimizer.agc import agc
-from pytorch_optimizer.chebyshev_schedule import get_chebyshev_schedule
-from pytorch_optimizer.diffgrad import DiffGrad
-from pytorch_optimizer.diffrgrad import DiffRGrad
-from pytorch_optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
-from pytorch_optimizer.gc import centralize_gradient
-from pytorch_optimizer.lamb import Lamb
-from pytorch_optimizer.lars import LARS
-from pytorch_optimizer.lookahead import Lookahead
-from pytorch_optimizer.madgrad import MADGRAD
-from pytorch_optimizer.nero import Nero
-from pytorch_optimizer.pcgrad import PCGrad
-from pytorch_optimizer.pnm import PNM
-from pytorch_optimizer.radam import RAdam
-from pytorch_optimizer.ralamb import RaLamb
-from pytorch_optimizer.ranger import Ranger
-from pytorch_optimizer.ranger21 import Ranger21
-from pytorch_optimizer.sam import SAM
-from pytorch_optimizer.sgdp import SGDP
-from pytorch_optimizer.shampoo import Shampoo
-from pytorch_optimizer.utils import (
+from pytorch_optimizer.base.types import OPTIMIZER, SCHEDULER
+from pytorch_optimizer.lr_scheduler.chebyshev import get_chebyshev_schedule
+from pytorch_optimizer.lr_scheduler.cosine_anealing import CosineAnnealingWarmupRestarts
+from pytorch_optimizer.optimizer.adabelief import AdaBelief
+from pytorch_optimizer.optimizer.adabound import AdaBound
+from pytorch_optimizer.optimizer.adamp import AdamP
+from pytorch_optimizer.optimizer.adan import Adan
+from pytorch_optimizer.optimizer.adapnm import AdaPNM
+from pytorch_optimizer.optimizer.agc import agc
+from pytorch_optimizer.optimizer.diffgrad import DiffGrad
+from pytorch_optimizer.optimizer.diffrgrad import DiffRGrad
+from pytorch_optimizer.optimizer.fp16 import DynamicLossScaler, SafeFP16Optimizer
+from pytorch_optimizer.optimizer.gc import centralize_gradient
+from pytorch_optimizer.optimizer.lamb import Lamb
+from pytorch_optimizer.optimizer.lars import LARS
+from pytorch_optimizer.optimizer.lookahead import Lookahead
+from pytorch_optimizer.optimizer.madgrad import MADGRAD
+from pytorch_optimizer.optimizer.nero import Nero
+from pytorch_optimizer.optimizer.pcgrad import PCGrad
+from pytorch_optimizer.optimizer.pnm import PNM
+from pytorch_optimizer.optimizer.radam import RAdam
+from pytorch_optimizer.optimizer.ralamb import RaLamb
+from pytorch_optimizer.optimizer.ranger import Ranger
+from pytorch_optimizer.optimizer.ranger21 import Ranger21
+from pytorch_optimizer.optimizer.sam import SAM
+from pytorch_optimizer.optimizer.sgdp import SGDP
+from pytorch_optimizer.optimizer.shampoo import Shampoo
+from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     get_optimizer_parameters,
     matrix_power,
     normalize_gradient,
     unit_norm,
 )
 
-OPTIMIZER_LIST: List[Type[Optimizer]] = [
+OPTIMIZER_LIST: List[OPTIMIZER] = [
     AdaBelief,
     AdaBound,
     AdamP,
@@ -56,10 +56,17 @@
     SGDP,
     Shampoo,
 ]
-OPTIMIZERS: Dict[str, Type[Optimizer]] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
+OPTIMIZERS: Dict[str, OPTIMIZER] = {str(optimizer.__name__).lower(): optimizer for optimizer in OPTIMIZER_LIST}
+
+LR_SCHEDULER_LIST: List[SCHEDULER] = [
+    CosineAnnealingWarmupRestarts,
+]
+LR_SCHEDULERS: Dict[str, SCHEDULER] = {
+    str(lr_scheduler.__name__).lower(): lr_scheduler for lr_scheduler in LR_SCHEDULER_LIST
+}
 
 
-def load_optimizer(optimizer: str) -> Type[Optimizer]:
+def load_optimizer(optimizer: str) -> OPTIMIZER:
     optimizer: str = optimizer.lower()
 
     if optimizer not in OPTIMIZERS:
@@ -68,5 +75,18 @@ def load_optimizer(optimizer: str) -> Type[Optimizer]:
     return OPTIMIZERS[optimizer]
 
 
-def get_supported_optimizers() -> List[Type[Optimizer]]:
+def load_lr_scheduler(lr_scheduler: str) -> SCHEDULER:
+    lr_scheduler: str = lr_scheduler.lower()
+
+    if lr_scheduler not in LR_SCHEDULERS:
+        raise NotImplementedError(f'[-] not implemented lr_scheduler : {lr_scheduler}')
+
+    return LR_SCHEDULERS[lr_scheduler]
+
+
+def get_supported_optimizers() -> List[OPTIMIZER]:
     return OPTIMIZER_LIST
+
+
+def get_supported_lr_schedulers() -> List[SCHEDULER]:
+    return LR_SCHEDULER_LIST

pytorch_optimizer/base/__init__.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from pytorch_optimizer.types import BETAS
+from pytorch_optimizer.base.types import BETAS
 
 
 class BaseOptimizer(ABC):

pytorch_optimizer/base_optimizer.py renamed to pytorch_optimizer/base/base_optimizer.py

@@ -1,10 +1,14 @@
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
 
 import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 
 CLOSURE = Optional[Callable[[], float]]
 LOSS = Optional[float]
 BETAS = Union[Tuple[float, float], Tuple[float, float, float]]
 DEFAULTS = Dict[str, Any]
 PARAMETERS = Optional[Union[Iterable[Dict[str, Any]], Iterable[torch.Tensor]]]
 STATE = Dict[str, Any]
+OPTIMIZER = Type[Optimizer]
+SCHEDULER = Type[_LRScheduler]

pytorch_optimizer/types.py renamed to pytorch_optimizer/base/types.py

@@ -0,0 +1,45 @@
+from torch import nn
+
+from pytorch_optimizer.base.types import PARAMETERS
+
+
+def deberta_v3_large_lr_scheduler(
+    model: nn.Module,
+    head_param_start: int = 390,
+    base_lr: float = 2e-5,
+    head_lr: float = 1e-4,
+    wd: float = 1e-2,
+) -> PARAMETERS:
+    """DeBERTa-v3 large layer-wise lr scheduler
+    Reference : https://github.com/gilfernandes/commonlit
+
+    :param model: nn.Module. model. based on Huggingface Transformers.
+    :param head_param_start: int. where the backbone ends (head starts)
+    :param base_lr: float. base lr
+    :param head_lr: float. head_lr
+    :param wd: float. weight decay
+    """
+    named_parameters = list(model.named_parameters())
+
+    backbone_parameters = named_parameters[:head_param_start]
+    regressor_parameters = named_parameters[head_param_start:]
+
+    regressor_group = [params for (_, params) in regressor_parameters]
+
+    parameters = [{'params': regressor_group, 'lr': head_lr}]
+
+    layer_low_threshold: int = 195  # start of the 12 layers
+    layer_middle_threshold: int = 323  # end of the 24 layers
+
+    for layer_num, (name, params) in enumerate(backbone_parameters):
+        weight_decay: float = 0.0 if ('bias' in name) or ('LayerNorm.weight' in name) else wd
+
+        lr = base_lr / 2.5  # 2e-5
+        if layer_num >= layer_middle_threshold:
+            lr = base_lr / 0.5  # 1e-4
+        elif layer_num >= layer_low_threshold:
+            lr = base_lr
+
+        parameters.append({'params': params, 'weight_decay': weight_decay, 'lr': lr})
+
+    return parameters