Wasserstein Autoencoders for MNIST

This is a re-implementation in Tensorflow 2 (2.2.0) of both the WAE-GAN and WAE-MMD models proposed in 1 on the MNIST dataset.

Random draws

These are a few samples decoded from the latent space. The quality of the generation is inferior to the original paper's implementation.

• WAE-GAN

• WAE-MMD

Reconstructions

Even rows contain original samples while odd rows their relative reconstruction.

• WAE-GAN

• WAE-MMD

Downloading & running

git clone https://github.com/w00zie/wae_mnist.git
cd wae_mnist

The code is structured as following:

• wae_mnist
  • models
  • saved_models
  • wae_gan
  • wae_mmd
  • data.py
  • train_gan.py
  • train_mmd.py
  • utils.py

The models dir contains the code for two models I've coded: model.py and func_model.py. They have the same architecture but the second one is in a functional form.

Both wae_gan and wae_mmd dirs contain the hyperparameters (config_gan.py and config_mmd.py) for their relative experiments. If you want to run a new training with your hparams you have to edit these files.

Files train_gan.py and train_mmd.py execute a training procedure with WAE-GAN and WAE-MMD algorithms respectively.

If you want to train a new WAE-GAN model just run

python train_gan.py

Inference

I provide a few (at the moment only one) trained models per algorithm. You can explore the saved_models/ directory, where, for each {experiment}, I provide the weights (models/), the hyperparameters (config.json) and the Tensorboard logs (events.out.tfevents....) that you can load and inspect by running tensorboard --logdir saved_model/{experiment}.

Files wae_gan/test_gan.py and wae_mmd/test_mmd.py are responsible for running an inference procedure. These two files create two .png images containing several samples that you can inspect (decoded from the latent space and reconstructed from the original distribution).

Let's say you have trained a new WAE-GAN model logged into ./runs/experiment_wgan_000 and you want to visualize the samples you are able to produce: just run

python wae_gan/test_gan.py --exp_dir ./runs/experiment_wgan_000

This will create two images (random.png and recons.png) under runs/experiment_wgan_000/img/

Requirements

tensorflow >= 2.2.0
matplotlib