From a1429fbfe4fade69a8cff6e27606f219cd206a58 Mon Sep 17 00:00:00 2001 From: nilshof01 Date: Thu, 30 Jul 2026 13:19:03 +0200 Subject: [PATCH] unet_segmentation: raise the epoch budget and add early stopping The pipeline trained a fixed 5 epochs with no early stopping and no best-epoch restore (the DINOv3 pipeline has both). Measured at 32 training chips (6 regions x 5 folds), the best epoch fell in the 19-30 range in every region; the best val F1 reachable within 5 epochs vs 30 epochs + early stopping: Banepa 48.2->69.1, Jakarta 54.8->61.6, Accra 37.1->48.3, Bogota 15.0->23.1, Manila 61.9->67.2, Nairobi 40.3->51.7. epochs default 5 -> 30; new early_stop_patience (default 5, 0 disables) monitoring validation loss with best-epoch weight restore. Compute cost stays bounded by early stopping, and warm starts that converge early (e.g. the building-pretrained base) stop well before the cap. Co-Authored-By: Claude Fable 5 --- models/unet_segmentation/pipeline.py | 20 ++++++++++++++++++++ models/unet_segmentation/stac-item.json | 15 ++++++++++++--- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/models/unet_segmentation/pipeline.py b/models/unet_segmentation/pipeline.py index c1baab2..2e526c2 100644 --- a/models/unet_segmentation/pipeline.py +++ b/models/unet_segmentation/pipeline.py @@ -276,6 +276,7 @@ def train_model( max_grad_norm = hyperparameters.get("max_grad_norm", 1.0) scheduler_name = hyperparameters.get("scheduler", "cosine") freeze_encoder = hyperparameters.get("freeze_encoder", True) + early_stop_patience = int(hyperparameters.get("early_stop_patience", 5)) seed = split_info["seed"] with mlflow_training_context(hyperparameters, model_name, base_model_id, dataset_id): @@ -327,6 +328,10 @@ def train_model( train_losses: list[float] = [] val_losses: list[float] = [] + best_val_loss = float("inf") + best_epoch = 0 + best_state: dict[str, Any] | None = None + epochs_without_improvement = 0 model.train() for epoch in range(epochs): @@ -358,6 +363,21 @@ def train_model( msg = f"epoch {epoch + 1}/{epochs} train_loss={avg_train_loss:.4f} val_loss={avg_val_loss:.4f}" print(msg, flush=True) + if avg_val_loss < best_val_loss: + best_val_loss = avg_val_loss + best_epoch = epoch + 1 + best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} + epochs_without_improvement = 0 + else: + epochs_without_improvement += 1 + if early_stop_patience > 0 and epochs_without_improvement >= early_stop_patience: + print(f"early stopping at epoch {epoch + 1} (best epoch {best_epoch})", flush=True) + break + + if best_state is not None: + model.load_state_dict(best_state) + log_metadata(metadata={"best_epoch": best_epoch, "best_val_loss": best_val_loss}) + from fair.zenml.metrics import log_loss_history log_loss_history(train_losses, val_losses) diff --git a/models/unet_segmentation/stac-item.json b/models/unet_segmentation/stac-item.json index 5d34dee..f78eadf 100644 --- a/models/unet_segmentation/stac-item.json +++ b/models/unet_segmentation/stac-item.json @@ -172,7 +172,8 @@ } ], "mlm:hyperparameters": { - "training.epochs": 5, + "training.epochs": 30, + "training.early_stop_patience": 5, "training.batch_size": 4, "training.learning_rate": 0.0001, "training.weight_decay": 0.0001, @@ -197,10 +198,18 @@ { "key": "epochs", "type": "int", - "default": 5, + "default": 30, "min": 1, "max": 500, - "description": "Number of training epochs" + "description": "Maximum number of training epochs; early stopping usually ends training sooner" + }, + { + "key": "early_stop_patience", + "type": "int", + "default": 5, + "min": 0, + "max": 100, + "description": "Stop training after this many epochs without validation-loss improvement and restore the best epoch's weights; 0 disables early stopping" }, { "key": "batch_size",