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

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/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,9 +8,76 @@ from lightning.pytorch.callbacks import EarlyStopping as _EarlyStopping
 from lightning.pytorch.utilities.rank_zero import rank_prefixed_message
 from typing_extensions import override
 
+from .base import CallbackConfigBase
+
 log = logging.getLogger(__name__)
 
 
+class EarlyStoppingConfig(CallbackConfigBase):
+    name: Literal["early_stopping"] = "early_stopping"
+
+    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):
+        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,
+        )
+
+
 class EarlyStopping(_EarlyStopping):
     def __init__(
         self,

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"
     """

{nshtrainer-0.12.0.dist-info → nshtrainer-0.12.1.dist-info}/METADATA

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

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

@@ -6,7 +6,7 @@ nshtrainer/_experimental/__init__.py,sha256=2tQIcrWT8U8no_AeBTYnozaTmxN40kuAJdGQ
 nshtrainer/_experimental/flops/__init__.py,sha256=edo9Ez3LlrnxkNRX9W6YBhPkRPKYGLpkpnl5gx7sEX8,1550
 nshtrainer/_experimental/flops/flop_counter.py,sha256=-sL0Fy6poXa__hyzUMdZScjPULp4coQELQpPU6p6dXU,25736
 nshtrainer/_experimental/flops/module_tracker.py,sha256=bUL-IRTd0aF_DwmXkZjHZAA31p4ZEhyqhc26XWKQUUY,4922
-nshtrainer/callbacks/__init__.py,sha256=k-DbpIlH2t5-oR3gHGHr8KiyCd_Twers4PcIUM1noqQ,2262
+nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
@@ -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=LGn3rdbvkFfUo9kwMzK4eMGlPAqD9uFdowDx6VdfozQ,3761
+nshtrainer/callbacks/early_stopping.py,sha256=m-YIWwQAmp1KUNYwJEiuQyNRyX-0pmzBbAQrkH1BLYI,5664
 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
@@ -54,7 +54,7 @@ nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy0
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
 nshtrainer/model/__init__.py,sha256=RlGW5a46DZcqK6cYICYxDaKpZIEj-8zLxoMrl432tno,1429
 nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
-nshtrainer/model/config.py,sha256=3qSaXNjpKnxF60LbfpONsVhjtrteLEdHkjb5KAhAnIk,52757
+nshtrainer/model/config.py,sha256=gvZRv8Gow6ViF0t8OIoR_IzxLn8GzVmbpvoVzpdykFc,50871
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -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.0.dist-info/METADATA,sha256=YKxr6_dhF7jHF1njUUJi0CZRR0BuOPDNrQiwiIO7iTU,860
-nshtrainer-0.12.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.12.0.dist-info/RECORD,,
+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,,