vdm - GitLite

KO | EN

gitlite — search

#javascript #hacktoberfest #typescript #golang #go #rust #deep-learning #linux #docker #java #vue #ios #openai #open-source

vdm

(★ 331)

No description available.

File Explorer

Download Latest Version (.zip)

- 2D_VDM_Example.ipynb
- SimpleDiffusionColab.ipynb
- cifar10.py
- cifar10_aug.py
- cifar10_aug_paper.py
- imagenet32.py
- imagenet64.py
- create-vms-tpuv3.sh
- setup-vm-tpuv3.sh
CONTRIBUTING.md
dataset.py
experiment.py
experiment_vdm.py
LICENSE
main.py
model_vdm.py
README.md
requirements.txt
train_state.py
utils.py

# Installation Guide

1. Get the code

git clone https://github.com/google-research/vdm

Downloads the entire project code from GitHub to your computer.

cd vdm

Moves into the project folder you just downloaded.

2. Official Install Script

Easy Recommended

Prerequisites

Python 3 Python is required to use pip.

pip3 install -U pip

Installs the package published on PyPI directly — no need to clone the source.

pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Installs the package published on PyPI directly — no need to clone the source.

pip3 install --upgrade -r requirements.txt

Installs the package published on PyPI directly — no need to clone the source.

pip3 install --upgrade jax jaxlib

Installs the package published on PyPI directly — no need to clone the source.

✅ After installing, open a new terminal and run the program's version command (e.g. --version) to confirm it worked.

Pulled directly from this repo's README.

3. Python

Easy

Prerequisites

Git Needed to download the project code from GitHub.
Python 3 On Windows, be sure to check 'Add Python to PATH' during installation.

pip3 install -U pip

Installs the package published on PyPI directly — no need to clone the source.

sudo pip uninstall -y jax jaxlib

Type this command into your terminal and run it.

pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Installs the package published on PyPI directly — no need to clone the source.

pip3 install --upgrade -r requirements.txt

Installs the package published on PyPI directly — no need to clone the source.

pip3 install --upgrade jax jaxlib

Installs the package published on PyPI directly — no need to clone the source.

✅ If it runs without errors and prints output in the terminal, it worked.

Pulled directly from this repo's README.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# 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](https://colab.research.google.com/github/google-research/vdm/blob/main/colab/SimpleDiffusionColab.ipynb). (Thanks a lot to [Alex Alemi](https://www.alexalemi.com/) and [Ben Poole](https://cs.stanford.edu/~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](https://colab.research.google.com/github/google-research/vdm/blob/main/colab/2D_VDM_Example.ipynb).

## 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](https://ruiqigao.github.io/) for substantial contributions to this public VDM codebase, and to [Alex Alemi](https://www.alexalemi.com/) and [Ben Poole](https://cs.stanford.edu/~poole/) for implementing `colab/SimpleDiffusionColab.ipynb`.

## Disclaimer

This is not an officially supported Google product.