README.md

metadata

license: apache-2.0
library_name: monai
tags:
  - medical-imaging
  - chest-x-ray
  - diffusion
  - ddpm
  - federated-learning
  - nvflare

Chest X-ray diffusion model: FLIP tutorial model

A 2-D denoising diffusion model (DDPM) that generates 128×128 chest X-rays, trained with the chest X-ray diffusion tutorial of FLIP, the Federated Learning Interoperability Platform (fl-tutorials/nvflare/image_generation/xray_diffusion/).

It was trained federated, with NVIDIA FLARE simulating two sites on one GPU, entirely on synthetic X-rays. It is a tutorial model, not for clinical use.

Contents

diffusion_model.safetensors   the weights, float32 as trained
provenance.json               commit, install line, launch command, versions, hardware, wall time, final losses
code/                         the app files exactly as the run executed them
evaluation/sample_xray.py     the script that drew the samples above
figures/

Model and training

MONAI DiffusionModelUNet in 2-D, unconditional (its only input besides the image is the noise step), 22.1 M parameters, trained from scratch: four levels of 64, 128, 192 and 256 channels, with self-attention at the two coarsest. It predicts the noise added at a random step of a 1000-step DDPM schedule, with an MSE loss.

Data	Ark+ tutorial splits, about 1.9k images per site, two sites
Preprocessing	grayscale, resized to 128×128, scaled to [0, 1]; light random affine in training
Schedule	2 rounds × 300 epochs, batch 6
Learning rate	1e-3 decaying to 5e-5 over the first 200 epochs
Final validation loss	0.0039 (site 1), 0.0036 (site 2), on each site's held-out 20 %
Wall time	5 h 05 min on one NVIDIA GeForce RTX 5090
Software	PyTorch 2.13.0, MONAI 1.6.0, NVIDIA FLARE 2.9.0

Samples saved during training looked like chest X-rays from about epoch 150.

Use

import json
from monai.networks.nets import DiffusionModelUNet
from safetensors.torch import load_file

net_config = json.load(open("code/config.json"))["net_config"]
model = DiffusionModelUNet(spatial_dims=2, **net_config["diffusion_model"])
weights = load_file("diffusion_model.safetensors")
model.load_state_dict({k.removeprefix("diffusion_model."): v for k, v in weights.items()})

Sample with MONAI's DDPMScheduler and DiffusionInferer (1000 steps), as in evaluation/sample_xray.py.

Reproducing it

Install flip-utils from the commit in provenance.json with its install line, check out FLIP at the same commit, copy code/ over the tutorial's app_files/, download the data with make -C fl-tutorials download-arkplus-finetuning-data, and run the launch command. The run predates the tutorial's move to image_generation/, so provenance.json gives the directory it ran from. Training on a GPU is not bit-for-bit repeatable, so expect samples of the same quality, not the same samples.

Data and licence

The Ark+ tutorial splits (aicentreflip/tutorials-arkplus-cxr-classification) are fully synthetic chest X-rays (DeCaf balanced_synthetic_split), with no patient data, licensed Apache-2.0. This model is licensed Apache-2.0 too.