C4_W3_Lab_1_VAE_MNIST.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "C4_W3_Lab_1_VAE_MNIST.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I2m-t1vpjAMj"
      },
      "source": [
        "# Ungraded Lab: Variational Autoencoders\n",
        "\n",
        "This lab will demonstrate all the concepts you learned this week. You will build a Variational Autoencoder (VAE) trained on the MNIST dataset and see how it is able to generate new images. This will be very useful for this week's assignment. Let's begin!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9dE3sNy_jKus"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "17DD2aRgudaO"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython import display"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGHahainjOji"
      },
      "source": [
        "## Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DBoDTLNXuFqT"
      },
      "source": [
        "# Define global constants to be used in this notebook\n",
        "BATCH_SIZE=128\n",
        "LATENT_DIM=2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VqZ-LiQbjaNX"
      },
      "source": [
        "## Prepare the Dataset\n",
        "\n",
        "You will just be using the `train` split of the MNIST dataset in this notebook. We've prepared a few helper functions below to help in downloading and preparing the dataset:\n",
        "\n",
        "* `map_image()` - normalizes and creates a tensor from the image, returning only the image. This will be used for the unsupervised learning in the autoencoder.\n",
        "\n",
        "* `get_dataset()` - loads MNIST from Tensorflow Datasets, fetching the `train` split by default, then prepares it using the mapping function. If `is_validation` is set to `True`, then it will get the `test` split instead. Training sets will also be shuffled."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iXgPMPNbteYU"
      },
      "source": [
        "def map_image(image, label):\n",
        "  '''returns a normalized and reshaped tensor from a given image'''\n",
        "  image = tf.cast(image, dtype=tf.float32)\n",
        "  image = image / 255.0\n",
        "  image = tf.reshape(image, shape=(28, 28, 1,))\n",
        "  \n",
        "  return image\n",
        "\n",
        "\n",
        "def get_dataset(map_fn, is_validation=False):\n",
        "  '''Loads and prepares the mnist dataset from TFDS.'''\n",
        "  if is_validation:\n",
        "    split_name = \"test\"\n",
        "  else:\n",
        "    split_name = \"train\"\n",
        "\n",
        "  dataset = tfds.load('mnist', as_supervised=True, split=split_name)\n",
        "  dataset = dataset.map(map_fn)\n",
        "  \n",
        "  if is_validation:\n",
        "    dataset = dataset.batch(BATCH_SIZE)\n",
        "  else:\n",
        "    dataset = dataset.shuffle(1024).batch(BATCH_SIZE)\n",
        "\n",
        "  return dataset"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ttCP6xrJGxY5"
      },
      "source": [
        "Please run this cell to download and prepare the `train` split of the MNIST dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jszTpjHVuJXO"
      },
      "source": [
        "train_dataset = get_dataset(map_image)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qedUCLa_jfeM"
      },
      "source": [
        "## Build the Model\n",
        "\n",
        "You will now be building your VAE model. The main parts are shown in the figure below:\n",
        "\n",
        "<img src=\"https://drive.google.com/uc?export=view&id=1YAZAeMGEJ1KgieYk1ju-S9DoshpMREeC\" width=\"60%\" height=\"60%\"/>\n",
        "\n",
        "Like the autoencoder last week, the VAE also has an encoder-decoder architecture with the main difference being the grey box in the middle which stands for the latent representation. In this layer, the model mixes a random sample and combines it with the outputs of the encoder. This mechanism makes it useful for generating new content. Let's build these parts one-by-one in the next sections."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaj1dg-FrWuG"
      },
      "source": [
        "### Sampling Class\n",
        "\n",
        "First, you will build the `Sampling` class. This will be a custom Keras layer that will provide the Gaussian noise input along with the mean (mu) and standard deviation (sigma) of the encoder's output. In practice, the output of this layer is given by the equation:\n",
        "\n",
        "$$z = \\mu + e^{0.5\\sigma} * \\epsilon  $$\n",
        "\n",
        "where $\\mu$ = mean, $\\sigma$ = standard deviation, and $\\epsilon$ = random sample"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ppLApb2VuzKZ"
      },
      "source": [
        "class Sampling(tf.keras.layers.Layer):\n",
        "  def call(self, inputs):\n",
        "    \"\"\"Generates a random sample and combines with the encoder output\n",
        "    \n",
        "    Args:\n",
        "      inputs -- output tensor from the encoder\n",
        "\n",
        "    Returns:\n",
        "      `inputs` tensors combined with a random sample\n",
        "    \"\"\"\n",
        "\n",
        "    # unpack the output of the encoder\n",
        "    mu, sigma = inputs\n",
        "\n",
        "    # get the size and dimensions of the batch\n",
        "    batch = tf.shape(mu)[0]\n",
        "    dim = tf.shape(mu)[1]\n",
        "\n",
        "    # generate a random tensor\n",
        "    epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
        "\n",
        "    # combine the inputs and noise\n",
        "    return mu + tf.exp(0.5 * sigma) * epsilon"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HCqWbPNvrplb"
      },
      "source": [
        "### Encoder\n",
        "\n",
        "Next, you will build the encoder part of the network. You will follow the architecture shown in class which looks like this. Note that aside from mu and sigma, you will also output the shape of features before flattening it. This will be useful when reconstructing the image later in the decoder.\n",
        "\n",
        "*Note:* You might encounter issues with using batch normalization with smaller batches, and sometimes the advice is given to avoid using batch normalization when training VAEs in particular. Feel free to experiment with adding or removing it from this notebook to explore the effects.\n",
        "\n",
        "<img src=\"https://drive.google.com/uc?export=view&id=1eoxFK_UVSHd3a_5EHcCU8F8QDZlPiXfW\" width=\"60%\" height=\"60%\"/>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nU5kZsj0u9jX"
      },
      "source": [
        "def encoder_layers(inputs, latent_dim):\n",
        "  \"\"\"Defines the encoder's layers.\n",
        "  Args:\n",
        "    inputs -- batch from the dataset\n",
        "    latent_dim -- dimensionality of the latent space\n",
        "\n",
        "  Returns:\n",
        "    mu -- learned mean\n",
        "    sigma -- learned standard deviation\n",
        "    batch_2.shape -- shape of the features before flattening\n",
        "  \"\"\"\n",
        "\n",
        "  # add the Conv2D layers followed by BatchNormalization\n",
        "  x = tf.keras.layers.Conv2D(filters=32, kernel_size=3, strides=2, padding=\"same\", activation='relu', name=\"encode_conv1\")(inputs)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.Conv2D(filters=64, kernel_size=3, strides=2, padding='same', activation='relu', name=\"encode_conv2\")(x)\n",
        "\n",
        "  # assign to a different variable so you can extract the shape later\n",
        "  batch_2 = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "  # flatten the features and feed into the Dense network\n",
        "  x = tf.keras.layers.Flatten(name=\"encode_flatten\")(batch_2)\n",
        "\n",
        "  # we arbitrarily used 20 units here but feel free to change and see what results you get\n",
        "  x = tf.keras.layers.Dense(20, activation='relu', name=\"encode_dense\")(x)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "\n",
        "  # add output Dense networks for mu and sigma, units equal to the declared latent_dim.\n",
        "  mu = tf.keras.layers.Dense(latent_dim, name='latent_mu')(x)\n",
        "  sigma = tf.keras.layers.Dense(latent_dim, name ='latent_sigma')(x)\n",
        "\n",
        "  return mu, sigma, batch_2.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gFiOzFnUnPMN"
      },
      "source": [
        "With the encoder layers defined, you can declare the encoder model that includes the `Sampling` layer with the function below:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CoLLpfBUvhBm"
      },
      "source": [
        "def encoder_model(latent_dim, input_shape):\n",
        "  \"\"\"Defines the encoder model with the Sampling layer\n",
        "  Args:\n",
        "    latent_dim -- dimensionality of the latent space\n",
        "    input_shape -- shape of the dataset batch\n",
        "\n",
        "  Returns:\n",
        "    model -- the encoder model\n",
        "    conv_shape -- shape of the features before flattening\n",
        "  \"\"\"\n",
        "\n",
        "  # declare the inputs tensor with the given shape\n",
        "  inputs = tf.keras.layers.Input(shape=input_shape)\n",
        "\n",
        "  # get the output of the encoder_layers() function\n",
        "  mu, sigma, conv_shape = encoder_layers(inputs, latent_dim=LATENT_DIM)\n",
        "\n",
        "  # feed mu and sigma to the Sampling layer\n",
        "  z = Sampling()((mu, sigma))\n",
        "\n",
        "  # build the whole encoder model\n",
        "  model = tf.keras.Model(inputs, outputs=[mu, sigma, z])\n",
        "\n",
        "  return model, conv_shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NkFgN22trttX"
      },
      "source": [
        "### Decoder\n",
        "\n",
        "Next, you will build the decoder part of the network which expands the latent representations back to the original image dimensions. As you'll see later in the training loop, you can feed random inputs to this model and it will generate content that resemble the training data."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6H_HoaAYvWZn"
      },
      "source": [
        "def decoder_layers(inputs, conv_shape):\n",
        "  \"\"\"Defines the decoder layers.\n",
        "  Args:\n",
        "    inputs -- output of the encoder \n",
        "    conv_shape -- shape of the features before flattening\n",
        "\n",
        "  Returns:\n",
        "    tensor containing the decoded output\n",
        "  \"\"\"\n",
        "\n",
        "  # feed to a Dense network with units computed from the conv_shape dimensions\n",
        "  units = conv_shape[1] * conv_shape[2] * conv_shape[3]\n",
        "  x = tf.keras.layers.Dense(units, activation = 'relu', name=\"decode_dense1\")(inputs)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  \n",
        "  # reshape output using the conv_shape dimensions\n",
        "  x = tf.keras.layers.Reshape((conv_shape[1], conv_shape[2], conv_shape[3]), name=\"decode_reshape\")(x)\n",
        "\n",
        "  # upsample the features back to the original dimensions\n",
        "  x = tf.keras.layers.Conv2DTranspose(filters=64, kernel_size=3, strides=2, padding='same', activation='relu', name=\"decode_conv2d_2\")(x)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.Conv2DTranspose(filters=32, kernel_size=3, strides=2, padding='same', activation='relu', name=\"decode_conv2d_3\")(x)\n",
        "  x = tf.keras.layers.BatchNormalization()(x)\n",
        "  x = tf.keras.layers.Conv2DTranspose(filters=1, kernel_size=3, strides=1, padding='same', activation='sigmoid', name=\"decode_final\")(x)\n",
        "  \n",
        "  return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VX2hjxYhxQyn"
      },
      "source": [
        "You can define the decoder model as shown below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VGZ5kqA4vuEy"
      },
      "source": [
        "def decoder_model(latent_dim, conv_shape):\n",
        "  \"\"\"Defines the decoder model.\n",
        "  Args:\n",
        "    latent_dim -- dimensionality of the latent space\n",
        "    conv_shape -- shape of the features before flattening\n",
        "\n",
        "  Returns:\n",
        "    model -- the decoder model\n",
        "  \"\"\"\n",
        "\n",
        "  # set the inputs to the shape of the latent space\n",
        "  inputs = tf.keras.layers.Input(shape=(latent_dim,))\n",
        "\n",
        "  # get the output of the decoder layers\n",
        "  outputs = decoder_layers(inputs, conv_shape)\n",
        "\n",
        "  # declare the inputs and outputs of the model\n",
        "  model = tf.keras.Model(inputs, outputs)\n",
        "  \n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQbtaVsHrxQ_"
      },
      "source": [
        "### Kullback–Leibler Divergence"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqG6oUC3r6Um"
      },
      "source": [
        "To improve the generative capability of the model, you have to take into account the random normal distribution introduced in the latent space. For that, the [Kullback–Leibler Divergence](https://arxiv.org/abs/2002.07514) is computed and added to the reconstruction loss. The formula is defined in the function below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "14Gla4K6vcLN"
      },
      "source": [
        "def kl_reconstruction_loss(inputs, outputs, mu, sigma):\n",
        "  \"\"\" Computes the Kullback-Leibler Divergence (KLD)\n",
        "  Args:\n",
        "    inputs -- batch from the dataset\n",
        "    outputs -- output of the Sampling layer\n",
        "    mu -- mean\n",
        "    sigma -- standard deviation\n",
        "\n",
        "  Returns:\n",
        "    KLD loss\n",
        "  \"\"\"\n",
        "  kl_loss = 1 + sigma - tf.square(mu) - tf.math.exp(sigma)\n",
        "  kl_loss = tf.reduce_mean(kl_loss) * -0.5\n",
        "\n",
        "  return kl_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IiAwutTjr6aQ"
      },
      "source": [
        "### VAE Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ymQdQTOJvOTR"
      },
      "source": [
        "You can now define the entire VAE model. Note the use of `model.add_loss()` to add the KL reconstruction loss. Computing this loss doesn't use `y_true` and `y_pred` so it can't be used in `model.compile()`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3hkx7OCqvzlb"
      },
      "source": [
        "def vae_model(encoder, decoder, input_shape):\n",
        "  \"\"\"Defines the VAE model\n",
        "  Args:\n",
        "    encoder -- the encoder model\n",
        "    decoder -- the decoder model\n",
        "    input_shape -- shape of the dataset batch\n",
        "\n",
        "  Returns:\n",
        "    the complete VAE model\n",
        "  \"\"\"\n",
        "\n",
        "  # set the inputs\n",
        "  inputs = tf.keras.layers.Input(shape=input_shape)\n",
        "\n",
        "  # get mu, sigma, and z from the encoder output\n",
        "  mu, sigma, z = encoder(inputs)\n",
        "  \n",
        "  # get reconstructed output from the decoder\n",
        "  reconstructed = decoder(z)\n",
        "\n",
        "  # define the inputs and outputs of the VAE\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=reconstructed)\n",
        "\n",
        "  # add the KL loss\n",
        "  loss = kl_reconstruction_loss(inputs, z, mu, sigma)\n",
        "  model.add_loss(loss)\n",
        "\n",
        "  return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S5FxUuopxa_I"
      },
      "source": [
        "We'll add a helper function to setup and get the different models from the functions you defined."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "piqZLzkHv3jw"
      },
      "source": [
        "def get_models(input_shape, latent_dim):\n",
        "  \"\"\"Returns the encoder, decoder, and vae models\"\"\"\n",
        "  encoder, conv_shape = encoder_model(latent_dim=latent_dim, input_shape=input_shape)\n",
        "  decoder = decoder_model(latent_dim=latent_dim, conv_shape=conv_shape)\n",
        "  vae = vae_model(encoder, decoder, input_shape=input_shape)\n",
        "  return encoder, decoder, vae"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hOp-yWg2v7uP"
      },
      "source": [
        "# Get the encoder, decoder and 'master' model (called vae)\n",
        "encoder, decoder, vae = get_models(input_shape=(28,28,1,), latent_dim=LATENT_DIM)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iLMU6YySmWKR"
      },
      "source": [
        "## Train the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SHVaw_kqyPQI"
      },
      "source": [
        "You can now setup the VAE model for training. Let's start by defining the reconstruction loss, optimizer and metric."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VMWqvQqvwEMK"
      },
      "source": [
        "# Define our loss functions and optimizers\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "loss_metric = tf.keras.metrics.Mean()\n",
        "bce_loss = tf.keras.losses.BinaryCrossentropy()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MpmZiTYQzIVH"
      },
      "source": [
        "You will want to see the progress of the image generation at each epoch. For that, you can use the helper function below. This will generate 16 images in a 4x4 grid."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gaG0h17cwUYM"
      },
      "source": [
        "def generate_and_save_images(model, epoch, step, test_input):\n",
        "  \"\"\"Helper function to plot our 16 images\n",
        "\n",
        "  Args:\n",
        "\n",
        "  model -- the decoder model\n",
        "  epoch -- current epoch number during training\n",
        "  step -- current step number during training\n",
        "  test_input -- random tensor with shape (16, LATENT_DIM)\n",
        "  \"\"\"\n",
        "\n",
        "  # generate images from the test input\n",
        "  predictions = model.predict(test_input)\n",
        "\n",
        "  # plot the results\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",
        "  fig.suptitle(\"epoch: {}, step: {}\".format(epoch, step))\n",
        "  plt.savefig('image_at_epoch_{:04d}_step{:04d}.png'.format(epoch, step))\n",
        "  plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VeArnhVI0HQx"
      },
      "source": [
        "The training loop is shown below. This will display generated images each epoch and will take around 30 minutes to complete. Notice too that we add the KLD loss to the binary crossentropy loss before we get the gradients and update the weights.\n",
        "\n",
        "As you might expect, the initial 16 images will look random but it will improve overtime as the network learns and you'll see images that resemble the MNIST dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y8o4ZPU1wSFc"
      },
      "source": [
        "# Training loop. \n",
        "\n",
        "# generate random vector as test input to the decoder\n",
        "random_vector_for_generation = tf.random.normal(shape=[16, LATENT_DIM])\n",
        "\n",
        "# number of epochs\n",
        "epochs = 100\n",
        "\n",
        "# initialize the helper function to display outputs from an untrained model\n",
        "generate_and_save_images(decoder, 0, 0, random_vector_for_generation)\n",
        "\n",
        "for epoch in range(epochs):\n",
        "  print('Start of epoch %d' % (epoch,))\n",
        "\n",
        "  # iterate over the batches of the dataset.\n",
        "  for step, x_batch_train in enumerate(train_dataset):\n",
        "    with tf.GradientTape() as tape:\n",
        "\n",
        "      # feed a batch to the VAE model\n",
        "      reconstructed = vae(x_batch_train)\n",
        "\n",
        "      # compute reconstruction loss\n",
        "      flattened_inputs = tf.reshape(x_batch_train, shape=[-1])\n",
        "      flattened_outputs = tf.reshape(reconstructed, shape=[-1])\n",
        "      loss = bce_loss(flattened_inputs, flattened_outputs) * 784\n",
        "      \n",
        "      # add KLD regularization loss\n",
        "      loss += sum(vae.losses)  \n",
        "\n",
        "    # get the gradients and update the weights\n",
        "    grads = tape.gradient(loss, vae.trainable_weights)\n",
        "    optimizer.apply_gradients(zip(grads, vae.trainable_weights))\n",
        "\n",
        "    # compute the loss metric\n",
        "    loss_metric(loss)\n",
        "\n",
        "    # display outputs every 100 steps\n",
        "    if step % 100 == 0:\n",
        "      display.clear_output(wait=False)    \n",
        "      generate_and_save_images(decoder, epoch, step, random_vector_for_generation)\n",
        "      print('Epoch: %s step: %s mean loss = %s' % (epoch, step, loss_metric.result().numpy()))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CrQB8Lr919hn"
      },
      "source": [
        "**Congratulations on completing this lab on Variational Autoencoders!**"
      ]
    }
  ]
}