The Official PyTorch Implementation of "NVAE: A Deep Hierarchical Variational Autoencoder" (NeurIPS 2020 Spotlight Paper)

NVAE is a deep hierarchical variational autoencoder that enables training SOTA likelihood-based generative models on several image datasets.

Requirements

NVAE is built in Python 3.7 using PyTorch 1.6.0. Use the following command to install the requirements:

pip install -r requirements.txt

Set up file paths and data

We have examined NVAE on several datasets. For large datasets, we store the data in LMDB datasets for I/O efficiency. Click below on each dataset to see how you can prepare your data. Below, $DATA_DIR indicates the path to a data directory that will contain all the datasets and $CODE_DIR refers to the code directory:

Running the main NVAE training and evaluation scripts

We use the following commands on each dataset for training NVAEs on each dataset for Table 1 in the paper. In all the datasets but MNIST normalizing flows are enabled. Check Table 6 in the paper for more information on training details. Note that for the multinode training (more than 8-GPU experiments), we use the mpirun command to run the training scripts on multiple nodes. Please adjust the commands below according to your setup. Below IP_ADDR is the IP address of the machine that will host the process with rank 0 (see here). NODE_RANK is the index of each node among all the nodes that are running the job.

If for any reason your training is stopped, use the exact same commend with the addition of --cont_training to continue training from the last saved checkpoint. If you observe NaN, continuing the training using this flag usually will not fix the NaN issue.

Known Issues

Monitoring the training progress

While running any of the commands above, you can monitor the training progress using Tensorboard:

Post-training sampling, evaluation, and checkpoints

How to construct smaller NVAE models

In the commands above, we are constructing big NVAE models that require several days of training in most cases. If you'd like to construct smaller NVAEs, you can use these tricks:

• Reduce the network width: --num_channels_enc and --num_channels_dec are controlling the number of initial channels in the bottom-up and top-down networks respectively. Recall that we halve the number of channels with every spatial downsampling layer in the bottom-up network, and we double the number of channels with every upsampling layer in the top-down network. By reducing --num_channels_enc and --num_channels_dec, you can reduce the overall width of the networks.

• Reduce the number of residual cells in the hierarchy: --num_cell_per_cond_enc and --num_cell_per_cond_dec control the number of residual cells used between every latent variable group in the bottom-up and top-down networks respectively. In most of our experiments, we are using two cells per group for both networks. You can reduce the number of residual cells to one to make the model smaller.

• Reduce the number of epochs: You can reduce the training time by reducing --epochs.

• Reduce the number of groups: You can make NVAE smaller by using a smaller number of latent variable groups. We use two schemes for setting the number of groups:

  1. An equal number of groups: This is set by --num_groups_per_scale which indicates the number of groups in each scale of latent variables. Reduce this number to have a small NVAE.

  2. An adaptive number of groups: This is enabled by --ada_groups. In this case, the highest resolution of latent variables will have --num_groups_per_scale groups and the smaller scales will get half the number of groups successively (see groups_per_scale in utils.py). We don't let the number of groups go below --min_groups_per_scale. You can reduce the total number of groups by reducing --num_groups_per_scale and --min_groups_per_scale when --ada_groups is enabled.

Understanding the implementation

If you are modifying the code, you can use the following figure to map the code to the paper.

Traversing the latent space

We can generate images by traversing in the latent space of NVAE. This sequence is generated using our model trained on CelebA HQ, by interpolating between samples generated with temperature 0.6. Some artifacts are due to color quantization in GIFs.

License

Please check the LICENSE file. NVAE may be used non-commercially, meaning for research or evaluation purposes only. For business inquiries, please contact researchinquiries@nvidia.com.

You should take into consideration that VAEs are trained to mimic the training data distribution, and, any bias introduced in data collection will make VAEs generate samples with a similar bias. Additional bias could be introduced during model design, training, or when VAEs are sampled using small temperatures. Bias correction in generative learning is an active area of research, and we recommend interested readers to check this area before building applications using NVAE.

Bibtex:

Please cite our paper, if you happen to use this codebase:

@inproceedings{vahdat2020NVAE,
  title={{NVAE}: A Deep Hierarchical Variational Autoencoder},
  author={Vahdat, Arash and Kautz, Jan},
  booktitle={Neural Information Processing Systems (NeurIPS)},
  year={2020}
}