Reimplementation of Wasserstein Auto Encoder (WAE) with Wasserstein GAN based penalty D_Z in Tensorflow
This repository contians an reimplmentation of WAE with Tensorflow.
I made one tweak on top of the paper. I used Wasserstein distance to penalize an encoder $Q\phi$.
In order to do that, I trained Discriminator $D\gamma$ through critic loss with gradient penalty as Gularjani etal. suggested in [improved WGAN].
I (personally) believe that this implementation is much clearer and easy to read (, and more importantly, the code almost exactly matches with the algorithm shows on the paper), so I hope it will help someone who wants to digin more! Enjoy !
python main.py
Check main.py file to change target dataset or to adjust hyperparmeters such as z_dim, and etc…
See the MNIST Plot.ipynb
and CelebA Plot.ipynb
with Jupyter Notebook.
A pretrained model for both MNIST is included on the repository while a model for CelebA is uploaded on this place.
Please download the zip file and decompress it on assets/pretrained_models/celeba/last*
. Or, you can easily modify a path at the first cell on the notebook.
Reconstruction Results
(top): original images from MNIST validation set, (bottom): reconstructed image
It seems not sharp as the authors suggest, but it might due to not enough training and untuned hyperparameters such as lambda, number of layers, or etc.