PyTorch VAE

A collection of Variational AutoEncoders (VAEs) implemented in PyTorch with focus on reproducibility. The aim of this project is to provide a quick and simple working example for many of the cool VAE models out there. All the models are trained on the CelebA dataset for consistency and comparison. The architecture of all the models are kept as similar as possible with the same layers, except for cases where the original paper necessitates a radically different architecture. Here are the results of each model.

Requirements

• Python >= 3.5
• PyTorch >= 1.3
• Pytorch Lightning >= 0.6.0 (GitHub Repo)
• CUDA enabled computing device

Installation

$ git clone https://github.com/AntixK/PyTorch-VAE
$ cd PyTorch-VAE
$ pip install -r requirements.txt

Usage

$ cd PyTorch-VAE
$ python run.py -c configs/<config-file-name.yaml>

Config file template

model_params:
  name: "<name of VAE model>"
  in_channels: 3
  latent_dim: 
    .         # Other parameters required by the model
    .
    .

exp_params:
  data_path: "<path to the celebA dataset>"
  img_size: 64    # Models are designed to work for this size
  batch_size: 64  # Better to have a square number
  LR: 0.005
  weight_decay:
    .         # Other arguments required for training, like scheduler etc.
    .
    .

trainer_params:
  gpus: 1         
  max_nb_epochs: 50
  gradient_clip_val: 1.5
    .
    .
    .

logging_params:
  save_dir: "logs/"
  name: "<experiment name>"
  manual_seed: 

---

Results

Model	Paper	Reconstruction	Samples
VAE (Code, Config)	Link
Conditional VAE (Code, Config)	Link
WAE - MMD (RBF Kernel) (Code, Config)	Link
WAE - MMD (IMQ Kernel) (Code, Config)	Link
Beta-VAE (Code, Config)	Link
Disentangled Beta-VAE (Code, Config)	Link
IWAE (K = 5) (Code, Config)	Link
MIWAE (K = 5, M = 3) (Code, Config)	Link
DFCVAE (Code, Config)	Link
MSSIM VAE (Code, Config)	Link
Categorical VAE (Code, Config)	Link
Joint VAE (Code, Config)	Link
Info VAE (Code, Config)	Link
LogCosh VAE (Code, Config)	Link
SWAE (200 Projections) (Code, Config)	Link
VQ-VAQ (K = 512, D = 64) (Code, Config)	Link		N/A

TODO

• VanillaVAE
• Beta VAE
• DFC VAE
• MSSIM VAE
• IWAE
• MIWAE
• WAE-MMD
• Conditional VAE
• Categorical VAE (Gumbel-Softmax VAE)
• Joint VAE
• Disentangled beta-VAE
• InfoVAE
• LogCosh VAE
• SWAE
• VQVAE (in progress)
• Ladder VAE (Doesn't work well)
• Gamma VAE (Doesn't work well)
• Vamp VAE (Doesn't work well)
• Beta TC-VAE
• PixelVAE

Contributing

If you have trained a better model, using these implementations, by fine-tuning the hyper-params in the config file, I would be happy to include your result (along with your config file) in this repo, citing your name 😊.

Additionally, if you would like to contribute some models, check out the TODO for models that are pending or doesn't seem to work well.

License

Apache License 2.0

Permissions	Limitations	Conditions
✔️ Commercial use	❌ Trademark use	ⓘ License and copyright notice
✔️ Modification	❌ Liability	ⓘ State changes
✔️ Distribution	❌ Warranty
✔️ Patent use
✔️ Private use

Citation

@misc{Subramanian2020,
  author = {Subramanian, A.K},
  title = {PyTorch-VAE},
  year = {2020},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/AntixK/PyTorch-VAE}}
}

---