From 3bde73740ca267871b5e220ac335438d9bc2020b Mon Sep 17 00:00:00 2001 From: Briarion Date: Mon, 16 Mar 2026 08:46:27 +0300 Subject: [PATCH] feat: add early stopping support to PyTorchModelTrainer Add optional early stopping to prevent overfitting in PyTorch-based FreqAI models. When `early_stopping_patience` is set in model_training_parameters, training will stop if validation loss does not improve for the specified number of epochs. Changes: - Add `early_stopping_patience` parameter (default 0 = disabled) - `estimate_loss()` now returns average loss (float | None) instead of None, enabling downstream use for schedulers and early stopping - Track best validation loss and patience counter across epochs Usage in config: ```json { "model_training_parameters": { "n_epochs": 100, "early_stopping_patience": 10 } } ``` The change is fully backward compatible - early stopping is disabled by default, and the return value of estimate_loss() can be safely ignored by existing subclasses. Co-Authored-By: Claude Opus 4.6 (1M context) --- freqtrade/freqai/torch/PyTorchModelTrainer.py | 40 ++++++++++++++++++- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/freqtrade/freqai/torch/PyTorchModelTrainer.py b/freqtrade/freqai/torch/PyTorchModelTrainer.py index 1c12f15fd..223f786e4 100644 --- a/freqtrade/freqai/torch/PyTorchModelTrainer.py +++ b/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]: