| --- |
| 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](https://github.com/londonaicentre/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 |
|
|
| ```text |
| 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 |
|
|
| ```python |
| 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`](https://huggingface.co/datasets/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. |
|
|
