Tutorial on generative diffusion models

A generative model is a type of machine learning model that can generate new data that is similar to the data it has been trained on ¹². This tutorial will give you a short introduction to diffusion models, which are a specific type of generative model.

A diffusion model is based on the idea of incrementally inverting a forward process of adding noise to data. In this forward process, the input data $x(0)$ is gradually transformed into a standard Gaussian distribution $x(T)$ by scaling and adding Gaussian noise at every time step $t \in \mathcal{N} \subset [0, T]$. Specifically, the forward process can be defined as:

$$ q(x_t | x_0) = \mathcal{N}(x_t; \alpha_t x_0, \sigma_t^2 I) $$

Overview

We will use denoising-diffusion-pytorch which is a diffusion model library written by lucidrains. For training we will use lightning.

Quickstart

First, clone and cd into the repository:

git clone https://github.com/weigertlab/diffusion_model_tutorial.git
cd diffusion_model_tutorial

If using BARD, simply open the notebook regularly in VSCode and choose the kernel diffusion.

Otherwise all the setup, including environment creation as well as data downloading, can be done by simply running this command:

source setup.sh

Example data

Generative models are typically trained on a dataset of (unlabeled) images. In this tutorial we will use two example datasets, showing images of

1. flywing membrane, or
2. dual color zebrafish retina.

Please choose one of the two datasets to train your own models. If you want to train a model on your own data, you can skip these steps. Note that the 2D images are given as a npz file, which is a compressed numpy array. For your custom data, you could as well use a folder with tiff files.

1. Flywing membrane³

2. Dual color zebrafish retina⁴

Training your own diffusion model

Please see diffusion.ipynb to see how to train a diffusion model on one of the two datasets.

References

Footnotes

1. Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851. ↩

2. Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. arXiv preprint arXiv:2011.13456, 2020. ↩

3. Prakash, M., Buchholz, T.-O., Schmidt, D., Krull, A., & Jug, F. (2020). Flywing (noise 0) dataset for microscopy image denoising and segmentation benchmark as used in DenoiSeg paper. Zenodo dataset ↩

4. Martin Weigert, Uwe Schmidt, Tobias Boothe, Andreas Müller, Alexandr Dibrov, Akanksha Jain, Benjamin Wilhelm, Deborah Schmidt, Coleman Broaddus, Sian Culley, Mauricio Rocha-Martins, Fabián Segovia-Miranda, Caren Norden, Ricardo Henriques, Marino Zerial, Michele Solimena, Jochen Rink, Pavel Tomancak, Loic Royer, Florian Jug, Eugene W. Myers. Content Aware Image Restoration: Pushing the Limits of Fluorescence Microscopy. Dataset ↩