Uncertainty-Driven Federated Brain Tumor Segmentation

This repository contains code for training a 3D U-Net for brain tumor segmentation on the FeTS 2022 dataset using federated learning with NVFlare and MONAI.

The repository supports two training modes:

• Standard Federated Learning using Federated Averaging (FedAvg)
• Uncertainty-Driven Sample Selection, where training focuses on the most informative (high-uncertainty) samples while temporarily excluding low-uncertainty cases from training.

---

Overview

Brain tumor segmentation is performed using four MRI modalities:

• T1
• T1ce
• T2
• FLAIR

The network predicts four segmentation classes:

• Background
• Non-enhancing tumor
• Edema
• Enhancing tumor

Each client trains locally on its own dataset and periodically sends model updates to a federated server. The server aggregates the client models using FedAvg before broadcasting the updated global model.

When uncertainty-driven training is enabled, clients dynamically adjust their training dataset by prioritizing uncertain samples and temporarily removing easy cases.

---

Features

• 3D U-Net implementation using MONAI
• Optional UNETR architecture
• Federated learning with NVFlare
• MONAI preprocessing and data augmentation
• DiceLoss or DiceCELoss
• DiceMetric evaluation
• Sliding-window inference
• Automatic class-weight computation
• Dynamic uncertainty-driven sample selection
• Memory bank for tracking case uncertainty over training rounds

---

Repository Structure

.
├── src/
│   ├── client.py                 # Standard federated client
│   ├── client_uncertainty.py     # Federated client with uncertainty-driven sample selection
│   ├── job.py                    # Launches the federated learning simulation
│   ├── model.py                  # U-Net and UNETR model definitions
│   └── utils.py                  # Utility functions, uncertainty estimation, memory bank
├── config.yml                    # Server configuration
├── client_config.yml             # Client configuration
├── environment.yml               # Conda environment
├── submit_job.sh                 # Example SLURM submission script
└── nvflare_sim.sh                # Launch script

---

Model

The default network is a MONAI 3D U-Net with:

• 4 input channels
• 4 output channels
• Residual units
• Batch normalization
• Dropout

The repository also includes support for UNETR, which can be enabled through the configuration files.

---

Data Format

Each subject should contain the following files:

CaseID/
├── CaseID_t1.nii.gz
├── CaseID_t1ce.nii.gz
├── CaseID_t2.nii.gz
├── CaseID_flair.nii.gz
└── CaseID_seg.nii.gz

The segmentation labels are converted into a four-channel one-hot representation during preprocessing.

---

Training Pipeline

Data preprocessing

Training uses MONAI transforms including:

• Image loading
• Intensity normalization
• Multi-channel label conversion
• Spatial padding
• Random 3D patch extraction
• Random flipping
• Intensity shifting
• Gaussian noise
• Intensity clipping

Loss

Supported loss functions:

• DiceLoss
• DiceCELoss

Optional voxel-frequency class weighting can be computed automatically.

Evaluation

Performance is measured using the Dice coefficient for:

• Non-enhancing tumor
• Edema
• Enhancing tumor

Validation and testing are performed using sliding-window inference.

---

Uncertainty-Driven Sample Selection

The uncertainty-aware training strategy is implemented in client_uncertainty.py.

Instead of training equally on every case throughout the entire experiment, the client periodically estimates the uncertainty of each training sample.

The uncertainty is computed using normalized voxel-wise entropy of the network predictions over tumor voxels only. Cases with higher entropy are considered more informative because the model is less confident in its predictions.

During training:

1. Case-level uncertainty scores are computed periodically.
2. A memory bank stores the latest uncertainty value for every case.
3. Cases whose uncertainty falls below a predefined threshold are temporarily removed from the active training set.
4. After a fixed number of federated rounds, excluded cases are automatically reintroduced into training.
5. A dynamic sampler updates the DataLoader so that subsequent epochs focus on the currently active cases.

This strategy encourages the model to spend more training time on difficult examples while reducing computation on samples that have already been learned.

---

Running an Experiment

Create the environment:

conda env create -f environment.yml
conda activate nvflare

Run a simulation:

python job.py --config config.yml

For HPC systems, example SLURM scripts are provided:

bash submit_job.sh

---

Configuration

Two configuration files control the experiments:

config.yml

Contains server-side settings:

• Number of clients
• Number of federated rounds
• Model architecture
• Dataset partitioning
• Learning rate
• Training script

client_config.yml

Contains client-side settings:

• Data augmentation
• Model parameters
• Loss function
• Metrics
• Optimizer
• Batch size
• Local epochs
• Validation interval
• Uncertainty-driven sampling options

---

Dependencies

Major libraries used:

• Python 3.10
• PyTorch
• MONAI
• NVFlare
• NumPy
• Pandas
• NiBabel
• SciPy

A complete list of dependencies is provided in environment.yml.

---

Citation

If you use this repository in your research, please cite the following resources:

• FeTS 2022 Challenge
• FeTS 2024 Challenge
• MONAI
• NVFlare

---

License

Please add an appropriate license before releasing this repository publicly.