RealNVP on MNIST (PyTorch, demo)

Normalizing flows (NFs) are an exciting new family of neural networks for density estimation, sampling, and likelihood inference [1]. A NF is built by combining invertible components with tractable Jacobian computations, and they are shown to generalize many known models, such as invertible ResNets or autoregressive architectures.

This code provides a self-contained demo in pure PyTorch of a real-world NF (RealNVP, [2]). We try to make our code as concise and simple as possible, with low dependency on external libraries.

The model

Facing directly an image generation task with a NF is challenging! For this reason, we simplify the model and train a RealNVP model on the intermediate embeddings generated by an auxiliary autoencoder:

We train a simple convolutional autoencoder and a classifier on MNIST, and RealNVP on the codes generated in this way. The RealNVP architecture is composed by coupling layers, alternating masks, and batch normalization layers to stabilize training. You can find more details on the notebook or in [1].

Running the code

You can run the example in the corresponding notebook. We first train the autoencoder on MNIST, then the RealNVP architecture. Below some generated samples after both trainings:

The code of the RealNVP can be found in realnvp.py:

• LinearBatchNorm, LinearCouplingLayer, and Permutation are the basic layers of the NF;
• SequentialFlow defines a generic NF as a sequence of layers, by extending torch.nn.Sequential;
• RealNVP builds the actual RealNVP used in our demo.

References

[1] Papamakarios, G., Nalisnick, E., Rezende, D.J., Mohamed, S. and Lakshminarayanan, B., 2019. Normalizing Flows for Probabilistic Modeling and Inference. https://arxiv.org/abs/1912.02762.

[2] Dinh, L., Sohl-Dickstein, J. and Bengio, S., 2016. Density estimation using Real NVP. https://arxiv.org/abs/1605.08803.