Unconditional Diffusion model

This repo presents our paper implementation of the unconditional diffusion model, with its popular optimizations, from Denoising Diffusion Probabilistic Models by Ho et al and Improved Denoising Diffusion Probabilistic Models by Nichol and Dhariwal. The pipeline contains training, sampling, and evaluation functions meant to be run on the HPC.

We show the correctness of our pipeline by training unconditional diffusion models on the landscapes (LHQ) and celebrity A (CelebAHQ) datasets, generating realistic images with a resolution of 128x128px.

Diffusion models are a class of generative models that offer a unique approach to modeling complex data distributions by simulating a stochastic process, known as a diffusion process, that gradually transforms data from a simple initial distribution into a complex data distribution. More specifically, the simple distribution is given by Gaussian Noise which is iteratively denoised into coherent images through modeling the entire data distribution present in the training set.

Sample Examples

Unconditional Landscape Generation:

Unconditional Celebrity Face Generation:

Recreating Results

We used the following modules:

• Python/3.10.4
• CUDA/11.8.0
• cuDNN/8.6.0.163-CUDA-11.8.0

The python libraries are stored in the requirements.txt file and can be installed with: pip install -r requirements.txt a virtual environment is recommended.

Model Training

To train the model, follow the steps:

1. Create an experiment folder with the experiment creator and the wished settings for training.
2. within the repository folder, run python main.py train <path to experimentfolder>/settings

Model Sampling

1. Make sure that the checkpoint file is within the ldm_trained folder within the experiment folder. Alternatively one can create this folder manually and add the checkpoint file.
2. Also make sure that the correct checkpoint name is given in the json file settings/sample_samplesettings.json otherwise the sampling will be done with randomly initialized weights.
3. within the repository folder, run python main.py sample <path to experimentfolder>/settings

Model Evaluation

To evaluate the performance of the unconditional diffusion model:

1. Ensure the settings are set to desired values in settings/evaluation_settings.json
2. within the repository folder, run python main.py evaluate <path to experimentfolder>/settings
3. For a detailed overview on evaluation metrics, refer to evaluation_readme.

Pipeline Description

This repository houses our comprehensive pipeline, designed to conveniently train, sample from, and evaluate our unconditional diffusion model. The pipeline is initiated via the experiment_creator.ipynb notebook, which is separately run our local machine. This notebook allows for the configuration of every aspect of the diffusion model, including all hyperparameters. These configurations extend to the underlying neural backbone UNet, as well as the training parameters, such as training from checkpoint, Weights & Biases run name for resumption, optimizer selection, adjustment of the CosineAnnealingLR learning rate schedule parameters, and more. Moreover, it includes parameters for evaluating a and sampling images via a trained diffusion models.

Upon execution, the notebook generates individual JSON files, encapsulating all the hyperparameter information. When running the model on the HPC, we can choose between the operations 'train', 'sample', and 'evaluate'. These operations automatically extract the necessary hyperparameters from the JSON files and perform their respective tasks. This process is managed by the main.py file. The remaining files contain all the necessary functions optimized for HPC to perform the aforementioned tasks.

Every uniquely trained diffusion model has its own experiment folder, given by its WANDB run name. It holds four different directories: settings, trained_ddpm, samples, and evaluations. The settings folder holds the JSON files specifying the diffusion model's configurations as well as the arguments for the training, sampling, and evaluation functions. The trained_ddpm folder contains .pth files storing the weights and biases of this experiment's diffusion model, which have been saved at different epoch milestones while training. Upon resuming training, the pipeline takes the specified model in trained_ddpm and continues training from there. When sampling images from these trained diffusion models, the samples are stored in different directories for the milestones under the names epoch_{i}. This is done so we know what epoch i version of the diffusion model was used to generate these samples. This is done analogously for the evaluation.