Merge pull request #12952 from Briarion/feat/pytorch-early-stopping

feat: add early stopping support to PyTorchModelTrainer

docs/freqai-parameter-table.md

@@ -106,6 +106,7 @@ Mandatory parameters are marked as **Required** and have to be set in one of the
 | `n_epochs`   | The `n_epochs` parameter is a crucial setting in the PyTorch training loop that determines the number of times the entire training dataset will be used to update the model's parameters. An epoch represents one full pass through the entire training dataset. Overrides `n_steps`. Either `n_epochs` or `n_steps` must be set. <br><br> **Datatype:** int. optional. <br> Default: `10`.
 | `n_steps`    | An alternative way of setting `n_epochs` -  the number of training iterations to run. Iteration here refer to the number of times we call `optimizer.step()`. Ignored if `n_epochs` is set. A simplified version of the function: <br><br> n_epochs = n_steps / (n_obs / batch_size) <br><br> The motivation here is that `n_steps` is easier to optimize and keep stable across different n_obs - the number of data points.  <br> <br> **Datatype:** int. optional. <br> Default: `None`.
 | `batch_size` | The size of the batches to use during training. <br><br> **Datatype:** int. <br> Default: `64`.
+| `early_stopping_patience` | Number of epochs with no improvement in validation loss before training is stopped early. This helps prevent overfitting by halting training when the model stops improving. Set to `0` to disable early stopping. Requires a test/validation split (`test_size > 0`). <br><br> **Datatype:** int. <br> Default: `0` (disabled).
 
 ### Additional parameters
 

freqtrade/freqai/torch/PyTorchModelTrainer.py

@@ -63,6 +63,11 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
         self.tb_logger = tb_logger
         self.test_batch_counter = 0
 
+        # Early stopping parameters
+        self.early_stopping_patience: int = kwargs.get("early_stopping_patience", 0)
+        self.best_val_loss: float = float("inf")
+        self.patience_counter: int = 0
+
     def fit(self, data_dictionary: dict[str, pd.DataFrame], splits: list[str]):
         """
         :param data_dictionary: the dictionary constructed by DataHandler to hold
@@ -99,15 +104,40 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
 
             # evaluation
             if "test" in splits:
-                self.estimate_loss(data_loaders_dictionary, "test")
+                val_loss = self.estimate_loss(data_loaders_dictionary, "test")
+
+                # Early stopping check
+                if self.early_stopping_patience > 0 and val_loss is not None:
+                    if val_loss < self.best_val_loss:
+                        self.best_val_loss = val_loss
+                        self.patience_counter = 0
+                    else:
+                        self.patience_counter += 1
+                        if self.patience_counter >= self.early_stopping_patience:
+                            logger.info(
+                                f"Early stopping triggered after {self.patience_counter} "
+                                f"epochs without improvement. "
+                                f"Best val_loss: {self.best_val_loss:.6f}"
+                            )
+                            break
 
     @torch.no_grad()
     def estimate_loss(
         self,
         data_loader_dictionary: dict[str, DataLoader],
         split: str,
-    ) -> None:
+    ) -> float | None:
+        """
+        Estimate loss on a data split.
+
+        :param data_loader_dictionary: dictionary of data loaders.
+        :param split: split to estimate loss on (e.g. "test").
+        :return: average loss over all batches, or None if no batches.
+        """
         self.model.eval()
+        total_loss = 0.0
+        num_batches = 0
+
         for _, batch_data in enumerate(data_loader_dictionary[split]):
             xb, yb = batch_data
             xb = xb.to(self.device)
@@ -115,11 +145,17 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
 
             yb_pred = self.model(xb)
             loss = self.criterion(yb_pred, yb)
+            total_loss += loss.item()
+            num_batches += 1
             self.tb_logger.log_scalar(f"{split}_loss", loss.item(), self.test_batch_counter)
             self.test_batch_counter += 1
 
         self.model.train()
 
+        if num_batches > 0:
+            return total_loss / num_batches
+        return None
+
     def create_data_loaders_dictionary(
         self, data_dictionary: dict[str, pd.DataFrame], splits: list[str]
     ) -> dict[str, DataLoader]: