nshtrainer 1.3.3__tar.gz → 1.3.4__tar.gz

{nshtrainer-1.3.3 → nshtrainer-1.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.3.3
+Version: 1.3.4
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.3.3 → nshtrainer-1.3.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.3.3"
+version = "1.3.4"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"

{nshtrainer-1.3.3 → nshtrainer-1.3.4}/src/nshtrainer/callbacks/early_stopping.py

@@ -43,6 +43,13 @@ class EarlyStoppingCallbackConfig(CallbackConfigBase):
     the training will be stopped.
     """
 
+    skip_first_n_epochs: int = 0
+    """
+    Number of initial epochs to skip before starting to monitor for early stopping.
+    This helps avoid false early stopping when the model might temporarily perform worse
+    during early training phases.
+    """
+
     strict: bool = True
     """
     Whether to enforce that the monitored quantity must improve by at least `min_delta`
@@ -94,6 +101,16 @@ class EarlyStoppingCallback(_EarlyStopping):
         if getattr(trainer, "fast_dev_run", False):
             return
 
+        # Skip early stopping check for the first n epochs
+        if trainer.current_epoch < self.config.skip_first_n_epochs:
+            if self.verbose and trainer.current_epoch == 0:
+                self._log_info(
+                    trainer,
+                    f"Early stopping checks are disabled for the first {self.config.skip_first_n_epochs} epochs",
+                    self.log_rank_zero_only,
+                )
+            return
+
         should_stop, reason = False, None
 
         if not should_stop: