nshtrainer 0.43.0__tar.gz → 0.44.1__tar.gz

{nshtrainer-0.43.0 → nshtrainer-0.44.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.43.0
+Version: 0.44.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.43.0 → nshtrainer-0.44.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.43.0"
+version = "0.44.1"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.43.0 → nshtrainer-0.44.1}/src/nshtrainer/lr_scheduler/_base.py

@@ -11,7 +11,7 @@ from lightning.pytorch.utilities.types import (
     LRSchedulerTypeUnion,
 )
 from torch.optim import Optimizer
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import Never, NotRequired, TypedDict
 
 if TYPE_CHECKING:
     from ..model.base import LightningModuleBase
@@ -44,20 +44,18 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
     @abstractmethod
     def create_scheduler_impl(
-        self,
-        optimizer: Optimizer,
-        lightning_module: "LightningModuleBase",
-        lr: float,
+        self, optimizer: Optimizer, lightning_module: LightningModuleBase
     ) -> LRSchedulerTypeUnion | LRSchedulerConfigType: ...
 
     def create_scheduler(
         self,
         optimizer: Optimizer,
-        lightning_module: "LightningModuleBase",
-        lr: float,
+        lightning_module: LightningModuleBase,
+        lr: Never
+        | None = None,  # Backward compatibility, should be removed in the future
     ) -> LRSchedulerConfigType:
         # Create the scheduler.
-        scheduler = self.create_scheduler_impl(optimizer, lightning_module, lr)
+        scheduler = self.create_scheduler_impl(optimizer, lightning_module)
 
         # If the scheduler is not a `LRSchedulerConfigType`, then make it one.
         if not isinstance(scheduler, Mapping):
@@ -89,9 +87,7 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
         return scheduler
 
-    def compute_num_steps_per_epoch(
-        self, lightning_module: "LightningModuleBase"
-    ) -> int:
+    def compute_num_steps_per_epoch(self, lightning_module: LightningModuleBase) -> int:
         trainer = lightning_module.trainer
         # Use the Lightning trainer to convert the epoch-based values to step-based values
         _ = trainer.estimated_stepping_batches

{nshtrainer-0.43.0 → nshtrainer-0.44.1}/src/nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -20,21 +20,21 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
         optimizer: Optimizer,
         warmup_epochs: int,
         max_epochs: int,
-        warmup_start_lr: float = 0.0,
-        eta_min: float = 0.0,
+        warmup_start_lr_factor: float = 0.0,
+        eta_min_factor: float = 0.0,
         last_epoch: int = -1,
         should_restart: bool = True,
     ) -> None:
         self.warmup_epochs = warmup_epochs
         self.max_epochs = max_epochs
-        self.warmup_start_lr = warmup_start_lr
-        self.eta_min = eta_min
+        self.warmup_start_lr_factor = warmup_start_lr_factor
+        self.eta_min_factor = eta_min_factor
         self.should_restart = should_restart
 
         super().__init__(optimizer, last_epoch)
 
     @override
-    def get_lr(self) -> list[float]:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def get_lr(self) -> list[float]:
         if not self._get_lr_called_within_step:
             warnings.warn(
                 "To get the last learning rate computed by the scheduler, "
@@ -43,25 +43,26 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
             )
 
         if self.last_epoch == 0:
-            return [self.warmup_start_lr] * len(self.base_lrs)
+            return [self.warmup_start_lr_factor * base_lr for base_lr in self.base_lrs]
         if self.last_epoch < self.warmup_epochs:
             return [
                 group["lr"]
-                + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
+                + (base_lr - self.warmup_start_lr_factor * base_lr)
+                / (self.warmup_epochs - 1)
                 for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
             ]
         if self.last_epoch == self.warmup_epochs:
             return self.base_lrs
 
         if not self.should_restart and self.last_epoch >= self.max_epochs:
-            return [self.eta_min] * len(self.base_lrs)
+            return [self.eta_min_factor * base_lr for base_lr in self.base_lrs]
 
         if (self.last_epoch - 1 - self.max_epochs) % (
             2 * (self.max_epochs - self.warmup_epochs)
         ) == 0:
             return [
                 group["lr"]
-                + (base_lr - self.eta_min)
+                + (base_lr - self.eta_min_factor * base_lr)
                 * (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs)))
                 / 2
                 for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
@@ -84,9 +85,9 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
                     / (self.max_epochs - self.warmup_epochs)
                 )
             )
-            * (group["lr"] - self.eta_min)
-            + self.eta_min
-            for group in self.optimizer.param_groups
+            * (group["lr"] - self.eta_min_factor * base_lr)
+            + self.eta_min_factor * base_lr
+            for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
         ]
 
 
@@ -121,12 +122,10 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
         }
 
     @override
-    def create_scheduler_impl(self, optimizer, lightning_module, lr):
+    def create_scheduler_impl(self, optimizer, lightning_module):
         num_steps_per_epoch = self.compute_num_steps_per_epoch(lightning_module)
         warmup_steps = self.warmup_duration.to_steps(num_steps_per_epoch).value
         max_steps = self.max_duration.to_steps(num_steps_per_epoch).value
-        warmup_start_lr = self.warmup_start_lr_factor * lr
-        min_lr = self.min_lr_factor * lr
 
         # Warmup and max steps should be at least 1.
         warmup_steps = max(warmup_steps, 1)
@@ -137,8 +136,8 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
             optimizer=optimizer,
             warmup_epochs=warmup_steps,
             max_epochs=max_steps,
-            warmup_start_lr=warmup_start_lr,
-            eta_min=min_lr,
+            warmup_start_lr_factor=self.warmup_start_lr_factor,
+            eta_min_factor=self.min_lr_factor,
             should_restart=self.annealing,
         )
         return scheduler

{nshtrainer-0.43.0 → nshtrainer-0.44.1}/src/nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -22,21 +22,21 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
     """Metric to monitor.
     If not provided, the primary metric of the runner will be used."""
 
-    patience: int = 10
+    patience: int
     r"""Number of epochs with no improvement after which learning rate will be reduced."""
 
-    factor: float = 0.1
+    factor: float
     r"""Factor by which the learning rate will be reduced. new_lr = lr * factor."""
 
+    cooldown: int = 0
+    r"""Number of epochs to wait before resuming normal operation after lr has been reduced."""
+
     min_lr: float | list[float] = 0.0
     r"""A scalar or a list of scalars. A lower bound on the learning rate of all param groups or each group respectively."""
 
     eps: float = 1.0e-8
     r"""Minimal decay applied to lr. If the difference between new and old lr is smaller than eps, the update is ignored."""
 
-    cooldown: int = 0
-    r"""Number of epochs to wait before resuming normal operation after lr has been reduced."""
-
     threshold: float = 1.0e-4
     r"""Threshold for measuring the new optimum, to only focus on significant changes."""
 
@@ -45,7 +45,7 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
 
     @override
     def create_scheduler_impl(
-        self, optimizer, lightning_module, lr
+        self, optimizer, lightning_module
     ) -> LRSchedulerConfigType:
         if (metric := self.metric) is None:
             lm_config = cast("BaseConfig", lightning_module.config)