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) <noreply@anthropic.com>
This commit is contained in:
@@ -63,6 +63,11 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
|
|||||||
self.tb_logger = tb_logger
|
self.tb_logger = tb_logger
|
||||||
self.test_batch_counter = 0
|
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]):
|
def fit(self, data_dictionary: dict[str, pd.DataFrame], splits: list[str]):
|
||||||
"""
|
"""
|
||||||
:param data_dictionary: the dictionary constructed by DataHandler to hold
|
:param data_dictionary: the dictionary constructed by DataHandler to hold
|
||||||
@@ -99,15 +104,40 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
|
|||||||
|
|
||||||
# evaluation
|
# evaluation
|
||||||
if "test" in splits:
|
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()
|
@torch.no_grad()
|
||||||
def estimate_loss(
|
def estimate_loss(
|
||||||
self,
|
self,
|
||||||
data_loader_dictionary: dict[str, DataLoader],
|
data_loader_dictionary: dict[str, DataLoader],
|
||||||
split: str,
|
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()
|
self.model.eval()
|
||||||
|
total_loss = 0.0
|
||||||
|
num_batches = 0
|
||||||
|
|
||||||
for _, batch_data in enumerate(data_loader_dictionary[split]):
|
for _, batch_data in enumerate(data_loader_dictionary[split]):
|
||||||
xb, yb = batch_data
|
xb, yb = batch_data
|
||||||
xb = xb.to(self.device)
|
xb = xb.to(self.device)
|
||||||
@@ -115,11 +145,17 @@ class PyTorchModelTrainer(PyTorchTrainerInterface):
|
|||||||
|
|
||||||
yb_pred = self.model(xb)
|
yb_pred = self.model(xb)
|
||||||
loss = self.criterion(yb_pred, yb)
|
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.tb_logger.log_scalar(f"{split}_loss", loss.item(), self.test_batch_counter)
|
||||||
self.test_batch_counter += 1
|
self.test_batch_counter += 1
|
||||||
|
|
||||||
self.model.train()
|
self.model.train()
|
||||||
|
|
||||||
|
if num_batches > 0:
|
||||||
|
return total_loss / num_batches
|
||||||
|
return None
|
||||||
|
|
||||||
def create_data_loaders_dictionary(
|
def create_data_loaders_dictionary(
|
||||||
self, data_dictionary: dict[str, pd.DataFrame], splits: list[str]
|
self, data_dictionary: dict[str, pd.DataFrame], splits: list[str]
|
||||||
) -> dict[str, DataLoader]:
|
) -> dict[str, DataLoader]:
|
||||||
|
|||||||
Reference in New Issue
Block a user