Pytorch implementation of Learning Disentangled Joint Continuous and Discrete Representations (NIPS 2018).

This repo contains an implementation of JointVAE, a framework for jointly disentangling continuous and discrete factors of variation in data in an unsupervised manner.

The train_model.ipynb notebook contains code for training a JointVAE model.

The load_model.ipynb notebook contains code for loading a trained model.

from jointvae.models import VAE
from jointvae.training import Trainer
from torch.optim import Adam
from viz.visualize import Visualizer as Viz
# Build a dataloader for your data
dataloader = get_my_dataloader(batch_size=32)
# Define latent distribution
latent_spec = {'cont': 20, 'disc': [10, 5, 5, 2]}
# Build a Joint-VAE model
model = VAE(img_size=(3, 64, 64), latent_spec=latent_spec)
# Build a trainer and train model
optimizer = Adam(model.parameters())
trainer = Trainer(model, optimizer,
                  cont_capacity=[0., 5., 25000, 30.],
                  disc_capacity=[0., 5., 25000, 30.])
trainer.train(dataloader, epochs=10)
# Visualize samples from the model
viz = Viz(model)
samples = viz.samples()
# Do all sorts of fun things with model
...

The trained models referenced in the paper are included in the trained_models folder. The load_model.ipynb ipython notebook provides code to load and use these trained models.

The MNIST and FashionMNIST datasets can be automatically downloaded using torchvision.

All CelebA images were resized to be 64 by 64. Data can be found here.

All Chairs images were center cropped and resized to 64 by 64. Data can be found here.

If you find this work useful in your research, please cite using:

@inproceedings{dupont2018learning,
  title={Learning disentangled joint continuous and discrete representations},
  author={Dupont, Emilien},
  booktitle={Advances in Neural Information Processing Systems},
  pages={707--717},
  year={2018}
}

MIT