This repository contains a PyTorch-only implementation of the research by Yang Song. All credit for the original research goes to him. Due to the complexity of setting up the project, I have ported all TensorFlow and JAX related code to PyTorch, with some assistance from AI tools. While the results may differ from the original implementation, I hope this will be helpful.
First, install PyTorch separately. Then, install the necessary Python packages by running:
pip install -r requirements.txt
Currently, only the following script and configuration have been tested. On two RTX 3090 Ti GPUs, the program runs at approximately 1 iteration per second. Training should complete in about 15 days. The expected image is
python main.py --config configs/subvp/cifar10_ddpm_continuous.py --mode train --workdir workdir
The following script requires to change the batch size (dropped to 12).
python main.py --config configs/subvp/cifar10_ddpmpp_continuous.py --mode train --workdir workdir
python main.py --mode train --workdir workdir --config configs/subvp/cifar10_ddpmpp_deep_continuous.py
python main.py --mode train --workdir workdir --config configs/subvp/cifar10_ncsnpp_continuous.py
python main.py --mode train --workdir workdir --config configs/subvp/cifar10_ncsnpp_deep_continuous.py
python main.py --mode train --workdir workdir --config configs/ve/cifar10_ddpm.py
The following section is derived from the original author's GitHub project.
This repo contains a PyTorch implementation for the paper Score-Based Generative Modeling through Stochastic Differential Equations
by Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole
We propose a unified framework that generalizes and improves previous work on score-based generative models through the lens of stochastic differential equations (SDEs). In particular, we can transform data to a simple noise distribution with a continuous-time stochastic process described by an SDE. This SDE can be reversed for sample generation if we know the score of the marginal distributions at each intermediate time step, which can be estimated with score matching. The basic idea is captured in the figure below:
Our work enables a better understanding of existing approaches, new sampling algorithms, exact likelihood computation, uniquely identifiable encoding, latent code manipulation, and brings new conditional generation abilities (including but not limited to class-conditional generation, inpainting and colorization) to the family of score-based generative models.
All combined, we achieved an FID of 2.20 and an Inception score of 9.89 for unconditional generation on CIFAR-10, as well as high-fidelity generation of 1024px Celeba-HQ images (samples below). In addition, we obtained a likelihood value of 2.99 bits/dim on uniformly dequantized CIFAR-10 images.
Aside from the NCSN++ and DDPM++ models in our paper, this codebase also re-implements many previous score-based models in one place, including NCSN from Generative Modeling by Estimating Gradients of the Data Distribution, NCSNv2 from Improved Techniques for Training Score-Based Generative Models, and DDPM from Denoising Diffusion Probabilistic Models.
It supports training new models, evaluating the sample quality and likelihoods of existing models. We carefully designed the code to be modular and easily extensible to new SDEs, predictors, or correctors.
Most models are now also available in ๐งจ Diffusers and accesible via the ScoreSdeVE pipeline.
Diffusers allows you to test score sde based models in PyTorch in just a couple lines of code.
You can install diffusers as follows:
pip install diffusers torch accelerate
And then try out the models with just a couple lines of code:
from diffusers import DiffusionPipeline
model_id = "google/ncsnpp-ffhq-1024"
# load model and scheduler
sde_ve = DiffusionPipeline.from_pretrained(model_id)
# run pipeline in inference (sample random noise and denoise)
image = sde_ve().images[0]
# save image
image[0].save("sde_ve_generated_image.png")
More models can be found directly on the Hub.
Please find a JAX implementation here, which additionally supports class-conditional generation with a pre-trained classifier, and resuming an evalution process after pre-emption.
In general, this PyTorch version consumes less memory but runs slower than JAX. Here is a benchmark on training an NCSN++ cont. model with VE SDE. Hardware is 4x Nvidia Tesla V100 GPUs (32GB)
Framework | Time (second per step) | Memory usage in total (GB) |
---|---|---|
PyTorch | 0.56 | 20.6 |
JAX (n_jitted_steps=1 ) |
0.30 | 29.7 |
JAX (n_jitted_steps=5 ) |
0.20 | 74.8 |
Run the following to install a subset of necessary python packages for our code
pip install -r requirements.txt
We provide the stats file for CIFAR-10. You can download cifar10_stats.npz
and save it to assets/stats/
. Check out #5 on how to compute this stats file for new datasets.
Train and evaluate our models through main.py
.
main.py:
--config: Training configuration.
(default: 'None')
--eval_folder: The folder name for storing evaluation results
(default: 'eval')
--mode: <train|eval>: Running mode: train or eval
--workdir: Working directory
-
config
is the path to the config file. Our prescribed config files are provided inconfigs/
. They are formatted according toml_collections
and should be quite self-explanatory.Naming conventions of config files: the path of a config file is a combination of the following dimensions:
- dataset: One of
cifar10
,celeba
,celebahq
,celebahq_256
,ffhq_256
,celebahq
,ffhq
. - model: One of
ncsn
,ncsnv2
,ncsnpp
,ddpm
,ddpmpp
. - continuous: train the model with continuously sampled time steps.
- dataset: One of
-
workdir
is the path that stores all artifacts of one experiment, like checkpoints, samples, and evaluation results. -
eval_folder
is the name of a subfolder inworkdir
that stores all artifacts of the evaluation process, like meta checkpoints for pre-emption prevention, image samples, and numpy dumps of quantitative results. -
mode
is either "train" or "eval". When set to "train", it starts the training of a new model, or resumes the training of an old model if its meta-checkpoints (for resuming running after pre-emption in a cloud environment) exist inworkdir/checkpoints-meta
. When set to "eval", it can do an arbitrary combination of the following-
Evaluate the loss function on the test / validation dataset.
-
Generate a fixed number of samples and compute its Inception score, FID, or KID. Prior to evaluation, stats files must have already been downloaded/computed and stored in
assets/stats
. -
Compute the log-likelihood on the training or test dataset.
These functionalities can be configured through config files, or more conveniently, through the command-line support of the
ml_collections
package. For example, to generate samples and evaluate sample quality, supply the--config.eval.enable_sampling
flag; to compute log-likelihoods, supply the--config.eval.enable_bpd
flag, and specify--config.eval.dataset=train/test
to indicate whether to compute the likelihoods on the training or test dataset. -
- New SDEs: inherent the
sde_lib.SDE
abstract class and implement all abstract methods. Thediscretize()
method is optional and the default is Euler-Maruyama discretization. Existing sampling methods and likelihood computation will automatically work for this new SDE. - New predictors: inherent the
sampling.Predictor
abstract class, implement theupdate_fn
abstract method, and register its name with@register_predictor
. The new predictor can be directly used insampling.get_pc_sampler
for Predictor-Corrector sampling, and all other controllable generation methods incontrollable_generation.py
. - New correctors: inherent the
sampling.Corrector
abstract class, implement theupdate_fn
abstract method, and register its name with@register_corrector
. The new corrector can be directly used insampling.get_pc_sampler
, and all other controllable generation methods incontrollable_generation.py
.
All checkpoints are provided in this Google drive.
Instructions: You may find two checkpoints for some models. The first checkpoint (with a smaller number) is the one that we reported FID scores in our paper's Table 3 (also corresponding to the FID and IS columns in the table below). The second checkpoint (with a larger number) is the one that we reported likelihood values and FIDs of black-box ODE samplers in our paper's Table 2 (also FID(ODE) and NNL (bits/dim) columns in the table below). The former corresponds to the smallest FID during the course of training (every 50k iterations). The later is the last checkpoint during training.
Per Google's policy, we cannot release our original CelebA and CelebA-HQ checkpoints. That said, I have re-trained models on FFHQ 1024px, FFHQ 256px and CelebA-HQ 256px with personal resources, and they achieved similar performance to our internal checkpoints.
Here is a detailed list of checkpoints and their results reported in the paper. FID (ODE) corresponds to the sample quality of black-box ODE solver applied to the probability flow ODE.
Checkpoint path | FID | IS | FID (ODE) | NNL (bits/dim) |
---|---|---|---|---|
ve/cifar10_ncsnpp/ |
2.45 | 9.73 | - | - |
ve/cifar10_ncsnpp_continuous/ |
2.38 | 9.83 | - | - |
ve/cifar10_ncsnpp_deep_continuous/ |
2.20 | 9.89 | - | - |
vp/cifar10_ddpm/ |
3.24 | - | 3.37 | 3.28 |
vp/cifar10_ddpm_continuous |
- | - | 3.69 | 3.21 |
vp/cifar10_ddpmpp |
2.78 | 9.64 | - | - |
vp/cifar10_ddpmpp_continuous |
2.55 | 9.58 | 3.93 | 3.16 |
vp/cifar10_ddpmpp_deep_continuous |
2.41 | 9.68 | 3.08 | 3.13 |
subvp/cifar10_ddpm_continuous |
- | - | 3.56 | 3.05 |
subvp/cifar10_ddpmpp_continuous |
2.61 | 9.56 | 3.16 | 3.02 |
subvp/cifar10_ddpmpp_deep_continuous |
2.41 | 9.57 | 2.92 | 2.99 |
Checkpoint path | Samples |
---|---|
ve/bedroom_ncsnpp_continuous |
![]() |
ve/church_ncsnpp_continuous |
![]() |
ve/ffhq_1024_ncsnpp_continuous |
![]() |
ve/ffhq_256_ncsnpp_continuous |
![]() |
ve/celebahq_256_ncsnpp_continuous |
![]() |
- When using the JAX codebase, you can jit multiple training steps together to improve training speed at the cost of more memory usage. This can be set via
config.training.n_jitted_steps
. For CIFAR-10, we recommend usingconfig.training.n_jitted_steps=5
when your GPU/TPU has sufficient memory; otherwise we recommend usingconfig.training.n_jitted_steps=1
. Our current implementation requiresconfig.training.log_freq
to be dividable byn_jitted_steps
for logging and checkpointing to work normally. - The
snr
(signal-to-noise ratio) parameter ofLangevinCorrector
somewhat behaves like a temperature parameter. Largersnr
typically results in smoother samples, while smallersnr
gives more diverse but lower quality samples. Typical values ofsnr
is0.05 - 0.2
, and it requires tuning to strike the sweet spot. - For VE SDEs, we recommend choosing
config.model.sigma_max
to be the maximum pairwise distance between data samples in the training dataset.
If you find the code useful for your research, please consider citing
@inproceedings{
song2021scorebased,
title={Score-Based Generative Modeling through Stochastic Differential Equations},
author={Yang Song and Jascha Sohl-Dickstein and Diederik P Kingma and Abhishek Kumar and Stefano Ermon and Ben Poole},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=PxTIG12RRHS}
}
This work is built upon some previous papers which might also interest you:
- Song, Yang, and Stefano Ermon. "Generative Modeling by Estimating Gradients of the Data Distribution." Proceedings of the 33rd Annual Conference on Neural Information Processing Systems. 2019.
- Song, Yang, and Stefano Ermon. "Improved techniques for training score-based generative models." Proceedings of the 34th Annual Conference on Neural Information Processing Systems. 2020.
- Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Proceedings of the 34th Annual Conference on Neural Information Processing Systems. 2020.