flwr new @flwrlabs/fed-med-seg

Federated Brain Tumor Segmentation with Flower and MONAI

This example demonstrates how to perform federated learning for 3D brain tumor segmentation using Flower, PyTorch, and MONAI. It uses the fed-brats dataset hosted on Hugging Face and supports both simulation and deployment workflows.

The project includes:

• A 3D U-Net model for volumetric segmentation
• Data loading pipelines using MONAI transforms
• Federated training with Flower
• Support for both IID and natural (site-based) partitioning

Fetch the App

Install Flower:

pip install flwr

Fetch the app:

flwr new @flwrlabs/fed-med-seg

Then, install dependencies:

cd fed-med-seg && pip install -e .

Project structure:

fed-med-seg
├── fed_med_seg
│   ├── __init__.py
│   ├── client_app.py   # Client-side training logic
│   ├── server_app.py   # Server-side orchestration and evaluation
│   └── task.py         # Model, data loading, training, evaluation
├── pyproject.toml      # Dependencies and configuration
└── README.md

Run the App

This Flower App supports both simulation mode and deployment mode without code changes.

Run with the Simulation Engine

In simulation mode:

• The fed-brats dataset is automatically downloaded from Hugging Face

• Data is partitioned across clients using:

  • iid (random split), or
  • natural (by hospital/site)

Run with default configuration:

flwr run .

Override configuration (example):

flwr run . --run-config "num-server-rounds=5 batch-size=2"

Key configurable parameters (from pyproject.toml):

• seed: base random seed for deterministic server, client, and MONAI execution (default: 42)
• num-server-rounds: number of FL rounds
• local-epochs: local training epochs per client
• batch-size: training batch size
• roi-x/y/z: 3D crop size for training
• partitioner: iid or natural
• learning-rate-max/min: cosine annealing schedule

For a reproducible run with a different seed:

PYTHONHASHSEED=2026 flwr run . --run-config "seed=2026"

The server uses the base seed, while each client round receives a stable seed derived from the base seed, client partition, and server round. PyTorch and MONAI deterministic execution are enabled, including random transforms and DataLoader workers. The run fails instead of silently using a known nondeterministic kernel. Exact replay also requires the same data, participating clients, dependency versions, hardware, and runtime configuration.

The seed must be an integer between 0 and 4294967295. The run seed controls model initialization, training randomness, and client sampling. Dataset partitioning keeps its existing fixed seed so seed sweeps compare the same client data. Sampling and aggregation use sorted Flower node IDs, so exact replay requires the same node-ID-to-partition mapping and client availability. For local data, the client seed uses partition-id when provided, otherwise data-path; give clients distinct stable identities to obtain distinct streams. Set PYTHONHASHSEED before starting the ServerApp and ClientApp runtimes (including remote workers), since setting it in an already running interpreter has no effect on that interpreter's hash seed. Timing and memory metrics are not deterministic.

Model

The model is a 3D U-Net implemented using MONAI:

• Input channels: 4 MRI modalities (t1n, t1c, t2w, t2f)
• Output channels: segmentation classes (default: 4)
• Architecture: encoder-decoder with residual units

Data Pipeline

Data is loaded from the Hugging Face dataset:

• Dataset: flwrlabs/fed-brats
• Automatically downloaded and cached locally
• Converted into MONAI-compatible format

Preprocessing includes:

• Resampling to 1mm spacing
• Intensity normalization
• Label remapping
• Random cropping and augmentation (training only)

Two modes:

• Simulation mode → uses FederatedDataset
• Deployment mode → loads pre-partitioned data from disk

Training

Each client:

• Receives the global model

• Trains locally using:

  • Dice + Cross Entropy loss (cross entropy treats voxels as classification examples to avoid the spatial CUDA loss reduction that rejects strict determinism; Dice still operates on the original volumes)
  • Adam optimizer

• Applies cosine annealing learning rate

Evaluation

Server-side evaluation:

• Uses centralized test split

• Applies sliding window inference

• Reports:

  • Loss
  • Mean Dice score

Run with the Deployment Engine

For deployment, you must provide local dataset partitions.

Step 1: Prepare data

Download and partition the dataset manually (or via Flower Datasets), then store partitions locally.

Step 2: Start SuperNodes

Each node must point to its local data:

flower-supernode \
    --insecure \
    --superlink <SUPERLINK-FLEET-API> \
    --node-config="data-path=/path/to/local_partition"

Step 3: Run the federation

flwr run . <SUPERLINK-CONNECTION> --stream

Benchmarking and System Metrics

This app writes a benchmark summary next to the standard Flower result pickle:

result_<run-name>_communication.json

The summary includes per-round and total communication volume:

• total_comm_bytes
• comm_bytes_total per training round

Enable system metric tracking with:

flwr run . <SUPERLINK-CONNECTION> --stream --run-config "benchmark-system-metrics=true"

When enabled, the benchmark summary also includes:

• client_train_time_sec
• server_aggregation_time_sec
• round_wall_clock_sec
• client_peak_cpu_memory_mb
• client_peak_gpu_memory_mb

Server-side centralized evaluation can be disabled for benchmark-only runs:

flwr run . <SUPERLINK-CONNECTION> --stream --run-config "benchmark-run-server-eval=false"

Deployment Fingerprints

Use the following fed-brats partition fingerprints when checking that each deployment client has the expected local data partition:

Client	Partition ID	Examples	Dataset fingerprint
0	0	379	dfad4dd6243e677e2cd6c4d4e20ffa00e03ca61c5f72cbef2d68b43a12b2ef19
1	1	25	5c1b805391929fb7a6fb787f540fe32f0383ebb6e88c2a8cfecbff85e18e1c37
2	2	256	d2a719dbbc5c7df5c60e812ab0d3b2737ebca94c1be12a3a9375032c074c4df2
3	3	160	95d047fbf66c27ad8bdb47655df92314cef493723b2b8bebdc557b093182a0af
4	4	477	177b73c74d44d6ad92fb511ae1d56a5aa7c6b68e93ea0838b6a3f634feab13c3

Notes

• GPU is automatically used if available
• Large 3D volumes are handled via sliding window inference
• Data loading uses MONAI CacheDataset for efficiency
• Supports both research (simulation) and real-world deployment