nshtrainer 0.11.13__tar.gz → 0.12.1__tar.gz

{nshtrainer-0.11.13 → nshtrainer-0.12.1}/PKG-INFO

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

{nshtrainer-0.11.13 → nshtrainer-0.12.1}/pyproject.toml

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

{nshtrainer-0.11.13 → nshtrainer-0.12.1}/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.11.13 → nshtrainer-0.12.1}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -55,7 +55,7 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
 
     @override
     def default_filename(self):
-        return f"epoch{{epoch:03d}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
+        return f"epoch{{epoch}}-step{{step}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
 
     @override
     def topk_sort_key(self, metadata: CheckpointMetadata):

{nshtrainer-0.11.13 → nshtrainer-0.12.1}/src/nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -28,7 +28,7 @@ class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
 
     @override
     def default_filename(self):
-        return "epoch{epoch:03d}-step{step:07d}"
+        return "epoch{epoch}-step{step}"
 
     @override
     def topk_sort_key(self, metadata: CheckpointMetadata):

{nshtrainer-0.11.13 → nshtrainer-0.12.1}/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,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-0.11.13 → nshtrainer-0.12.1}/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"
     """