File size: 3,378 Bytes
380c9c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | ---
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.
|