nshtrainer 0.12.1__py3-none-any.whl → 0.13.0__py3-none-any.whl

nshtrainer/callbacks/early_stopping.py

@@ -8,6 +8,7 @@ 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__)
@@ -16,18 +17,12 @@ log = logging.getLogger(__name__)
 class EarlyStoppingConfig(CallbackConfigBase):
     name: Literal["early_stopping"] = "early_stopping"
 
-    monitor: str | None = None
+    metric: MetricConfig | 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.
@@ -52,64 +47,30 @@ class EarlyStoppingConfig(CallbackConfigBase):
 
     @override
     def create_callbacks(self, root_config):
-        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"
-
-        yield EarlyStopping(
-            monitor=monitor,
-            mode=mode,
-            patience=self.patience,
-            min_delta=self.min_delta,
-            min_lr=self.min_lr,
-            strict=self.strict,
-        )
+        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(
@@ -152,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
@@ -167,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.1.dist-info → nshtrainer-0.13.0.dist-info}/METADATA

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

{nshtrainer-0.12.1.dist-info → nshtrainer-0.13.0.dist-info}/RECORD

@@ -15,7 +15,7 @@ nshtrainer/callbacks/checkpoint/_base.py,sha256=yBUZ63P4a3m5jviIUGxBd2WXq_TigpKt
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
 nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
-nshtrainer/callbacks/early_stopping.py,sha256=m-YIWwQAmp1KUNYwJEiuQyNRyX-0pmzBbAQrkH1BLYI,5664
+nshtrainer/callbacks/early_stopping.py,sha256=2vhlXvp6gIFaJXjmPFVK-nbuBUva3200rve5s0krs2c,4618
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=gJC_RUr3ais3FJI0uB6wUZnDdE3WRwCix3ppA3PwQXA,2077
 nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
@@ -82,6 +82,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.12.1.dist-info/METADATA,sha256=5pwyF0q1y7DxG3HNgRPJ3vyVG3XuctB9KyR2XlrW9m0,860
-nshtrainer-0.12.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.12.1.dist-info/RECORD,,
+nshtrainer-0.13.0.dist-info/METADATA,sha256=1OsZvE4GYPu85jiPrBJffrqGvsmdUuQU96TGIYnlIoo,860
+nshtrainer-0.13.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.13.0.dist-info/RECORD,,