vaesbd icon indicating copy to clipboard operation
vaesbd copied to clipboard

Variational Autoencoder with Spatial Broadcast Decoder

VAE-SBD

PyTorch implementation of the Variational Autoencoder with Spatial Broadcast Decoder.

Instructions

Create a conda environment with all the requirements (edit environment.yml if you want to change the name of the environment):

conda env create -f environment.yml

Activate the environment

source activate pytorch

Generate the data

python generate.py

Train the model

python train.py

Training on GPU for 100 epochs should take about 3 hours.

Acknowledgements

Thanks to Nicholas Watters for very helpful pointers for the implementation.