Adaptation of Diffusion Probabilistic Models for Inpainting Missing Data

This repository extends the original DPM code to support inpainting datasets with general missing data supplied in the form of mask images.

Also, it includes the small NORB dataset which wasn't in the original paper.

Usage

Training a model on MNIST follows the instructions in the original paper- after installing all dependencies, simply run python train.py.

To train on NORB- after installing all dependencies, run python train_norb.py --data_dir /path/. data_dir should be the path to the directory containing train and test image folders, as generated by the generate_norb_small script. i.e- path/to/data_dir/aug_norb_small_2D_train_image_data path/to/data_dir/norb_small_2D_test_image_data

Inpainting images with missing data locations given by a mask image is meant to be used with a dataset in the format created by this script.

The missing data dataset should have one index.txt file listing each corrupted image file with its label in Caffe ImageData format, and one index_mask.txt file in the same format listing the mask file corresponding to each corrupted image file.

To inpaint such a dataset, run for MNIST python inpaint_from_model.py --resume_file path/to/model.pkl --missing_dataset_path path/to/dataset for NORB python inpaint_from_model_norb.py --resume_file path/to/model.pkl --missing_dataset_path path/to/dataset

Where model.pkl is the model generated through training.

The new dataset will be created at the same location as the missing dataset, under the same name with the addition of _dpm_ip.

Diffusion Probabilistic Models

This repository provides a reference implementation of the method described in the paper:

Deep Unsupervised Learning using Nonequilibrium Thermodynamics.
Sohl-Dickstein, Jascha and Weiss, Eric A. and Maheswaranathan, Niru and Ganguli, Surya.
International Conference on Machine Learning. 2015
http://arxiv.org/abs/1503.03585

This implementation builds a generative model of data by training a Gaussian diffusion process to transform a noise distribution into a data distribution in a fixed number of time steps. The mean and covariance of the diffusion process are parameterized using deep supervised learning. The resulting model is tractable to train, easy to exactly sample from, allows the probability of datapoints to be cheaply evaluated, and allows straightforward computation of conditional and posterior distributions.

Using the Software

In order to train a diffusion probabilistic model on the default dataset of MNIST, simply install dependencies (see below), and then run python train.py.

Dependencies

1. Install Blocks and its dependencies following these instructions
2. Setup Fuel and download MNIST following these instructions.

As of October 16, 2015 this code requires the bleeding edge, rather than stable, versions of both Blocks and Fuel. (thanks to David Hofmann for pointing out that the stable release will not work due to an interface change)

Output

The objective function being minimized is the bound on the negative log likelihood in bits per pixel, minus the negative log likelihood under an identity-covariance Gaussian model. That is, it is the negative of the number in the rightmost column in Table 1 in the paper.

Logging information is printed to the console once per training epoch, including the current value of the objective on the training set.

Figures showing samples from the model, parameters, gradients, and training progress are also output periodically (every 25 epochs by default -- see train.py).

The samples from the model are of three types -- standard samples, samples inpainting the left half of masked images, and samples denoising images with Gaussian noise added (by default, the signal-to-noise ratio is 1). This demonstrate the straightforward way in which inpainting, denoising, and sampling from a posterior in general can be performed using this framework.

Here are samples generated by this code after 825 training epochs on MNIST, trained using the command run train.py:

Here are samples generated by this code after 1700 training epochs on CIFAR-10, trained using the command run train.py --batch-size 200 --dataset CIFAR10 --model-args "n_hidden_dense_lower=1000,n_hidden_dense_lower_output=5,n_hidden_conv=100,n_layers_conv=6,n_layers_dense_lower=6,n_layers_dense_upper=4,n_hidden_dense_upper=100":

Miscellaneous

Different nonlinearities - In the paper, we used softplus units in the convolutional layers, and tanh units in the dense layers. In this implementation, I use leaky ReLU units everywhere.

Original source code - This repository is a refactoring of the code used to run the experiments in the published paper. In the spirit of reproducibility, if you email me a request I am willing to share the original source code. It is poorly commented and held together with duct tape though. For most applications, you will be better off using the reference implementation provided here.

Contact - I would love to hear from you. Let me know what goes right/wrong! jaschasd@google.com