Awesome
Variational Diffusion Models
Jax/Flax Code for reproducing some key results of Variational Diffusion Models (https://arxiv.org/abs/2107.00630).
Standalone Colabs
At colab/SimpleDiffusionColab.ipynb
you will find an independent and stand-alone Colab implementation of a Variational Diffusion Model (VDM), serving as an easy-to-understand demonstration of the code and principles behind the paper. Link to open in Colab. (Thanks a lot to Alex Alemi and Ben Poole for this implementation.)
At colab/2D_VDM_Example.ipynb
you will find an even more basic implementation, on a 2D swirl dataset and using MLPs. Link to open in Colab.
Setup: Installing required libraries
This code was tested on a TPU-v3 machine. For instructions on how to launch such a machine, see 'Setting up a v3-8
machine' below.
To install the required libraries on a TPU machine:
pip3 install -U pip
sudo pip uninstall -y jax jaxlib
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip3 install --upgrade -r requirements.txt
Alternatively, on a GPU machine, this should work:
pip3 install -U pip
pip3 install --upgrade jax jaxlib
pip3 install --upgrade -r requirements.txt
Train/evaluate: CIFAR-10 without data augmentation
The commands below assume that the code is checked out at the ./vdm
directory, such that this README is located at `./vdm/README'.
To evaluate from a pre-trained checkpoint:
python3 -m vdm.main --mode=eval --config=vdm/configs/cifar10.py --workdir=[workdir] --checkpoint='gs://gresearch/vdm/cifar10/checkpoint-final/checkpoints-0'
where [workdir]
is a directory to write results to, such as '/tmp/vdm-workdir'
, or a Google Cloud Storage address (gs://[your-address]
). Running the command above will print out a bunch of statistics, including eval_bpd=2.637
, which matches the result in the paper (2.64).
To train:
python3 -m vdm.main --config=vdm/configs/cifar10.py --workdir=[workdir]
CIFAR-10 with data augmentation
We also provide code for training a continuous-time VDM of CIFAR-10 with data augmentation. The model has some minor differences with the one described in the paper, but achieves similar performance.
To evaluate:
python3 -m vdm.main --mode=eval --config=vdm/configs/cifar10_aug.py --workdir=[workdir] --checkpoint='gs://gresearch/vdm/cifar10/aug-checkpoint-final/checkpoints-0'
This reports a bpd results slightly worse than the paper: eval_bpd=2.522, versus the paper's result of 2.49. We suspect this is due to small the fact that this open source version of the model was trained on a smaller batch size, in addition to a small difference in the score network implementation, namely how it is conditioned.
To train:
python3 -m vdm.main --config=vdm/configs/cifar10_aug.py --workdir=[workdir]
At the time of writing, training this model is too slow and memory-intensive to run on a single TPU or GPU machine. Therefore, we run on a v3-64
TPU pod, as explained below.
(Optional) Setting up a v3-8
machine on Google Cloud Engine (GCE)
This code was tested on a v3-8
TPU machine on Google Cloud Engine (GCE). The machine was created using:
gcloud alpha compute tpus tpu-vm create --project=[project] --zone=[zone] [machine-name] --accelerator-type=v3-8 --version=v2-alpha
SSH to the instance:
gcloud alpha compute tpus tpu-vm ssh --project=[project] --zone=[zone] [machine-name]
Then, we copied (or cloned through git) the code to a directory ~/vdm
on the machine. Then, on the instance, we installed libraries:
cd ~; bash vdm/sh/setup-vm-tpuv3.sh
Training on a TPU pod on GCE
Create the machine with:
gcloud alpha compute tpus tpu-vm create --project=[project] --zone=[zone] [machine-name] --accelerator-type=v3-64 --version=v2-alpha
Copy (or clone through git) the code to a directory ~/vdm
on the machine. Then, install libraries:
gcloud alpha compute tpus tpu-vm ssh --project=[project] --zone=[zone] [machine-name] --worker=all --command="$(<~/vdm/sh/setup-tpu-machine-v3.sh)"
Start training with:
gcloud alpha compute tpus tpu-vm ssh --project=[project] --zone=[zone] [machine-name] --worker=all --command="~/.local/bin/ipython vdm/main.py -- --workdir=[dir_to_save_logs] --config=vdm/configs/cifar10_aug.py"
Stop training on all machines by running:
gcloud alpha compute tpus tpu-vm ssh --project=[project] --zone=[zone] [machine-name] --worker=all --command="killall main.py"
Acknowledgments
We thank Ruiqi Gao for substantial contributions to this public VDM codebase, and to Alex Alemi and Ben Poole for implementing colab/SimpleDiffusionColab.ipynb
.
Disclaimer
This is not an officially supported Google product.