Merge pull request #98 from kozistr/test/coverage

[Test] Coverage

Author: kozistr

docs/scheduler_api.rst

@@ -56,3 +56,11 @@ OneCycleLR
 
 .. autoclass:: pytorch_optimizer.OneCycleLR
     :members:
+
+.. _deberta_v3_large_lr_scheduler:
+
+deberta_v3_large_lr_scheduler
+-----------------------------
+
+.. autoclass:: pytorch_optimizer.deberta_v3_large_lr_scheduler
+    :members:

pyproject.toml

@@ -52,8 +52,8 @@ url = "https://download.pytorch.org/whl/cpu"
 secondary = true
 
 [tool.ruff]
-select = ["A", "B", "C4", "E", "F", "G", "I", "N", "S", "T", "ISC", "W", "INP", "PIE", "T20", "RET", "SIM", "ARG"]
-ignore = []
+select = ["A", "B", "C4", "D", "E", "F", "G", "I", "N", "S", "T", "ISC", "ICN", "W", "INP", "PIE", "T20", "RET", "SIM", "TID", "ARG", "ERA", "RUF", "YTT", "PL"]
+ignore = ["D100", "D102", "D104", "D105", "D107", "D203", "D213", "PIE790", "PLR2004"]
 fixable = ["A", "B", "C", "D", "E", "F"]
 unfixable = ["F401"]
 exclude = [
@@ -77,20 +77,21 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 target-version = "py39"
 
 [tool.ruff.per-file-ignores]
-"./hubconf.py" = ["INP001"]
-"./tests/test_utils.py" = ["S101"]
-"./tests/test_gradients.py" = ["S101"]
-"./tests/test_optimizers.py" = ["S101"]
-"./tests/test_optimizer_parameters.py" = ["S101"]
-"./tests/test_load_optimizers.py" = ["S101"]
-"./tests/test_load_lr_schedulers.py" = ["S101"]
-"./tests/test_lr_scheduler_parameters.py" = ["S101"]
+"./hubconf.py" = ["D", "INP001"]
+"./tests/__init__.py" = ["D"]
+"./tests/constants.py" = ["D"]
+"./tests/utils.py" = ["D"]
+"./tests/test_utils.py" = ["D", "S101"]
+"./tests/test_gradients.py" = ["D", "S101"]
+"./tests/test_optimizers.py" = ["D", "S101"]
+"./tests/test_optimizer_parameters.py" = ["D", "S101"]
+"./tests/test_load_optimizers.py" = ["D", "S101"]
+"./tests/test_load_lr_schedulers.py" = ["D", "S101"]
+"./tests/test_lr_schedulers.py" = ["D"]
+"./tests/test_lr_scheduler_parameters.py" = ["D", "S101"]
 "./pytorch_optimizer/__init__.py" = ["F401"]
 "./pytorch_optimizer/lr_scheduler/__init__.py" = ["F401"]
 
-[tool.ruff.mccabe]
-max-complexity = 10
-
 [tool.coverage.run]
 omit = [
     "./pytorch_optimizer/optimizer/gsam.py",

pytorch_optimizer/__init__.py

@@ -2,6 +2,7 @@
 from typing import Dict, List
 
 from pytorch_optimizer.base.types import OPTIMIZER, SCHEDULER
+from pytorch_optimizer.experimental.deberta_v3_lr_scheduler import deberta_v3_large_lr_scheduler
 from pytorch_optimizer.lr_scheduler import (
     ConstantLR,
     CosineAnnealingLR,

pytorch_optimizer/base/exception.py

@@ -1,5 +1,5 @@
 class NoSparseGradientError(Exception):
-    """Raised when the gradient is sparse gradient
+    """Raised when the gradient is sparse gradient.
 
     :param optimizer_name: str. optimizer name.
     :param note: str. special conditions to note (default '').
@@ -12,23 +12,23 @@ def __init__(self, optimizer_name: str, note: str = ''):
 
 
 class ZeroParameterSizeError(Exception):
-    """Raised when the parameter size is 0"""
+    """Raised when the parameter size is 0."""
 
     def __init__(self):
         self.message: str = '[-] parameter size is 0'
         super().__init__(self.message)
 
 
 class NoClosureError(Exception):
-    """Raised when there's no closure function"""
+    """Raised when there's no closure function."""
 
     def __init__(self, optimizer_name: str):
         self.message: str = f'[-] {optimizer_name} requires closure.'
         super().__init__(self.message)
 
 
 class NegativeLRError(Exception):
-    """Raised when learning rate is negative"""
+    """Raised when learning rate is negative."""
 
     def __init__(self, lr: float, lr_type: str = ''):
         self.note: str = 'learning rate' if lr_type == '' else lr_type
@@ -37,7 +37,7 @@ def __init__(self, lr: float, lr_type: str = ''):
 
 
 class NegativeStepError(Exception):
-    """Raised when step is negative"""
+    """Raised when step is negative."""
 
     def __init__(self, num_steps: int, step_type: str = ''):
         self.note: str = 'step' if step_type == '' else step_type

pytorch_optimizer/base/optimizer.py

@@ -7,6 +7,8 @@
 
 
 class BaseOptimizer(ABC):
+    r"""Base optimizer class."""
+
     @staticmethod
     def validate_learning_rate(learning_rate: float):
         if learning_rate < 0.0:
@@ -90,7 +92,7 @@ def validate_reduction(reduction: str):
     @staticmethod
     def validate_update_frequency(update_frequency: int):
         if update_frequency < 1:
-            raise NegativeStepError(f'[-] update_frequency {update_frequency} must be positive')
+            raise NegativeStepError(update_frequency, step_type='update_frequency')
 
     @staticmethod
     def validate_norm(norm: float):

pytorch_optimizer/base/scheduler.py

@@ -6,7 +6,9 @@
 
 
 class BaseLinearWarmupScheduler(ABC):
-    r"""BaseLinearWarmupScheduler class. The LR Scheduler class based on this class has linear warmup strategy.
+    r"""BaseLinearWarmupScheduler class.
+
+        The LR Scheduler class based on this class has linear warmup strategy.
 
     :param optimizer: Optimizer. OPTIMIZER. It will set learning rate to all trainable parameters in optimizer.
     :param t_max: int. total steps to train.

pytorch_optimizer/experimental/deberta_v3_lr_scheduler.py

@@ -5,31 +5,33 @@
 
 def deberta_v3_large_lr_scheduler(
     model: nn.Module,
+    layer_low_threshold: int = 195,
+    layer_middle_threshold: int = 323,
     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
+    """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
+    :param layer_low_threshold: int. start of the 12 layers.
+    :param layer_middle_threshold: int. end of the 24 layers.
+    :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]
+    head_parameters = named_parameters[head_param_start:]
 
-    parameters = [{'params': regressor_group, 'lr': head_lr}]
+    head_group = [params for (_, params) in head_parameters]
 
-    layer_low_threshold: int = 195  # start of the 12 layers
-    layer_middle_threshold: int = 323  # end of the 24 layers
+    parameters = [{'params': head_group, 'lr': head_lr}]
 
     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

pytorch_optimizer/lr_scheduler/chebyshev.py

@@ -2,28 +2,29 @@
 
 
 def chebyshev_steps(small_m: float, big_m: float, num_epochs: int) -> np.ndarray:
-    r"""chebyshev_steps
+    r"""Chebyshev steps.
 
     :param small_m: float. stands for 'm' notation.
     :param big_m:  float. stands for 'M' notation.
     :param num_epochs: int. stands for 'T' notation.
     :return: np.array. chebyshev_steps.
     """
-
     c, r = (big_m + small_m) / 2.0, (big_m - small_m) / 2.0
     thetas = (np.arange(num_epochs) + 0.5) / num_epochs * np.pi
 
     return 1.0 / (c - r * np.cos(thetas))
 
 
 def chebyshev_perm(num_epochs: int) -> np.ndarray:
+    r"""Chebyshev permutation."""
     perm = np.array([0])
     while len(perm) < num_epochs:
         perm = np.vstack([perm, 2 * len(perm) - 1 - perm]).T.flatten()
     return perm
 
 
 def get_chebyshev_schedule(num_epochs: int) -> np.ndarray:
+    r"""Get Chebyshev schedules."""
     steps: np.ndarray = chebyshev_steps(0.1, 1, num_epochs - 2)
     perm: np.ndarray = chebyshev_perm(num_epochs - 2)
     return steps[perm]

pytorch_optimizer/lr_scheduler/cosine_anealing.py

@@ -7,7 +7,7 @@
 
 
 class CosineAnnealingWarmupRestarts(_LRScheduler):
-    r"""CosineAnnealingWarmupRestarts
+    r"""CosineAnnealingWarmupRestarts.
 
     :param optimizer: Optimizer. wrapped optimizer instance.
     :param first_cycle_steps: int. first cycle step size.

pytorch_optimizer/lr_scheduler/linear_warmup.py

@@ -6,20 +6,24 @@
 
 
 class LinearScheduler(BaseLinearWarmupScheduler):
+    r"""Linear LR Scheduler w/ linear warmup."""
+
     def _step(self) -> float:
         return self.max_lr + (self.min_lr - self.max_lr) * (self.step_t - self.warmup_steps) / (
             self.total_steps - self.warmup_steps
         )
 
 
 class CosineScheduler(BaseLinearWarmupScheduler):
+    r"""Cosine LR Scheduler w/ linear warmup."""
+
     def _step(self) -> float:
         phase: float = (self.step_t - self.warmup_steps) / (self.total_steps - self.warmup_steps) * math.pi
         return self.min_lr + (self.max_lr - self.min_lr) * (np.cos(phase) + 1.0) / 2.0
 
 
 class PolyScheduler(BaseLinearWarmupScheduler):
-    r"""Poly LR Scheduler
+    r"""Poly LR Scheduler.
 
     :param: poly_order: float. lr scheduler decreases with steps.
     """