Awesome Open Source
Awesome Open Source

VQ-VAE (Neural Discrete Representation Learning) Tensorflow


This repository implements the paper, Neural Discrete Representation Learning (VQ-VAE) in Tensorflow.

⚠️ This is not an official implementation, and might have some glitch (,or a major defect).


  • Python 3.5
  • Tensorflow (v1.3 or higher)
  • numpy, better_exceptions, tqdm, etc.
  • ffmpeg

Updated Result: ImageNet

  • [x] ImageNet

    Validation Set Images Reconstructed Images
    Imagenet original images Imagenet Reconstructed Images
    • Class Conditioned Sampled Image (Not cherry-picked, just random sample)



      coral reef


      brown bear

      pickup truck

    • I could not reproduce as sharp images as the author produced.

    • But, some of results seems understandable.

    • Usually, natural scene images having consistent pixel orders shows better result, such as Alp or coral reef.

    • More tuning might provide better result.

Updated Result: Sampling with PixelCNN

  • [x] Pixel CNN

    • MNIST Sampled Image (Conditioned on class labels)

      MNIST Sampled Images

    • Cifar10 Sampled Image (Conditioned on class labels)

      Cifar10 Sampled Imagesl

      From top row to bottom, the sampled images for classes (airplane,auto,bird,cat,deer,dog,frog,horse,ship,truck)

      Not that satisfying so far; I guess hyperparameters for VQ-VAE should be tuned first to generate more sharper result.


All training is done with Quadro M4000 GPU. Training MNIST only takes less than 10 minutes.

  • [x] MNIST

    Original Images Reconstructed Images
    MNIST original images MNIST Reconstructed Images

    The result on MNIST test dataset. (K=20, D=64, latent space=3 by 3)

    I also observed its latent space by changing single value for each latent space from one of the observed latent code. The result is shown below. MNIST Latent Observation

    It seems that spatial location of latent code is improtant. By changing latent code on a specific location, the pixel matches with the location is disturbed.

    MNIST Latent Observation - Random Walk

    This results shows the 1000 generated images starting from knwon latent codes and changing aa single latent code at radnom location by +1 or -1. Most of the images are redundant (unrealistic), so it indicates that there are much room for compression.

    If you want to further explore the latent space, then try to play with notebook files I provided.

  • [x] CIFAR 10

    Original Images Reconstructed Images
    MNIST original images MNIST Reconstructed Images

    I was able to get 4.65 bits/dims. (K=10, D=256, latent space=8 by 8)


It will download required datasets on the directory ./datasets/{mnist,cifar10} by itself. Hence, just run the code will do the trick.

Run train

  • Run mnist: python
  • Run cifar10: python

Change the hyperparameters accordingly as you want. Please check at the bottom of each script.


I provide the model and the code for generating (,or reconstructing) images in the form of Jupyter notebook. Run jupyter notebook server, then run it to see more results with provided models.

If you want to test NLL, then run test() function on by uncomment the line. You can find it at the bottom of the file.


  • [ ] WaveNet?

Contributions are welcome!

Thoughts and Help request

  • The results seems correct, but there is a chance that the implmentation is not perfectly correct (especially, gradient copying...). If you find any glitches (or, a major defect) then, please let me know!
  • I am currently not sure how exactly NLL should be computed. Anyone who wants me a proper explantion on this?


Get A Weekly Email With Trending Projects For These Topics
No Spam. Unsubscribe easily at any time.
jupyter-notebook (6,156
tensorflow (2,138
generative-model (81
mnist (70
vae (69
cifar10 (39