Patch-Based Cross-Validation and Final Model Training

Cross-validation gives a more reliable estimate of how well a model generalizes than a single

Google Colab
from fastMONAI.vision_all import *

from monai.apps import DecathlonDataset
from sklearn.model_selection import KFold
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.losses import DiceCELoss
import torchio as tio
import gc

Configuration

We use reduced settings so the notebook runs quickly as a demo. Production runs use many more epochs. The reference script research/vs_seg/patch_based_dev/train_cv.py trains each fold for 500 epochs.

EPOCHS = 20
BS = 4
LR = 1e-3
EXPERIMENT = "Task02_Heart_5Fold_CV"

Download the dataset

We use MONAI’s DecathlonDataset to download the Heart MRI dataset from the Medical Segmentation Decathlon. Cross-validation folds over all the cases we load, so we keep them in one DataFrame.

path = Path('../data')
path.mkdir(exist_ok=True)
task = "Task02_Heart"
training_data = DecathlonDataset(root_dir=path, task=task, section="training",
    download=True, cache_num=0, num_workers=3)
df = pd.DataFrame(training_data.data)
df.shape

Analyze the dataset

MedDataset analyzes the masks and recommends a target voxel spacing for resampling. We also keep med_dataset.fingerprint (a content hash of the dataset) to tag every MLflow run for reproducibility.

med_dataset = MedDataset(img_list=df.label.tolist(), dtype=MedMask, max_workers=12)
data_info_df = med_dataset.summary()
suggestion = med_dataset.get_suggestion()
target_spacing = suggestion['target_spacing']
target_spacing
stats = med_dataset.get_size_statistics(target_spacing=target_spacing)
print(f"Image sizes (after resampling to {target_spacing}):")
print(f"  Min:    {stats['min']}")
print(f"  Median: {stats['median']}")
print(f"  Max:    {stats['max']}")

suggested_size = suggest_patch_size(med_dataset, target_spacing=target_spacing)
print(f"\nSuggested patch size: {suggested_size}")

Configure patch-based training

PatchConfig centralizes every patch-related parameter so that training and inference stay consistent. We use the same configuration for all five folds and for the final model, so the cross-validation estimate reflects the deployed model.

  • patch_size: size of extracted patches [x, y, z] (divisible by 16 for UNet)
  • samples_per_volume: patches extracted per volume per epoch
  • sampler_type / label_probabilities: 'label' sampling foreground-weights patches
  • patch_overlap / aggregation_mode: sliding-window inference settings ('hann' = smooth)
  • target_spacing: resampling grid (from the dataset analysis above)
  • normalization: set once here and applied at both training and inference
patch_config = PatchConfig(
    patch_size=[160, 160, 80],
    samples_per_volume=8,
    sampler_type='label',
    label_probabilities={0: 0.5, 1: 0.5},
    patch_overlap=0.5,
    keep_largest_component=True,
    target_spacing=target_spacing,
    aggregation_mode='hann',
    # normalization is set once on the config and applied at both training and inference
    # (no need to re-specify it at inference time).
    normalization=[ZNormalization(masking_method='foreground')]
)

print(f"Patch config: {patch_config}")
model_spec = make_model_spec('monai.unet', {
    'spatial_dims': 3, 'in_channels': 1, 'out_channels': 2,
    'channels': [16, 32, 64, 128, 256], 'strides': [2, 2, 2, 2],
    'num_res_units': 2, 'norm': 'INSTANCE',
})
output_spec = make_output_spec('multiclass_segmentation', classes=2)

Define augmentations

Normalization is already set on the PatchConfig (applied at both training and inference). For the remaining augmentations we use gpu_patch_augmentations, which runs batched on the GPU (roughly 25x faster than the per-sample CPU path) and is applied to each training batch. It is mutually exclusive with patch_tfms. We reuse the same augmentation across all folds and the final model.

gpu_aug = gpu_patch_augmentations(patch_config.patch_size, patch_config.target_spacing)

Part 1: Cross-validation (estimate)

We now estimate generalization performance with 5-fold cross-validation:

  1. Split the subjects into 5 folds.
  2. For each fold, train a fresh model on the other 4 folds and evaluate on the held-out fold.
  3. Aggregate the per-case metrics into a mean +/- std summary across all held-out cases.

Each fold gets its own MLflow run, and every model is trained from scratch so no information leaks between folds.

Create cross-validation folds

KFold assigns each subject to exactly one of 5 folds. We store the fold number in a fold column (1 to 5); a subject’s fold is its validation fold (it trains in the other four).

kf = KFold(n_splits=5, shuffle=True, random_state=42)
df = df.reset_index(drop=True)
df['fold'] = -1

for fold_num, (_, val_idx) in enumerate(kf.split(df), start=1):
    df.loc[val_idx, 'fold'] = fold_num

df['fold'].value_counts().sort_index()

Choosing a fold strategy. Plain KFold is fine when subjects are independent and roughly homogeneous. Real datasets often need:

  • StratifiedKFold to balance a property across folds, for example binning tumor volume into quartiles so each fold has a similar tumor-size distribution (as the VS dataset does).
  • GroupKFold to keep all scans from the same patient in the same fold, preventing leakage when a subject contributes multiple volumes.

Define the per-fold train-and-evaluate function

from fastMONAI.vision_metrics import (calculate_dsc, calculate_surface_metrics,
                                       calculate_confusion_metrics,
                                       calculate_lesion_detection_rate,
                                       calculate_signed_rve)
def train_one_fold(fold_num, df, patch_config, gpu_aug, fingerprint,
                   epochs=EPOCHS, bs=BS, lr=LR, experiment=EXPERIMENT):
    # Split: the held-out fold is validation, the other four folds are training.
    fold_df = df.copy()
    fold_df['is_val'] = fold_df['fold'] == fold_num

    dls = MedPatchDataLoaders.from_df(
        df=fold_df, img_col='image', mask_col='label', valid_col='is_val',
        patch_config=patch_config, gpu_augmentation=gpu_aug, bs=bs)

    # A fresh model + loss + Learner every fold (no weight leakage between folds).
    model = UNet(spatial_dims=3, in_channels=1, out_channels=2,
                 channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2),
                 num_res_units=2, norm=Norm.INSTANCE)
    loss_func = CustomLoss(loss_func=DiceCELoss(
        to_onehot_y=True, softmax=True, include_background=False, batch=True))
    learn = Learner(dls, model, loss_func=loss_func, metrics=[AccumulatedDice(n_classes=2)])

    best_fname = f'best_fold_{fold_num}'
    save_best = EMACheckpoint(monitor='accumulated_dice', momentum=0.9,
                              comp=np.greater, fname=best_fname, with_opt=False)
    mlflow_cb = create_mlflow_callback(
        learn, experiment_name=experiment, run_name=f'fold_{fold_num}',
        extra_tags={'fold': str(fold_num)}, dataset_version=fingerprint,
        model_spec=model_spec, output_spec=output_spec)

    learn.fit_one_cycle(epochs, lr, cbs=[mlflow_cb, save_best])
    learn.load(best_fname)

    # Full-volume sliding-window inference on the held-out fold (the real evaluation).
    val_df = fold_df[fold_df['fold'] == fold_num].reset_index(drop=True)
    val_img_paths = val_df['image'].tolist()
    val_mask_paths = val_df['label'].tolist()

    predictions = patch_inference(learner=learn, config=patch_config,
                                  file_paths=val_img_paths,
                                  save_dir=f'predictions/cv_fold_{fold_num}',
                                  progress=True, tta=True)

    results = []
    for i, pred in enumerate(predictions):
        # patch_inference returns predictions in native voxel space, and MedMask.create loads
        # the GT natively too, so both share the original grid. Surface metrics need the
        # per-case voxel spacing read from the GT file (the cohort spacing is non-uniform).
        gt = MedMask.create(val_mask_paths[i])
        pred_5d = pred.unsqueeze(0).float()
        gt_5d = gt.data.unsqueeze(0).float()
        spacing_mm = tio.LabelMap(val_mask_paths[i]).spacing
        sm = calculate_surface_metrics(pred_5d, gt_5d, spacing_mm=spacing_mm)
        results.append({
            'fold': fold_num,
            'image': Path(val_img_paths[i]).name,
            'dsc': calculate_dsc(pred_5d, gt_5d).mean().item(),
            'sensitivity': calculate_confusion_metrics(pred_5d, gt_5d, 'sensitivity').nanmean().item(),
            'precision': calculate_confusion_metrics(pred_5d, gt_5d, 'precision').nanmean().item(),
            'ldr': calculate_lesion_detection_rate(pred_5d, gt_5d).nanmean().item(),
            'rve': calculate_signed_rve(pred_5d, gt_5d).nanmean().item(),
            'hd95_mm': sm['hd95_mm'],
            'assd_mm': sm['assd_mm'],
        })

    results_df = pd.DataFrame(results)
    mlflow_cb.log_metrics_table(results_df.drop(columns=['fold']), display=False)
    mlflow_cb.log_dataframe(results_df)
    print(f"[Fold {fold_num}] DSC: {results_df['dsc'].mean():.4f} +/- {results_df['dsc'].std():.4f}")

    del learn, model, dls
    torch.cuda.empty_cache()
    gc.collect()
    return results_df

Run the 5-fold loop

This trains and evaluates all five folds sequentially. Each call returns a per-fold results DataFrame; we collect them for aggregation. (On a single GPU this is the slow part: five full trainings.)

This notebook uses 16 cases. MONAI’s DecathlonDataset(section="training") reserves 20% of Task02_Heart’s 20 labeled cases as a separate validation section we do not use, so KFold runs over the remaining 16, giving just 3-4 per fold. Per-fold scores and their spread will be noisy here, so treat the aggregate as illustrative rather than a precise benchmark. Real datasets give many more cases per fold.

all_results = [
    train_one_fold(f, df, patch_config, gpu_aug, med_dataset.fingerprint)
    for f in range(1, 6)
]

Aggregate cross-validation results

The cross-validation estimate is the mean +/- std across all held-out cases (pooled over the five folds) for each metric; the per-fold DSC block below is where you see fold-to-fold variation. We exclude non-finite hd95_mm / assd_mm values (which occur when exactly one of the prediction or ground truth is empty for a case) from those means, mirroring aggregate_results in train_cv.py.

cv_summary = pd.concat(all_results, ignore_index=True)
cv_summary.to_csv('cv_summary.csv', index=False)

metrics = ['dsc', 'sensitivity', 'precision', 'ldr', 'rve', 'hd95_mm', 'assd_mm']
inf_excluded = {'hd95_mm', 'assd_mm'}  # empty-case inf must not poison the mean

print("=" * 46)
print("  CROSS-VALIDATION SUMMARY")
print("=" * 46)
print(f"  Folds completed: {sorted(int(x) for x in cv_summary['fold'].unique())}")
print(f"  Total subjects:  {len(cv_summary)}\n")
print(f"  {'Metric':<14}{'Mean':>10}{'Std':>10}")
print(f"  {'-' * 34}")
for m in metrics:
    col = cv_summary[m]
    note = ""
    if m in inf_excluded:
        col = col[np.isfinite(col)]
        n_skip = len(cv_summary) - len(col)
        if n_skip:
            note = f"  ({n_skip} non-finite excl.)"
    mean = col.mean() if len(col) else float('nan')
    std = col.std() if len(col) else float('nan')
    print(f"  {m:<14}{mean:>10.4f}{std:>10.4f}{note}")

print("\n  Per-fold DSC:")
for fold_num, group in cv_summary.groupby('fold'):
    print(f"    Fold {fold_num}: {group['dsc'].mean():.4f} +/- {group['dsc'].std():.4f}")

print("\n  Saved to cv_summary.csv")
cv_summary

Part 2: Final model on all data (deploy)

Cross-validation in Part 1 gave us a performance estimate.

There are two reasonable paths from here. One is to keep the five cross-validation models and deploy them as an ensemble, averaging their predicted probability maps. This is nnU-Net’s default (Isensee et al., 2021): averaging across the folds reduces prediction variance and generally improves accuracy, at the cost of running five networks at inference.

The other path, which we take here, is to train a single model on all the available data so it sees every example. The training pipeline still needs a non-empty validation set to build its DataLoaders and to watch the loss and pseudo-Dice curves, but we have deliberately held nothing out. We resolve this by not selecting the model on validation performance. The cross-validation runs already tell us roughly how long training takes to plateau, so we fix the number of epochs from Part 1, train on all data for that fixed budget, and keep the final-epoch weights.

Build the all-data DataLoaders

Every subject is tagged is_val=False (training). We then append a duplicated random subset tagged is_val=True for monitoring. MedPatchDataLoaders.from_df splits on is_val, wiring reorder / resample / normalization from the same PatchConfig used in Part 1, so preprocessing stays identical.

NOMINAL_VAL_PCT = 0.15

# All subjects go into training; a small random subset is duplicated into a nominal
# validation split because the DataLoaders need a non-empty validation set (its scores
# are optimistic, since those subjects are also in training).
n_val = max(2, int(len(df) * NOMINAL_VAL_PCT))
nominal_val_df = df.sample(n=n_val).copy()

all_train_df = df.copy()
all_train_df['is_val'] = False
nominal_val_df['is_val'] = True
final_df = pd.concat([all_train_df, nominal_val_df], ignore_index=True)

dls = MedPatchDataLoaders.from_df(
    df=final_df, img_col='image', mask_col='label', valid_col='is_val',
    patch_config=patch_config, gpu_augmentation=gpu_aug, bs=BS)

print(f"Training subjects (all data):                  {len(dls.train.subjects_dataset)}")
print(f"Nominal validation subjects (also in training): {len(dls.valid.subjects_dataset)}")

Train the final model

A fresh UNet, loss, and Learner (identical architecture and configuration to the fold models), now trained on all data. The MLflow run is tagged training_type=final to distinguish it from the fold runs.

model = UNet(spatial_dims=3, in_channels=1, out_channels=2,
             channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2),
             num_res_units=2, norm=Norm.INSTANCE)

loss_func = CustomLoss(loss_func=DiceCELoss(
    to_onehot_y=True, softmax=True, include_background=False, batch=True))
learn = Learner(dls, model, loss_func=loss_func, metrics=[AccumulatedDice(n_classes=2)])

mlflow_cb = create_mlflow_callback(
    learn, experiment_name=EXPERIMENT, run_name='final_all_data',
    extra_tags={'training_type': 'final'}, dataset_version=med_dataset.fingerprint,
    model_spec=model_spec, output_spec=output_spec)
learn.fit_one_cycle(EPOCHS, LR, cbs=[mlflow_cb])

Export for deployment

Save the all-data model as Safetensors with its trusted architecture and complete patch-inference contract. Keep the separate .pth artifact for rebuilding a Learner and continuing training with a fresh optimizer; do not export a new Learner pickle.

Path('models').mkdir(exist_ok=True)
inference_config = patch_inference_config(patch_config, output_spec)
save_safetensors_model(
    learn.model, 'models/final_model.safetensors', model_spec,
    inference_config, artifact_role='final')
print("Saved models/final_model.safetensors with its embedded inference configuration")

View experiment tracking

All five fold runs (tagged fold=1..5) and the final run (tagged training_type=final) are logged under the Task02_Heart_5Fold_CV experiment.

mlflow_ui = MLflowUIManager()
mlflow_ui.start_ui()

Summary

This tutorial taught the full train-and-ship methodology for patch-based 3D segmentation:

Part 1, cross-validation (estimate)

  1. Fold generation: KFold assigns each subject to one validation fold.
  2. Per-fold loop: a fresh Learner per fold, with no weight leakage.
  3. Per-fold metrics: full-volume inference and metrics in one MLflow run per fold.
  4. Aggregate: mean and spread across held-out cases estimate generalization.

Part 2, final model (deploy)

  1. Train on all data: one deployable model is trained on every subject.
  2. Export: final_model.safetensors stores tensors plus the explicit model, preprocessing, sliding-window, and output contract. No separate training JSON or Learner .pkl is created.
mlflow_ui.stop()

References

  • Isensee, F., et al. (2021). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods. https://doi.org/10.1038/s41592-020-01008-z