joint-vae's Introduction

Learning Disentangled Joint Continuous and Discrete Representations

Pytorch implementation of Learning Disentangled Joint Continuous and Discrete Representations (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.

Examples

MNIST

CelebA

FashionMNIST

Discrete and continuous factors on MNIST

Usage

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

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

Example usage

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
...

Trained models

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.