nshtrainer 0.12.0__tar.gz → 0.13.0__tar.gz

{nshtrainer-0.12.0 → nshtrainer-0.13.0}/PKG-INFO

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

{nshtrainer-0.12.0 → nshtrainer-0.13.0}/pyproject.toml

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

{nshtrainer-0.12.0 → nshtrainer-0.13.0}/src/nshtrainer/callbacks/__init__.py

@@ -13,6 +13,7 @@ from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )
 from .early_stopping import EarlyStopping as EarlyStopping
+from .early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
 from .ema import EMA as EMA
 from .ema import EMAConfig as EMAConfig
 from .finite_checks import FiniteChecksCallback as FiniteChecksCallback
@@ -34,7 +35,8 @@ from .wandb_watch import WandbWatchCallback as WandbWatchCallback
 from .wandb_watch import WandbWatchConfig as WandbWatchConfig
 
 CallbackConfig = Annotated[
-    ThroughputMonitorConfig
+    EarlyStoppingConfig
+    | ThroughputMonitorConfig
     | EpochTimerConfig
     | PrintTableMetricsConfig
     | FiniteChecksConfig

{nshtrainer-0.12.0 → nshtrainer-0.13.0}/src/nshtrainer/callbacks/early_stopping.py

@@ -1,5 +1,6 @@
 import logging
 import math
+from typing import Literal
 
 from lightning.fabric.utilities.rank_zero import _get_rank
 from lightning.pytorch import Trainer
@@ -7,41 +8,69 @@ from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
 from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
 from typing_extensions import override
 
+from ..metrics._config import MetricConfig
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+class EarlyStoppingConfig(CallbackConfigBase):
+    name: Literal["early_stopping"] = "early_stopping"
+
+    metric: MetricConfig | None = None
+    """
+    The metric to monitor for early stopping.
+    If None, the primary metric will be used.
+    """
+
+    patience: int
+    """
+    Number of epochs with no improvement after which training will be stopped.
+    """
+
+    min_delta: float = 1.0e-8
+    """
+    Minimum change in the monitored quantity to qualify as an improvement.
+    """
+
+    min_lr: float | None = None
+    """
+    Minimum learning rate. If the learning rate of the model is less than this value,
+    the training will be stopped.
+    """
+
+    strict: bool = True
+    """
+    Whether to enforce that the monitored quantity must improve by at least `min_delta`
+    to qualify as an improvement.
+    """
+
+    @override
+    def create_callbacks(self, root_config):
+        if (metric := self.metric) is None and (
+            metric := root_config.primary_metric
+        ) is None:
+            raise ValueError(
+                "Either `metric` or `root_config.primary_metric` must be set."
+            )
+
+        yield EarlyStopping(self, metric)
+
+
 class EarlyStopping(_EarlyStopping):
-    def __init__(
-        self,
-        monitor: str,
-        min_delta: float = 0,
-        min_lr: float | None = None,
-        patience: int = 3,
-        verbose: bool = True,
-        mode: str = "min",
-        strict: bool = True,
-        check_finite: bool = True,
-        stopping_threshold: float | None = None,
-        divergence_threshold: float | None = None,
-        check_on_train_epoch_end: bool | None = None,
-        log_rank_zero_only: bool = False,
-    ):
+    def __init__(self, config: EarlyStoppingConfig, metric: MetricConfig):
+        self.config = config
+        self.metric = metric
+        del config, metric
+
         super().__init__(
-            monitor,
-            min_delta,
-            patience,
-            verbose,
-            mode,
-            strict,
-            check_finite,
-            stopping_threshold,
-            divergence_threshold,
-            check_on_train_epoch_end,
-            log_rank_zero_only,
+            monitor=self.metric.validation_monitor,
+            mode=self.metric.mode,
+            patience=self.patience,
+            min_delta=self.min_delta,
+            strict=self.strict,
         )
 
-        self.min_lr = min_lr
-
     @override
     @staticmethod
     def _log_info(
@@ -84,7 +113,7 @@ class EarlyStopping(_EarlyStopping):
     def _evaluate_stopping_criteria_min_lr(
         self, trainer: Trainer
     ) -> tuple[bool, str | None]:
-        if self.min_lr is None:
+        if self.config.min_lr is None:
             return False, None
 
         # Get the maximum LR across all param groups in all optimizers
@@ -99,13 +128,13 @@ class EarlyStopping(_EarlyStopping):
             return False, None
 
         # If the maximum LR is less than the minimum LR, stop training
-        if model_max_lr >= self.min_lr:
+        if model_max_lr >= self.config.min_lr:
             return False, None
 
         return True, (
             "Stopping threshold reached: "
             f"The maximum LR of the model across all param groups is {model_max_lr:.2e} "
-            f"which is less than the minimum LR {self.min_lr:.2e}"
+            f"which is less than the minimum LR {self.config.min_lr:.2e}"
         )
 
     def on_early_stopping(self, trainer: Trainer):

{nshtrainer-0.12.0 → nshtrainer-0.13.0}/src/nshtrainer/model/config.py

@@ -38,6 +38,7 @@ from .._checkpoint.loader import CheckpointLoadingConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
+    EarlyStoppingConfig,
     LastCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
     WandbWatchConfig,
@@ -1043,73 +1044,6 @@ class LightningTrainerKwargs(TypedDict, total=False):
     """
 
 
-class EarlyStoppingConfig(CallbackConfigBase):
-    monitor: str | None = None
-    """
-    The metric to monitor for early stopping.
-    If None, the primary metric will be used.
-    """
-
-    mode: Literal["min", "max"] | None = None
-    """
-    The mode for the metric to monitor for early stopping.
-    If None, the primary metric mode will be used.
-    """
-
-    patience: int
-    """
-    Number of epochs with no improvement after which training will be stopped.
-    """
-
-    min_delta: float = 1.0e-8
-    """
-    Minimum change in the monitored quantity to qualify as an improvement.
-    """
-
-    min_lr: float | None = None
-    """
-    Minimum learning rate. If the learning rate of the model is less than this value,
-    the training will be stopped.
-    """
-
-    strict: bool = True
-    """
-    Whether to enforce that the monitored quantity must improve by at least `min_delta`
-    to qualify as an improvement.
-    """
-
-    @override
-    def create_callbacks(self, root_config: "BaseConfig"):
-        from ..callbacks.early_stopping import EarlyStopping
-
-        monitor = self.monitor
-        mode = self.mode
-        if monitor is None:
-            assert mode is None, "If `monitor` is not provided, `mode` must be None."
-
-            primary_metric = root_config.primary_metric
-            if primary_metric is None:
-                raise ValueError(
-                    "No primary metric is set, so `monitor` must be provided in `early_stopping`."
-                )
-            monitor = primary_metric.validation_monitor
-            mode = primary_metric.mode
-
-        if mode is None:
-            mode = "min"
-
-        return [
-            EarlyStopping(
-                monitor=monitor,
-                mode=mode,
-                patience=self.patience,
-                min_delta=self.min_delta,
-                min_lr=self.min_lr,
-                strict=self.strict,
-            )
-        ]
-
-
 class SanityCheckingConfig(C.Config):
     reduce_lr_on_plateau: Literal["disable", "error", "warn"] = "error"
     """