BioDiffusion: A Versatile Diffusion Model for Biomedical Signal Synthesis

Description

This repository serves as the primary codebase for implementing a Diffusion model, specifically designed for the generation of synthetic signals. The model is a pivotal component used in the research paper titled "BioDiffusion: A Versatile Diffusion Model for Biomedical Signal Synthesis" (accessible here).

Setup

Before running the code, ensure that you have the following prerequisites installed:

• Python 3.x
• PyTorch
• Nvidia CUDA toolkit and cuDNN (for GPU acceleration)

pip install torch torchvision
conda install cudatoolkit

Conda Virtual Environment

Create the Conda virtual environment using the environment file:

conda env create -f environment.yml

# dynamically set python path for the environment
conda activate BioDiffusion
conda env config vars set PYTHONPATH=$(pwd):$(pwd)/src

Training

In order to use our model first you will need to retrain it:

Unconditional diffusion model for 3 channel data:

python ddpm.py

Conditional diffusion model for 3 channel data:

python ddpm_conditional.py

Classifier-free diffusion model for 1 channel data:

python load_dataset.py #KaggleKey #KaggleName
python ddpm1d_cls_free.py

Signal conditional diffusion model for 1 channel data:

python load_dataset.py #KaggleKey #KaggleName
python ddpm1d_sign_cond.py

Sampling

After training, you can sample from the trained model using the following steps:

Unconditional Diffusion Model for 3 Channel Data:

# Set device and load the pre-trained UNet model
device = "cuda:2"
model = UNet().to(device)
ckpt = torch.load("../../src/models/DDPM_Unconditional/ckpt.pt")
model.load_state_dict(ckpt)
# Create a Diffusion model instance and sample from it
diffusion = Diffusion(img_size=32, device=device)
x = diffusion.sample(model, 10)

Conditional Diffusion Model for 3 Channel Data:

# Set the number of samples and device
n = 10
device = "cuda:3"
# Create a Diffusion model instance and load the trained model checkpoint
diffusion = Diffusion(img_size=32, device=device)
diffusion.load("../../src/models/DDPM_conditional")
# Prepare labels and sample from the diffusion model
labels = torch.full((n,), 1).long().to(diffusion.device)
sampled_images = diffusion.sample(use_ema=False, labels=labels)

Classifier-Free Diffusion Model for 1 Channel Data:

# Set the number of samples, device, and create the Unet1D_cls_free model
n = 10
device = "cuda:3"
model = Unet1D_cls_free(
    dim=64,
    dim_mults=(1, 2, 4, 8),
    num_classes=5,
    cond_drop_prob=0.5,
    channels=1
)
# Load the pre-trained model checkpoint
ckpt = torch.load("../../src/signal/checkpoint/DDPM1D_cls_free_MITBIH/checkpoint.pt")
model.load_state_dict(ckpt['model_state_dict'])
model = model.to(device)
# Create the GaussianDiffusion1D_cls_free model and sample from it
diffusion = GaussianDiffusion1D_cls_free(
    model,
    seq_length=128,
    timesteps=1000
).to(device)
y = torch.Tensor([0] * n).long().to(device)
x = diffusion.sample(classes=y, cond_scale=3.)

Self-Conditional Diffusion Model for 1 Channel Data:

# Set the device and create the Unet1D model with self-conditioning
device = "cuda:3"
model = Unet1D(
    dim=64,
    self_condition=True,
    dim_mults=(1, 2, 4, 8),
    channels=1
)
# Load the pre-trained model checkpoint
ckpt = torch.load("../../src/signal/checkpoint/DDPM1D_Selfconditional_maskedCond/checkpoint.pt")
model.load_state_dict(ckpt['model_state_dict'])
model = model.to(device)
# Create the GaussianDiffusion1D model and sample from it
# seq_length must be able to be divided by dim_mults
diffusion = GaussianDiffusion1D(
    model,
    seq_length=128,
    timesteps=1000,
    objective='pred_v'
).to(device)

Make sure to adjust the file paths and model names as needed.

Contributing

Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.

Citation

Feel free to cite our paper using this .bibtex or .cff formats in this repository.

@misc{li2024biodiffusion,
      title={BioDiffusion: A Versatile Diffusion Model for Biomedical Signal Synthesis}, 
      author={Xiaomin Li and Mykhailo Sakevych and Gentry Atkinson and Vangelis Metsis},
      year={2024},
      eprint={2401.10282},
      archivePrefix={arXiv},
      primaryClass={eess.SP}
}