Training Loop

Training and validation loop with mixed precision, gradient accumulation, and MLflow integration.

training.train

prepare_output_for_comparison(outputs: torch.Tensor, target_size: tuple[int, int], output_size: int | None = None) -> torch.Tensor

Prepare model outputs for comparison with target mask.

When using context window (model output larger than output_size), center-crops to output_size first, then upsamples to target_size.

Parameters:

Name	Type	Description	Default
outputs	Tensor	Model predictions with shape (B, C, H, W)	required
target_size	tuple[int, int]	Target spatial size (height, width) to match mask	required
output_size	int | None	Expected output spatial size for center-cropping. If provided and output is larger, center-crops to this size. Use sentinel_patch_size when using context window.	None

Returns:

Type	Description
Tensor	Tensor with shape (B, C, target_size[0], target_size[1])

Source code in src/training/train.py

def prepare_output_for_comparison(
    outputs: torch.Tensor,
    target_size: tuple[int, int],
    output_size: int | None = None,
) -> torch.Tensor:
    """Prepare model outputs for comparison with target mask.

    When using context window (model output larger than output_size), center-crops
    to output_size first, then upsamples to target_size.

    Args:
        outputs: Model predictions with shape (B, C, H, W)
        target_size: Target spatial size (height, width) to match mask
        output_size: Expected output spatial size for center-cropping.
            If provided and output is larger, center-crops to this size.
            Use sentinel_patch_size when using context window.

    Returns:
        Tensor with shape (B, C, target_size[0], target_size[1])

    """
    if outputs.shape[-2:] == target_size:
        return outputs

    out_h, out_w = outputs.shape[-2:]

    # Center-crop if using context window
    if output_size is not None and out_h > output_size:
        crop_margin_h = (out_h - output_size) // 2
        crop_margin_w = (out_w - output_size) // 2
        outputs = outputs[
            :,
            :,
            crop_margin_h : crop_margin_h + output_size,
            crop_margin_w : crop_margin_w + output_size,
        ]

    # Upsample to target size
    if outputs.shape[-2:] != target_size:
        outputs = F.interpolate(
            outputs,
            size=target_size,
            mode="bilinear",
            align_corners=False,
        )

    return outputs

train(model: torch.nn.Module, train_loader: torch.utils.data.DataLoader, val_loader: torch.utils.data.DataLoader, criterion: torch.nn.Module, optimizer: torch.optim.Optimizer, device: torch.device, scheduler: LRScheduler | None = None, epochs: int = 100, patience: int = 20, num_classes: int = 13, other_class_index: int | None = None, accumulation_steps: int = 1, early_stopping_criterion: str = 'loss', *, use_amp: bool = False, apply_augmentations: bool = True, data_config: dict[str, Any] | None = None, log_evaluation_metrics: bool = True, log_model: bool = True, pruning_callback: Any | None = None, output_size: int | None = None, gradient_clip_val: float | None = None, sentinel_augmenter: Any | None = None) -> dict[str, list[float] | float]

Train a segmentation model, monitoring validation loss and saving the best model.

Detailed metrics should be calculated separately after training using an evaluation function.

Parameters:

Name	Type	Description	Default
model	Module	PyTorch model.	required
train_loader	DataLoader	DataLoader for training data.	required
val_loader	DataLoader	DataLoader for validation data.	required
criterion	Module	Loss function.	required
optimizer	Optimizer	Optimizer for training.	required
device	device	Device (CPU/GPU).	required
scheduler	LRScheduler | None	Optional learning rate scheduler.	None
apply_augmentations	bool	Whether to apply augmentations to the training data. Defaults to True.	True
data_config	dict[str, Any] | None	Full data configuration dict (contains augmentation config, normalization settings, and channel selections).	None
epochs	int	Maximum number of epochs to train. Defaults to 100.	100
patience	int	Early stopping patience. Defaults to 20.	20
num_classes	int	Number of classes in the segmentation task. Defaults to 13.	13
accumulation_steps	int	Number of steps to accumulate gradients before updating. Defaults to 1.	1
use_amp	bool	Whether to use Automatic Mixed Precision (AMP). Defaults to False.	False
log_evaluation_metrics	bool	Whether to log metrics and models to MLflow. Defaults to True.	True
log_model	bool	Whether to log the best model to MLflow. Defaults to True.	True

Returns:

Name	Type	Description
dict	dict[str, list[float] | float]	History of training and validation losses/mIoUs, and best values. The best model is logged as an MLflow artifact 'best_model' if log_to_mlflow is True.

Source code in src/training/train.py

def train(
    model: torch.nn.Module,
    train_loader: torch.utils.data.DataLoader,
    val_loader: torch.utils.data.DataLoader,
    criterion: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    scheduler: LRScheduler | None = None,
    epochs: int = 100,
    patience: int = 20,
    num_classes: int = 13,
    other_class_index: int | None = None,
    accumulation_steps: int = 1,
    early_stopping_criterion: str = "loss",
    *,
    use_amp: bool = False,
    apply_augmentations: bool = True,
    data_config: dict[str, Any] | None = None,
    log_evaluation_metrics: bool = True,
    log_model: bool = True,
    pruning_callback: Any | None = None,
    output_size: int | None = None,
    gradient_clip_val: float | None = None,
    sentinel_augmenter: Any | None = None,
) -> dict[str, list[float] | float]:
    """Train a segmentation model, monitoring validation loss and saving the best model.

    Detailed metrics should be calculated separately after training using an
    evaluation function.

    Args:
        model: PyTorch model.
        train_loader: DataLoader for training data.
        val_loader: DataLoader for validation data.
        criterion: Loss function.
        optimizer: Optimizer for training.
        device: Device (CPU/GPU).
        scheduler: Optional learning rate scheduler.
        apply_augmentations: Whether to apply augmentations to the training data.
            Defaults to True.
        data_config: Full data configuration dict (contains augmentation config,
            normalization settings, and channel selections).
        epochs: Maximum number of epochs to train. Defaults to 100.
        patience: Early stopping patience. Defaults to 20.
        num_classes: Number of classes in the segmentation task. Defaults to 13.
        accumulation_steps: Number of steps to accumulate gradients before updating.
            Defaults to 1.
        use_amp: Whether to use Automatic Mixed Precision (AMP). Defaults to False.
        log_evaluation_metrics: Whether to log metrics and models to MLflow.
            Defaults to True.
        log_model: Whether to log the best model to MLflow. Defaults to True.

    Returns:
        dict: History of training and validation losses/mIoUs, and best values.
            The best model is logged as an MLflow artifact 'best_model' if
            log_to_mlflow is True.

    """
    if early_stopping_criterion not in ("loss", "miou"):
        msg = f"early_stopping_criterion must be 'loss' or 'miou', got {early_stopping_criterion!r}"
        raise ValueError(msg)
    model.to(device)

    sample_inputs, sample_batch_positions = _get_sample_batch(train_loader)
    sample_input_shape = tuple(int(x) for x in sample_inputs.shape)

    if log_evaluation_metrics:
        _log_model_description(
            model,
            device,
            sample_input_shape,
            sample_inputs=sample_inputs,
            batch_positions=sample_batch_positions,
        )

    best_val_loss = float("inf")
    best_val_miou = 0.0
    no_improve = 0
    losses_train: list[float] = []
    losses_val: list[float] = []
    mious_val: list[float] = []

    augmenter = FlairAugmentation(data_config) if apply_augmentations and data_config else None
    chessmix = None
    if apply_augmentations and data_config:
        aug_config = data_config.get("data_augmentation", {}).get("augmentations", {})
        if "chessmix" in aug_config:
            cm_cfg = aug_config["chessmix"]
            chessmix = ChessMix(
                prob=cm_cfg.get("prob", 0.5),
                grid_sizes=cm_cfg.get("grid_sizes", [4]),
                ignore_index=cm_cfg.get("ignore_index", 12),
                class_counts=cm_cfg.get("class_counts", None),
                num_classes=data_config.get("num_classes", 13),
            )
    best_model_state = None

    is_temporal_model = sample_inputs.ndim == TEMPORAL_MODEL_NDIM

    if is_temporal_model:
        logger.info("Detected temporal model (5D input). Using temporal training loop.")
        train_epoch_fn = _train_epoch_temporal
        validate_epoch_fn = _validate_epoch_temporal
    else:
        logger.info("Detected standard model. Using standard training loop.")
        train_epoch_fn = _train_epoch_standard
        validate_epoch_fn = _validate_epoch_standard

    is_step_scheduler = scheduler is not None and not isinstance(scheduler, ReduceLROnPlateau)
    step_scheduler = scheduler if is_step_scheduler else None

    for epoch in range(epochs):
        if is_temporal_model:
            loss_epoch = train_epoch_fn(
                model,
                train_loader,
                criterion,
                optimizer,
                device,
                accumulation_steps,
                use_amp,
                step_scheduler,
                output_size,
                gradient_clip_val,
                sentinel_augmenter,
            )
        else:
            loss_epoch = train_epoch_fn(
                model,
                train_loader,
                criterion,
                optimizer,
                device,
                augmenter,
                chessmix,
                accumulation_steps,
                use_amp,
                step_scheduler,
            )
        losses_train.append(loss_epoch)
        logger.info("Epoch %d/%d: Training Loss: %.4f", epoch + 1, epochs, loss_epoch)

        val_loss, val_miou = validate_epoch_fn(
            model,
            val_loader,
            criterion,
            device,
            num_classes,
            other_class_index,
        )
        losses_val.append(val_loss)
        mious_val.append(val_miou)
        logger.info(
            "Epoch %d/%d: Validation Loss: %.4f, Validation mIoU: %.4f",
            epoch + 1,
            epochs,
            val_loss,
            val_miou,
        )

        if log_evaluation_metrics:
            log_metrics_to_mlflow(
                metrics={"train_loss": loss_epoch, "val_loss": val_loss, "val_miou": val_miou},
                step=epoch,
            )

        if pruning_callback is not None:
            report_value = val_miou if early_stopping_criterion == "miou" else val_loss
            pruning_callback(report_value, epoch)

        if early_stopping_criterion == "miou":
            improved = val_miou > best_val_miou
        else:
            improved = val_loss < best_val_loss

        if improved:
            best_val_miou = val_miou
            best_val_loss = val_loss
            no_improve = 0
            if early_stopping_criterion == "miou":
                logger.info("Validation mIoU improved to %.4f", best_val_miou)
            else:
                logger.info("Validation loss improved to %.4f", best_val_loss)
            best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        else:
            no_improve += 1
            logger.info("No improvement for %d epochs.", no_improve)
            if no_improve >= patience:
                logger.info("Early stopping at epoch %d.", epoch + 1)
                break

        if scheduler is not None and isinstance(scheduler, ReduceLROnPlateau):
            scheduler.step(val_loss)

        if log_evaluation_metrics:
            current_lr = optimizer.param_groups[0]["lr"]
            mlflow.log_metric("learning_rate", current_lr, step=epoch)
            logger.info("Current learning rate: %.6f", current_lr)

    if log_model and best_model_state is not None:
        model.load_state_dict(best_model_state)
        model.to(device)
        log_model_to_mlflow(
            model=model,
            train_loader=train_loader,
            sample_input_shape=sample_input_shape,
            num_classes=num_classes,
        )

    return {
        "train_loss": losses_train,
        "val_loss": losses_val,
        "val_miou": mious_val,
        "best_val_loss": best_val_loss,
        "best_val_miou": best_val_miou,
    }