080_convolutional_variational_autoencoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "cvae.ipynb",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ndo4ERqnwQOU"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xfNT-mlFwxVM"
      },
      "source": [
        "# Convolutional Variational Autoencoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ITZuApL56Mny"
      },
      "source": [
        "This notebook demonstrates how train a Variational Autoencoder (VAE) ([1](https://arxiv.org/abs/1312.6114), [2](https://arxiv.org/abs/1401.4082)). on the MNIST dataset. A VAE is a probabilistic take on the autoencoder, a model which takes high dimensional input data compress it into a smaller representation. Unlike a traditional autoencoder, which maps the input onto a latent vector, a VAE maps the input data into the parameters of a probability distribution, such as the mean and variance of a Gaussian. This approach produces a continuous, structured latent space, which is useful for image generation.\n",
        "\n",
        "![CVAE image latent space](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/generative/images/cvae_latent_space.jpg?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e1_Y75QXJS6h"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P-JuIu2N_SQf"
      },
      "source": [
        "!pip install tensorflow-probability\n",
        "\n",
        "# to generate gifs\n",
        "!pip install imageio\n",
        "!pip install git+https://github.com/tensorflow/docs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YfIk2es3hJEd"
      },
      "source": [
        "from IPython import display\n",
        "\n",
        "import glob\n",
        "import imageio\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import PIL\n",
        "import tensorflow as tf\n",
        "import tensorflow_probability as tfp\n",
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iYn4MdZnKCey"
      },
      "source": [
        "## Load the MNIST dataset\n",
        "Each MNIST image is originally a vector of 784 integers, each of which is between 0-255 and represents the intensity of a pixel. We model each pixel with a Bernoulli distribution in our model, and we statically binarize the dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a4fYMGxGhrna"
      },
      "source": [
        "(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NFC2ghIdiZYE"
      },
      "source": [
        "def preprocess_images(images):\n",
        "  images = images.reshape((images.shape[0], 28, 28, 1)) / 255.\n",
        "  return np.where(images > .5, 1.0, 0.0).astype('float32')\n",
        "\n",
        "train_images = preprocess_images(train_images)\n",
        "test_images = preprocess_images(test_images)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S4PIDhoDLbsZ"
      },
      "source": [
        "train_size = 60000\n",
        "batch_size = 32\n",
        "test_size = 10000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIGN6ouoQxt3"
      },
      "source": [
        "## Use *tf.data* to batch and shuffle the data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-yKCCQOoJ7cn"
      },
      "source": [
        "train_dataset = (tf.data.Dataset.from_tensor_slices(train_images)\n",
        "                 .shuffle(train_size).batch(batch_size))\n",
        "test_dataset = (tf.data.Dataset.from_tensor_slices(test_images)\n",
        "                .shuffle(test_size).batch(batch_size))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THY-sZMiQ4UV"
      },
      "source": [
        "## Define the encoder and decoder networks with *tf.keras.Sequential*\n",
        "\n",
        "In our VAE example, we use two small ConvNets for the encoder and decoder networks. In the literature, these networks are also referred to as inference/recognition and generative models respectively. We use `tf.keras.Sequential` to simplify implementation. Let $x$ and $z$ denote the observation and latent variable respectively in the following descriptions.\n",
        "\n",
        "### Encoder network\n",
        "This defines the approximate posterior distribution $q(z|x)$, which takes as input an observation and outputs a set of parameters for specifying the conditional distribution of the latent representation $z$. \n",
        "In this example, we simply model the distribution as a diagonal Gaussian, and the network outputs the mean and log-variance parameters of a factorized Gaussian. \n",
        "We output log-variance instead of the variance directly for numerical stability.\n",
        "\n",
        "### Decoder network \n",
        "This defines the conditional distribution of the observation $p(x|z)$, which takes a latent sample $z$ as input and outputs the parameters for a conditional distribution of the observation.\n",
        "We model the latent distribution prior $p(z)$ as a unit Gaussian.\n",
        "\n",
        "### Reparameterization trick\n",
        "To generate a sample $z$ for the decoder during training, we can sample from the latent distribution defined by the parameters outputted by the encoder, given an input observation $x$.\n",
        "However, this sampling operation creates a bottleneck because backpropagation cannot flow through a random node.\n",
        "\n",
        "To address this, we use a reparameterization trick.\n",
        "In our example, we approximate $z$ using the decoder parameters and another parameter $\\epsilon$ as follows:\n",
        "\n",
        "$$z = \\mu + \\sigma \\odot \\epsilon$$\n",
        "\n",
        "where $\\mu$ and $\\sigma$ represent the mean and standard deviation of a Gaussian distribution respectively. They can be derived from the decoder output. The $\\epsilon$ can be thought of as a random noise used to maintain stochasticity of $z$. We generate $\\epsilon$ from a standard normal distribution.\n",
        "\n",
        "The latent variable $z$ is now generated by a function of $\\mu$, $\\sigma$ and $\\epsilon$, which would enable the model to backpropagate gradients in the encoder through $\\mu$ and $\\sigma$ respectively, while maintaining stochasticity through $\\epsilon$.\n",
        "\n",
        "### Network architecture\n",
        "For the encoder network, we use two convolutional layers followed by a fully-connected layer. In the decoder network, we mirror this architecture by using a fully-connected layer followed by three convolution transpose layers (a.k.a. deconvolutional layers in some contexts). Note, it's common practice to avoid using batch normalization when training VAEs, since the additional stochasticity due to using mini-batches may aggravate instability on top of the stochasticity from sampling.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VGLbvBEmjK0a"
      },
      "source": [
        "class CVAE(tf.keras.Model):\n",
        "  \"\"\"Convolutional variational autoencoder.\"\"\"\n",
        "\n",
        "  def __init__(self, latent_dim):\n",
        "    super(CVAE, self).__init__()\n",
        "    self.latent_dim = latent_dim\n",
        "    self.encoder = tf.keras.Sequential(\n",
        "        [\n",
        "            tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),\n",
        "            tf.keras.layers.Conv2D(\n",
        "                filters=32, kernel_size=3, strides=(2, 2), activation='relu'),\n",
        "            tf.keras.layers.Conv2D(\n",
        "                filters=64, kernel_size=3, strides=(2, 2), activation='relu'),\n",
        "            tf.keras.layers.Flatten(),\n",
        "            # No activation\n",
        "            tf.keras.layers.Dense(latent_dim + latent_dim),\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    self.decoder = tf.keras.Sequential(\n",
        "        [\n",
        "            tf.keras.layers.InputLayer(input_shape=(latent_dim,)),\n",
        "            tf.keras.layers.Dense(units=7*7*32, activation=tf.nn.relu),\n",
        "            tf.keras.layers.Reshape(target_shape=(7, 7, 32)),\n",
        "            tf.keras.layers.Conv2DTranspose(\n",
        "                filters=64, kernel_size=3, strides=2, padding='same',\n",
        "                activation='relu'),\n",
        "            tf.keras.layers.Conv2DTranspose(\n",
        "                filters=32, kernel_size=3, strides=2, padding='same',\n",
        "                activation='relu'),\n",
        "            # No activation\n",
        "            tf.keras.layers.Conv2DTranspose(\n",
        "                filters=1, kernel_size=3, strides=1, padding='same'),\n",
        "        ]\n",
        "    )\n",
        "\n",
        "  @tf.function\n",
        "  def sample(self, eps=None):\n",
        "    if eps is None:\n",
        "      eps = tf.random.normal(shape=(100, self.latent_dim))\n",
        "    return self.decode(eps, apply_sigmoid=True)\n",
        "\n",
        "  def encode(self, x):\n",
        "    mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)\n",
        "    return mean, logvar\n",
        "\n",
        "  def reparameterize(self, mean, logvar):\n",
        "    eps = tf.random.normal(shape=mean.shape)\n",
        "    return eps * tf.exp(logvar * .5) + mean\n",
        "\n",
        "  def decode(self, z, apply_sigmoid=False):\n",
        "    logits = self.decoder(z)\n",
        "    if apply_sigmoid:\n",
        "      probs = tf.sigmoid(logits)\n",
        "      return probs\n",
        "    return logits"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0FMYgY_mPfTi"
      },
      "source": [
        "## Define the loss function and the optimizer\n",
        "\n",
        "VAEs train by maximizing the evidence lower bound (ELBO) on the marginal log-likelihood:\n",
        "\n",
        "$$\\log p(x) \\ge \\text{ELBO} = \\mathbb{E}_{q(z|x)}\\left[\\log \\frac{p(x, z)}{q(z|x)}\\right].$$\n",
        "\n",
        "In practice, we optimize the single sample Monte Carlo estimate of this expectation:\n",
        "\n",
        "$$\\log p(x| z) + \\log p(z) - \\log q(z|x),$$\n",
        "where $z$ is sampled from $q(z|x)$.\n",
        "\n",
        "**Note**: we could also analytically compute the KL term, but here we incorporate all three terms in the Monte Carlo estimator for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iWCn_PVdEJZ7"
      },
      "source": [
        "optimizer = tf.keras.optimizers.Adam(1e-4)\n",
        "\n",
        "\n",
        "def log_normal_pdf(sample, mean, logvar, raxis=1):\n",
        "  log2pi = tf.math.log(2. * np.pi)\n",
        "  return tf.reduce_sum(\n",
        "      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),\n",
        "      axis=raxis)\n",
        "\n",
        "\n",
        "def compute_loss(model, x):\n",
        "  mean, logvar = model.encode(x)\n",
        "  z = model.reparameterize(mean, logvar)\n",
        "  x_logit = model.decode(z)\n",
        "  cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)\n",
        "  logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])\n",
        "  logpz = log_normal_pdf(z, 0., 0.)\n",
        "  logqz_x = log_normal_pdf(z, mean, logvar)\n",
        "  return -tf.reduce_mean(logpx_z + logpz - logqz_x)\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def train_step(model, x, optimizer):\n",
        "  \"\"\"Executes one training step and returns the loss.\n",
        "\n",
        "  This function computes the loss and gradients, and uses the latter to\n",
        "  update the model's parameters.\n",
        "  \"\"\"\n",
        "  with tf.GradientTape() as tape:\n",
        "    loss = compute_loss(model, x)\n",
        "  gradients = tape.gradient(loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rw1fkAczTQYh"
      },
      "source": [
        "## Training\n",
        "\n",
        "* We start by iterating over the dataset\n",
        "* During each iteration, we pass the image to the encoder to obtain a set of mean and log-variance parameters of the approximate posterior $q(z|x)$\n",
        "* We then apply the *reparameterization trick* to sample from $q(z|x)$\n",
        "* Finally, we pass the reparameterized samples to the decoder to obtain the logits of the generative distribution $p(x|z)$\n",
        "* **Note:** Since we use the dataset loaded by keras with 60k datapoints in the training set and 10k datapoints in the test set, our resulting ELBO on the test set is slightly higher than reported results in the literature which uses dynamic binarization of Larochelle's MNIST.\n",
        "\n",
        "### Generating images\n",
        "\n",
        "* After training, it is time to generate some images\n",
        "* We start by sampling a set of latent vectors from the unit Gaussian prior distribution $p(z)$\n",
        "* The generator will then convert the latent sample $z$ to logits of the observation, giving a distribution $p(x|z)$\n",
        "* Here we plot the probabilities of Bernoulli distributions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NS2GWywBbAWo"
      },
      "source": [
        "epochs = 10\n",
        "# set the dimensionality of the latent space to a plane for visualization later\n",
        "latent_dim = 2\n",
        "num_examples_to_generate = 16\n",
        "\n",
        "# keeping the random vector constant for generation (prediction) so\n",
        "# it will be easier to see the improvement.\n",
        "random_vector_for_generation = tf.random.normal(\n",
        "    shape=[num_examples_to_generate, latent_dim])\n",
        "model = CVAE(latent_dim)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RmdVsmvhPxyy"
      },
      "source": [
        "def generate_and_save_images(model, epoch, test_sample):\n",
        "  mean, logvar = model.encode(test_sample)\n",
        "  z = model.reparameterize(mean, logvar)\n",
        "  predictions = model.sample(z)\n",
        "  fig = plt.figure(figsize=(4, 4))\n",
        "\n",
        "  for i in range(predictions.shape[0]):\n",
        "    plt.subplot(4, 4, i + 1)\n",
        "    plt.imshow(predictions[i, :, :, 0], cmap='gray')\n",
        "    plt.axis('off')\n",
        "\n",
        "  # tight_layout minimizes the overlap between 2 sub-plots\n",
        "  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "swCyrbqQQ-Ri"
      },
      "source": [
        "# Pick a sample of the test set for generating output images\n",
        "assert batch_size >= num_examples_to_generate\n",
        "for test_batch in test_dataset.take(1):\n",
        "  test_sample = test_batch[0:num_examples_to_generate, :, :, :]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2M7LmLtGEMQJ"
      },
      "source": [
        "generate_and_save_images(model, 0, test_sample)\n",
        "\n",
        "for epoch in range(1, epochs + 1):\n",
        "  start_time = time.time()\n",
        "  for train_x in train_dataset:\n",
        "    train_step(model, train_x, optimizer)\n",
        "  end_time = time.time()\n",
        "\n",
        "  loss = tf.keras.metrics.Mean()\n",
        "  for test_x in test_dataset:\n",
        "    loss(compute_loss(model, test_x))\n",
        "  elbo = -loss.result()\n",
        "  display.clear_output(wait=False)\n",
        "  print('Epoch: {}, Test set ELBO: {}, time elapse for current epoch: {}'\n",
        "        .format(epoch, elbo, end_time - start_time))\n",
        "  generate_and_save_images(model, epoch, test_sample)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P4M_vIbUi7c0"
      },
      "source": [
        "### Display a generated image from the last training epoch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WfO5wCdclHGL"
      },
      "source": [
        "def display_image(epoch_no):\n",
        "  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5x3q9_Oe5q0A"
      },
      "source": [
        "plt.imshow(display_image(epoch))\n",
        "plt.axis('off')  # Display images"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NywiH3nL8guF"
      },
      "source": [
        "### Display an animated GIF of all the saved images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IGKQgENQ8lEI"
      },
      "source": [
        "anim_file = 'cvae.gif'\n",
        "\n",
        "with imageio.get_writer(anim_file, mode='I') as writer:\n",
        "  filenames = glob.glob('image*.png')\n",
        "  filenames = sorted(filenames)\n",
        "  for filename in filenames:\n",
        "    image = imageio.imread(filename)\n",
        "    writer.append_data(image)\n",
        "  image = imageio.imread(filename)\n",
        "  writer.append_data(image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ZqAEtdqUmJF"
      },
      "source": [
        "import tensorflow_docs.vis.embed as embed\n",
        "embed.embed_file(anim_file)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PeunRU6TSumT"
      },
      "source": [
        "### Display a 2D manifold of digits from the latent space\n",
        "\n",
        "Running the code below will show a continuous distribution of the different digit classes, with each digit morphing into another across the 2D latent space. We use [TensorFlow Probability](https://www.tensorflow.org/probability) to generate a standard normal distribution for the latent space."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "code",
        "id": "mNcaaYPBS3mj"
      },
      "source": [
        "def plot_latent_images(model, n, digit_size=28):\n",
        "  \"\"\"Plots n x n digit images decoded from the latent space.\"\"\"\n",
        "\n",
        "  norm = tfp.distributions.Normal(0, 1)\n",
        "  grid_x = norm.quantile(np.linspace(0.05, 0.95, n))\n",
        "  grid_y = norm.quantile(np.linspace(0.05, 0.95, n))\n",
        "  image_width = digit_size*n\n",
        "  image_height = image_width\n",
        "  image = np.zeros((image_height, image_width))\n",
        "\n",
        "  for i, yi in enumerate(grid_x):\n",
        "    for j, xi in enumerate(grid_y):\n",
        "      z = np.array([[xi, yi]])\n",
        "      x_decoded = model.sample(z)\n",
        "      digit = tf.reshape(x_decoded[0], (digit_size, digit_size))\n",
        "      image[i * digit_size: (i + 1) * digit_size,\n",
        "            j * digit_size: (j + 1) * digit_size] = digit.numpy()\n",
        "\n",
        "  plt.figure(figsize=(10, 10))\n",
        "  plt.imshow(image, cmap='Greys_r')\n",
        "  plt.axis('Off')\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F-ZG69QCZnGY"
      },
      "source": [
        "plot_latent_images(model, 20)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrJRef8Ln945"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "This tutorial has demonstrated how to implement a convolutional variational autoencoder using TensorFlow. \n",
        "\n",
        "As a next step, you could try to improve the model output by increasing the network size. \n",
        "For instance, you could try setting the `filter` parameters for each of the `Conv2D` and `Conv2DTranspose` layers to 512. \n",
        "Note that in order to generate the final 2D latent image plot, you would need to keep `latent_dim` to 2. Also, the training time would increase as the network size increases.\n",
        "\n",
        "You could also try implementing a VAE using a different dataset, such as CIFAR-10.\n",
        "\n",
        "VAEs can be implemented in several different styles and of varying complexity. You can find additional implementations in the following sources:\n",
        "- [Variational AutoEncoder (keras.io)](https://keras.io/examples/generative/vae/)\n",
        "- [VAE example from \"Writing custom layers and models\" guide (tensorflow.org)](https://www.tensorflow.org/guide/keras/custom_layers_and_models#putting_it_all_together_an_end-to-end_example)\n",
        "- [TFP Probabilistic Layers: Variational Auto Encoder](https://www.tensorflow.org/probability/examples/Probabilistic_Layers_VAE)\n",
        "\n",
        "If you'd like to learn more about the details of VAEs, please refer to [An Introduction to Variational Autoencoders](https://arxiv.org/abs/1906.02691)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v8-Deke3uMq5"
      },
      "source": [
        "**Contact** : [FarrokhKarimi](https://zil.ink/farrokhkarimi)  \n",
        "**Reference** : [tensorflow.com](https://www.tensorflow.org/)  \n",
        "Special thanks to [François Chollet](https://fchollet.com/)"
      ]
    }
  ]
}