[FIX] ResetEarlyStopping state in cross-validation when refit=True (#1529)

Author: marcopeix

neuralforecast/common/_base_model.py

@@ -8,7 +8,7 @@
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 import fsspec
 import numpy as np
@@ -295,6 +295,8 @@ def __init__(
             raise Exception("max_epochs is deprecated, use max_steps instead.")
 
         # Callbacks
+        self._early_stop_kwargs: Optional[dict] = None
+        self._early_stop_cb: Optional[EarlyStopping] = None
         if early_stop_patience_steps > 0:
             valid_monitors = ["ptl/val_loss", "valid_loss", "train_loss"]
             if val_monitor not in valid_monitors:
@@ -304,11 +306,12 @@ def __init__(
                 )
             if "callbacks" not in trainer_kwargs:
                 trainer_kwargs["callbacks"] = []
-            trainer_kwargs["callbacks"].append(
-                EarlyStopping(
-                    monitor=val_monitor, patience=early_stop_patience_steps
-                )
-            )
+            self._early_stop_kwargs = {
+                "monitor": val_monitor,
+                "patience": early_stop_patience_steps,
+            }
+            self._early_stop_cb = EarlyStopping(**self._early_stop_kwargs)
+            trainer_kwargs["callbacks"].append(self._early_stop_cb)
 
         # Add GPU accelerator if available
         if trainer_kwargs.get("accelerator", None) is None:
@@ -604,6 +607,17 @@ def _fit(
         self.trainer_kwargs["val_check_interval"] = int(val_check_interval)
         self.trainer_kwargs["check_val_every_n_epoch"] = None
 
+        # The EarlyStopping callback we add in __init__ accumulates state across fits.
+        # Rebuild it from the stored kwargs.
+        # We match by identity, so user-supplied EarlyStopping callbacks are left untouched.
+        if self._early_stop_cb is not None and self._early_stop_kwargs is not None:
+            callbacks = self.trainer_kwargs.get("callbacks") or []
+            for i, cb in enumerate(callbacks):
+                if cb is self._early_stop_cb:
+                    self._early_stop_cb = EarlyStopping(**self._early_stop_kwargs)
+                    callbacks[i] = self._early_stop_cb
+                    break
+
         if is_local:
             model = self
             trainer = pl.Trainer(**model.trainer_kwargs)

tests/test_core.py

@@ -253,6 +253,64 @@ def test_neural_forecast_refit(setup_airplane_data):
             )
 
 
+# Cross_validation(refit=True, val_size=...) must give each refit window a fresh
+# EarlyStopping state. Otherwise the prior window's wait_count and best_score
+# carry over and subsequent refits stop on the first validation check.
+def test_cross_validation_refit_resets_early_stopping(setup_airplane_data):
+    import pytorch_lightning as pl
+    from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+
+    AirPassengersPanel_train, _ = setup_airplane_data
+
+    class _CaptureEarlyStoppingState(pl.Callback):
+        def __init__(self):
+            self.snapshots = []
+
+        def on_train_start(self, trainer, pl_module):
+            es = next(
+                (cb for cb in trainer.callbacks if isinstance(cb, EarlyStopping)),
+                None,
+            )
+            assert es is not None, "EarlyStopping callback missing from trainer"
+            self.snapshots.append(
+                {
+                    "wait_count": es.wait_count,
+                    "stopped_epoch": es.stopped_epoch,
+                    "best_score": float(es.best_score),
+                }
+            )
+
+    model = NHITS(
+        h=12,
+        input_size=24,
+        max_steps=4,
+        val_check_steps=1,
+        early_stop_patience_steps=1,
+        callbacks=[_CaptureEarlyStoppingState()],
+        enable_progress_bar=False,
+    )
+    nf = NeuralForecast(models=[model], freq="M")
+    capture = next(
+        cb
+        for cb in nf.models[0].trainer_kwargs["callbacks"]
+        if isinstance(cb, _CaptureEarlyStoppingState)
+    )
+
+    nf.cross_validation(
+        df=AirPassengersPanel_train,
+        val_size=12,
+        n_windows=3,
+        refit=True,
+    )
+
+    # One fit per refit window.
+    assert len(capture.snapshots) == 3, capture.snapshots
+    for i, snap in enumerate(capture.snapshots):
+        assert snap["wait_count"] == 0, (i, snap)
+        assert snap["stopped_epoch"] == 0, (i, snap)
+        assert snap["best_score"] == float("inf"), (i, snap)
+
+
 def test_neural_forecast_scaling(setup_airplane_data):
     """Test scaling functionality for NeuralForecast models."""
     AirPassengersPanel_train, AirPassengersPanel_test = setup_airplane_data