This repository contains the pipeline for training, sampling, and evaluation of our Conditional Diffusion Model on the HPC. It builds upon our code from the unconditional-diffusion repository.
We demonstrate our results by training conditional diffusion models to solve the tasks of class conditional image generation and inpainting. For the class labeled dataset, we chose to use the Animal Face (AFHQ) dataset containing three classes; dog, cat, and wildlife, each of them with a representation of 5000 images. For the inpainting dataset, we train the model on the same Landscape dataset (LHQ) as with the unconditional diffusion model, and generate our own labels by randomly drawing black rectangle masks which the model learns to inpaint.
For this purpose, we implement class and image conditioning mechanisms into our UNet. For class conditioning, we also make use of our implementation of classifier-free guided diffusion to achieve better sampling quality results.
These techniques are presented in the papers **GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models** by Nichol et al. and **Classifier-Free Diffusion Guidance** by Ho and Salamatin.
## Motivation
Conditional diffusion models are ...
## Sample Examples
...
## 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 these steps:
1. Create an experiment folder with the [experiment creator](experiment_creator.ipynb) and the desired 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 **trained_cdm** 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_settings.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
...
## Description
This repository houses our comprehensive pipeline, designed to conveniently train, sample from, and evaluate our Conditional diffusion model.
The pipeline is initiated via the experiment_creator.ipynb notebook, which may be separately run on 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 and sampling images via a trained diffusion model.
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.
