nnUNet-MLflow

Adds MLflow experiment tracking to nnUNet via a drop-in trainer. Just add -tr nnUNetTrainerMLflow when using the nnUNetv2_train command.

What it does

A custom nnUNet trainer that extends the base trainer with MLflow logging for all training metrics, system information, and artifacts (checkpoints, plots, validation results, logs, and nnUNet configuration files). nnUNet is installed as normal, and we simply copy in our trainer without changing any training logic.

Quick Start

# 1. Create and activate conda environment
conda create -n nnunet-mlflow python=3.12.8
conda activate nnunet-mlflow

# 2. Install PyTorch first. Find your specific install: https://pytorch.org/get-started/locally/
# Example for CUDA 12.8
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128

# 3. Install this package
git clone https://github.com/TomTonroe/nnUNet-MLflow.git
cd nnUNet-MLflow
pip install -e .

# 4. Install the trainer into nnUNet
nnunet-mlflow-install

# 5. Train with tracking (just add -tr nnUNetTrainerMLflow)
nnUNetv2_train DATASETID CONFIG FOLD -tr nnUNetTrainerMLflow

# 6. Monitor results
mlflow ui
# Open http://localhost:5000

Viewing Results

# Local access only
mlflow ui
# Network access (other machines on local network)
mlflow ui --host 0.0.0.0 --port 5000
# Open http://YOUR-HOST-IP:5000

Data Storage

MLflow stores all tracking data locally in the ./mlruns folder (created automatically). Each experiment gets its own subfolder with metrics, parameters, and artifacts. You can backup or move this folder to preserve your training history.

Backup your data:

nnunet-mlflow-backup -m ./mlruns -b ./my_backups

MLflow can be configured to store data in databases or cloud storage. See the MLflow documentation for setup details.

Examples - Use nnUNetv2_train as normal

# Single fold
nnUNetv2_train 005 3d_lowres 0 -tr nnUNetTrainerMLflow
# All folds
nnUNetv2_train 010 2d all -tr nnUNetTrainerMLflow
# Custom plans
nnUNetv2_train 027 3d_fullres all -tr nnUNetTrainerMLflow -p nnUNetResEncUNetLPlans

Everything else works exactly like standard nnUNet - see the official documentation.

Acknowledgements

Please cite the nnU-Net paper in your research:

Isensee, F., Jaeger, P.F., Kohl, S.A.A., Petersen, J., Maier-Hein, K.H. (2021). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods 18, 203–211. https://doi.org/10.1038/s41592-020-01008-z

License

Licensed under the Apache License, Version 2.0. See LICENSE.txt for details.