Lab_9_Generative_Models.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Lab_9_Generative_Models.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "REaV1GblZovs"
      },
      "source": [
        "#Generative models\n",
        "In this Lab we will be experimenting with\n",
        "\n",
        "- Convolutional autoencoders\n",
        "- Latent space visualization and interpolation\n",
        "- Upsampling techniques\n",
        "- Variational autoencoders\n",
        "- Deep Convolutional GANs (DCGANs)\n",
        "\n",
        "If you want to experiment with Denoising Autoencoders, revisit Lab 3 (task 5):\n",
        "https://github.com/aivclab/dlcourse/blob/master/Lab3_FunWithMNIST.ipynb\n",
        "\n",
        "**Before we start - remember to set runtime to GPU**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VSIpfu6wUG1r"
      },
      "source": [
        "**NOTE:** In case you have trouble running Keras/TensorFlow in Colab, try one of the following:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fZ01fAojUGCV"
      },
      "source": [
        "# Try this\n",
        "#!pip install --upgrade tensorflow==1.8.0\n",
        "\n",
        "# ... or this\n",
        "#%tensorflow_version 1.x\n",
        "\n",
        "# Check TensorFlow version\n",
        "#import tensorflow as tf\n",
        "#print(tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qc8RJkIJaxap"
      },
      "source": [
        "##1. Download the MNIST dataset\n",
        "As usual:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r-_9ZwQkiGQg"
      },
      "source": [
        "from __future__ import print_function\n",
        "import keras\n",
        "from keras.datasets import mnist\n",
        "from keras import backend as K\n",
        "\n",
        "num_classes = 10\n",
        "\n",
        "# input image dimensions\n",
        "img_rows, img_cols = 28, 28\n",
        "\n",
        "# the data, split between train and test sets\n",
        "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
        "\n",
        "if K.image_data_format() == 'channels_first':\n",
        "    x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)\n",
        "    x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)\n",
        "    input_shape = (1, img_rows, img_cols)\n",
        "else:\n",
        "    x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)\n",
        "    x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)\n",
        "    input_shape = (img_rows, img_cols, 1)\n",
        "\n",
        "# Pre-process inputs\n",
        "x_train = x_train.astype('float32')\n",
        "x_test = x_test.astype('float32')\n",
        "x_train /= 255\n",
        "x_test /= 255\n",
        "\n",
        "# Convert class indices to one-hot vectors\n",
        "y_train = keras.utils.to_categorical(y_train, num_classes)\n",
        "y_test = keras.utils.to_categorical(y_test, num_classes)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9Son2qjrbQTq",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4be322fb-25c4-4f8c-da54-9d146fa06d3c"
      },
      "source": [
        "# Input shape: 28 x 28 x 1 = image with one color channel\n",
        "print('input_shape :',input_shape)\n",
        "\n",
        "# Pre-process inputs\n",
        "print('x_train shape:', x_train.shape)\n",
        "print(x_train.shape[0], 'train samples')\n",
        "print(x_test.shape[0], 'test samples')\n",
        "\n",
        "# to_categorical converts class indices to one-hot vectors\n",
        "print('y_train shape:', y_train.shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "input_shape : (28, 28, 1)\n",
            "x_train shape: (60000, 28, 28, 1)\n",
            "60000 train samples\n",
            "10000 test samples\n",
            "y_train shape: (60000, 10)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6YJnv8WI7b4B"
      },
      "source": [
        "##2. Task 1: Convolutional Autoencoder\n",
        "Recall that *generative models* try to learn something useful about the distribution of the data. A succesful generative model is one that allows us to draw realistically looking samples from the original distribution.\n",
        "\n",
        "In today's lab, we will be using MNIST, but we will assume that we don't know the labels and see if we can detect interesting structures in the data - without using the labels. Hence, all the generative models below represent examples of *unsupervised learning*.\n",
        "\n",
        "The simplest type of generative model is an *autoencoder* (AEs). Traditional AEs are fully connected (i.e., inputs and outputs are vectors), which we know is not very useful when dealing with image data. Therefore, we will be using the convolutional variant, called a Convolutional Autoencoder (CAE).\n",
        "\n",
        "There are many ways to implement CAEs. The one below is designed to map the input data down to a 2D *latent space*, so that you can plot the latent vectors in 2D. In this way we can use to CAE to visualize the distribution of our data in 2D.\n",
        "\n",
        "**Note** that below we define the encoder and the decoder separately and combine them afterwards to form the final CAE model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aABSvT1cME5"
      },
      "source": [
        "###2.1 Your task\n",
        "An autoencoder learns an identity function, meaning that we require the input and output of the model to have exactly the same shape.\n",
        "\n",
        "Your task is to specify the kernel size and padding size of the last layer, such that shape of the decoder's output matches the input shape (28x28x1). That is, fill in the missing code marked with `???`:\n",
        "\n",
        "```\n",
        "decoded = Conv2D(1, kernel_size=???, padding=???, activation='sigmoid')(x)\n",
        "```\n",
        "\n",
        "Recall that in Keras, padding can be set to either `same` or `valid` (what's the difference?)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIBw4LuzVEEh"
      },
      "source": [
        "from keras.layers import Input, Dropout, Flatten, Dense, Conv2D, MaxPooling2D\n",
        "from keras.models import Model\n",
        "from keras.layers import UpSampling2D, ZeroPadding2D, Conv2DTranspose, Reshape\n",
        "\n",
        "# Number of latent dimensions\n",
        "latent_dim = 2\n",
        "\n",
        "# Encoder (convolutional base)\n",
        "inputs = Input(shape=(28, 28, 1))\n",
        "x = ZeroPadding2D(padding=(2, 2))(inputs)\n",
        "x = Conv2D(8, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(16, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(32, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "\n",
        "# shape info needed to build decoder model\n",
        "shape = K.int_shape(x)\n",
        "\n",
        "x = Flatten()(x) # vectorize last feature map\n",
        "encoded = Dense(latent_dim)(x) # map to vector of length \"latent_dim\"\n",
        "encoder = Model(inputs, encoded)\n",
        "encoder.summary()\n",
        "print((\"shape of encoded\", K.int_shape(encoded)))\n",
        "\n",
        "# Decoder (upsamling)\n",
        "encoding = Input(shape=(1, 1, latent_dim))\n",
        "x = Dense(shape[1] * shape[2] * shape[3], activation='relu')(encoding)\n",
        "x = Reshape((shape[1], shape[2], shape[3]))(x)\n",
        "x = Conv2DTranspose(32, (3,3), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(16, (3,3), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(8, (3,3), strides=(2,2), padding='same')(x)\n",
        "decoded = Conv2D(1, kernel_size=???, padding=???, activation='sigmoid')(x) # Fix this line !!!\n",
        "decoder = Model(encoding, decoded)\n",
        "decoder.summary()\n",
        "print((\"shape of decoded\", K.int_shape(decoded)))\n",
        "\n",
        "x = encoder(inputs)\n",
        "predictions = decoder(x)\n",
        "autoencoder = Model(inputs, predictions)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1gIKS_zSWZng"
      },
      "source": [
        "###2.2 Questions:\n",
        "1. What does the `ZeroPadding2D` layer do?\n",
        "2. What is the shape of the data before and after zero padding? (Note: for downsampling and upsampling it is more convenient if the shape of the data is a power of 2).\n",
        "3. What is the purpose of the `Reshape`layer in the decoder?\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bKXQlf8tBnmc"
      },
      "source": [
        "###2.3 Training\n",
        "Let's train the autoencoder for 30 epochs (add more epochs to improve results):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vw_OqAHD5k1-"
      },
      "source": [
        "autoencoder.compile(optimizer='adam', loss='mse')\n",
        "history = autoencoder.fit(x_train, x_train, epochs=30, batch_size=256,\n",
        "               shuffle=True, validation_data=(x_test, x_test), verbose=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q_91-fqbZSPp"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(history.history['loss'])\n",
        "plt.plot(history.history['val_loss'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gcd5EgZuZI1l"
      },
      "source": [
        "###2.4 Plot the latent space representation\n",
        "To get some intuition about what our autoencoder has learned, we can plot the latent representation of the training data:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G7Idi6hcEhjx"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import matplotlib as mpl\n",
        "mpl.rc('image', cmap='jet')\n",
        "\n",
        "# Get latent representation\n",
        "z = encoder.predict(x_train,batch_size=32)\n",
        "\n",
        "# Plot\n",
        "plt.figure(figsize=(12, 10))\n",
        "plt.scatter(z[:, 0], z[:, 1], c=np.argmax(y_train,axis=1))\n",
        "plt.colorbar()\n",
        "plt.xlabel(\"z[0]\")\n",
        "plt.ylabel(\"z[1]\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sujyFZhV7YCz"
      },
      "source": [
        "Each color represents a different class. What do you see from this plot? Which classes are clearly separated, and which are mixed up?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dlR9C90cfUVu"
      },
      "source": [
        "Recall that training was *unsupervised*, i.e., it was performed without knowing the labels. The autoencoder is a generative model, so we are also interested in seing how well it is able to draw samples from the data distribution it has learned. To do that we will use the autoencoder to *generate* new samples. We do this by generating latent vectors that span a 2D grid (defined by `grid_x` and `grid_y` below) and then feed each latent vector on the grid into the decoder to generate an image:\n",
        "\n",
        "**Note** you will have to adjust the limits of `grid_x` and `grid_y` to match the range of latent features that your model has learned."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nqk1t_r8_UNu"
      },
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "n = 20\n",
        "digit_size = 28\n",
        "figure = np.zeros((digit_size * n, digit_size * n))\n",
        "\n",
        "# linearly spaced coordinates corresponding to the 2D plot\n",
        "# of digit classes in the latent space\n",
        "grid_x = np.linspace(-5, 5, n)        # Task : Set range according to your latent representation\n",
        "grid_y = np.linspace(-5, 5, n)[::-1]  # Task : Set range according to your latent representation\n",
        "\n",
        "for i, yi in enumerate(grid_y):\n",
        "    for j, xi in enumerate(grid_x):\n",
        "        z_sample = np.array([[xi, yi]])\n",
        "        x_decoded = decoder.predict(z_sample.reshape(1,1,1,2))\n",
        "        digit = x_decoded[0].reshape(digit_size, digit_size)\n",
        "        figure[i * digit_size: (i + 1) * digit_size,\n",
        "                j * digit_size: (j + 1) * digit_size] = digit\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "start_range = digit_size // 2\n",
        "end_range = n * digit_size + start_range + 1\n",
        "pixel_range = np.arange(start_range, end_range, digit_size)\n",
        "sample_range_x = np.round(grid_x, 1)\n",
        "sample_range_y = np.round(grid_y, 1)\n",
        "plt.xticks(pixel_range, sample_range_x)\n",
        "plt.yticks(pixel_range, sample_range_y)\n",
        "plt.xlabel(\"z[0]\")\n",
        "plt.ylabel(\"z[1]\")\n",
        "plt.imshow(figure, cmap='Greys_r')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gPirfGdCmZ2v"
      },
      "source": [
        "The grid above shows the \"digits\" generated by the autoencoder for different combinations of 2D latent vectors (corresponding to the values on the x- and y-axis)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FnnBuylfeHO_"
      },
      "source": [
        "###2.5 Question\n",
        "1. Which digits can the autoencoder generate faithfully, which digits does it have trouble generating? Can you explain why?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EwXofIvQhOdN"
      },
      "source": [
        "###2.6 Encoding, decoding and latent space interpolation\n",
        "Now that we have trained an autoencoder, we can use it to encode existing images and generate new images (from a latent representation). With the latent representation we can also start doing interpolation between training samples.\n",
        "\n",
        "Your task is to \n",
        "\n",
        "1. Encode an image of a 7 and an image of a 9 (or any other pair if you refer)\n",
        "2. Decode the encodings to generate reconstructed images\n",
        "3. Interpolate between the two digits in latent space\n",
        "\n",
        "You just need to figure out what the missing shapes `???` should be."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qoBUEtgdSu_i"
      },
      "source": [
        "# Draw two samples (a 7 and a 9) and display them\n",
        "y_test_category = np.argmax(y_test,axis=1)\n",
        "ix7 = np.where(y_test_category==7)[0][1] # Pick a 7\n",
        "ix9 = np.where(y_test_category==9)[0][1] # Pick a 9\n",
        "\n",
        "print('Original input images')\n",
        "plt.subplot(221);plt.imshow(x_test[ix7,:].squeeze(),cmap='gray')\n",
        "plt.subplot(222);plt.imshow(x_test[ix9,:].squeeze(),cmap='gray')\n",
        "plt.show()\n",
        "\n",
        "# Subtask 1 (encoding):\n",
        "# Calculate the latent representation of each sample using the encoder\n",
        "z7 = encoder.predict(x_test[ix7,:].reshape(1,???,???,1))\n",
        "z9 = encoder.predict(x_test[ix9,:].reshape(1,???,???,1))\n",
        "\n",
        "# Subtask 2 (decoding):\n",
        "# Reconstruct images from the two latent vectors using the decoder\n",
        "x_hat_7 = decoder.predict(z7.reshape(1,1,1,???))\n",
        "x_hat_9 = decoder.predict(z9.reshape(1,1,1,???))\n",
        "\n",
        "# Show reconstruction\n",
        "print('Images reconstructewd from 2D latens representation')\n",
        "plt.subplot(223);plt.imshow(x_hat_7.squeeze(),cmap='gray')\n",
        "plt.subplot(224);plt.imshow(x_hat_9.squeeze(),cmap='gray')\n",
        "plt.show()\n",
        "\n",
        "# Subtask 3 (interpolate):\n",
        "# Just run - no changes required)\n",
        "N = 8\n",
        "interp_features = np.zeros((N,latent_dim))\n",
        "for i in range(latent_dim):\n",
        "  interp_features[:,i] = np.linspace(z7[0,i].squeeze(),z9[0,i].squeeze(),N)\n",
        "\n",
        "print('Interpolated images')\n",
        "plt.figure(figsize=(20,6))\n",
        "for i in range(N):\n",
        "  x = interp_features[i,:].reshape(1,1,1,latent_dim)\n",
        "  out = decoder.predict(x)\n",
        "  plt.subplot(1,N,i+1)\n",
        "  plt.imshow(out.squeeze(),cmap='gray')  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DXzGH3sj_QFI"
      },
      "source": [
        "**Note** if the reconstructions of the 7 and the 9 do not look nice enough, try loading different instances from the test set by modifying the indices indicated `A` and `B` below:\n",
        "\n",
        "```\n",
        "ix7 = np.where(y_test_category==7)[0][A] # Pick a 7\n",
        "ix9 = np.where(y_test_category==9)[0][B] # Pick a 9\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8QQDl2AC22AF"
      },
      "source": [
        "**Perspectives:** Given a dataset of facial images, you could use latent space interpolation to generate images like these:\n",
        "\n",
        "![alt text](https://github.com/davidsandberg/facenet/wiki/20170708-150701-add_smile.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gE_--D90TNsz"
      },
      "source": [
        "##3. Task 2: Implement a CAE from scratch\n",
        "The purpose of this task is to test if you can implement a CAE from scratch. **I recommend you skip ahead and complete the tasks on variational encoders and GANs first, and then return to this task later**.\n",
        "\n",
        "Your task is to implement this CAE archtecture for MNIST:\n",
        "\n",
        "![alt text](https://github.com/aivclab/dlcourse/raw/master/data/Lab9_CAE_architecture.png)\n",
        "\n",
        "**Explanation**:\n",
        "- \"Conv 1\", \"Conv 2\", \"Conv 3\", \"D Conv 1\", \"D Conv 2\", \"D Conv 3\", and \"D Conv 4\" are *all* regular 2D convolutions: [Conv2D](https://keras.io/layers/convolutional/#conv2d).\n",
        "- \"M.P\" is short for Max Pooling\n",
        "- \"U.S\" is short for upsampling. You must use [UpSampling2D](https://keras.io/layers/convolutional/#upsampling2d) and **not** [Conv2DTranspose](https://keras.io/layers/convolutional/#conv2dtranspose). (What's the difference by the way?)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PKKWQM-xVlfI"
      },
      "source": [
        "##4. Task 3: Variational Autoencoder\n",
        "Recall that variational autoencoders (VAE) are designed to learn smooth latent space representation. The problem with traditional autoencoders is that they tend to generate gaps in the latent space, making interpolation impossible. The purpose of this task is to see if this is actually the case in practise.\n",
        "\n",
        "Below is an implementation of a convolutional VAE."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "odSqFFagAaCR"
      },
      "source": [
        "from keras.layers import UpSampling2D, ZeroPadding2D, Conv2DTranspose, Lambda, Reshape\n",
        "\n",
        "batch_size = 256\n",
        "latent_dim = 2\n",
        "\n",
        "def sampling(args):\n",
        "    z_mean, z_log_sigma = args\n",
        "    epsilon = K.random_normal(shape=(batch_size, latent_dim))\n",
        "    return z_mean + K.exp(z_log_sigma) * epsilon\n",
        "\n",
        "# VAE model = encoder + decoder\n",
        "# build encoder model\n",
        "inputs = Input(shape=(28, 28, 1),name='encoder_input')\n",
        "x = ZeroPadding2D(padding=(2, 2))(inputs)\n",
        "x = Conv2D(8, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(16, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(32, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(64, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "x = Conv2D(64, kernel_size=(3, 3), strides=(2,2), activation='relu', padding='same')(x)\n",
        "\n",
        "# shape info needed to build decoder model\n",
        "shape = K.int_shape(x)\n",
        "\n",
        "# generate latent vector Q(z|X)\n",
        "x = Flatten()(x)\n",
        "z_mean = Dense(latent_dim, name='z_mean')(x)\n",
        "z_log_var = Dense(latent_dim, name='z_log_var')(x)\n",
        "\n",
        "# Use reparameterization trick to push the sampling out as input.\n",
        "# Note that \"output_shape\" isn't necessary with the TensorFlow backend\n",
        "z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])\n",
        "\n",
        "# instantiate encoder model\n",
        "encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')\n",
        "encoder.summary()\n",
        "\n",
        "# build decoder model\n",
        "latent_inputs = Input(shape=(latent_dim,), name='z_sampling')\n",
        "x = Dense(shape[1] * shape[2] * shape[3], activation='relu')(latent_inputs)\n",
        "x = Reshape((shape[1], shape[2], shape[3]))(x)\n",
        "\n",
        "x = Conv2DTranspose(64, (1,1), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(32, (3,3), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(16, (3,3), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(8, (3,3), strides=(2,2), padding='same')(x)\n",
        "x = Conv2DTranspose(8, (3,3), strides=(2,2), padding='same')(x)\n",
        "outputs = Conv2D(1, kernel_size=(5, 5), padding='valid', activation='sigmoid')(x)\n",
        "\n",
        "# instantiate decoder model\n",
        "decoder = Model(latent_inputs, outputs, name='decoder')\n",
        "decoder.summary()\n",
        "\n",
        "# instantiate VAE model\n",
        "outputs = decoder(encoder(inputs)[2])\n",
        "vae = Model(inputs, outputs, name='vae')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1oRlpirW_4g"
      },
      "source": [
        "###4.1 Questions\n",
        "1. What do you think the `sampling` function does?\n",
        "2. The encoder outputs three variables: `[z_mean, z_log_var, z]`. What do you think they represent?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yIYCApwOTA11"
      },
      "source": [
        "###4.2 Loss function\n",
        "The loss consists of two terms:\n",
        "\n",
        "- A reconstruction term (or similarity term)\n",
        "- and a KL divergence term\n",
        "\n",
        "You can read more about it here: https://towardsdatascience.com/intuitively-understanding-variational-autoencoders-1bfe67eb5daf\n",
        "\n",
        "The KL term is:\n",
        "\n",
        "![alt text](https://miro.medium.com/max/520/1*uEAxCmyVKxzZOJG6afkCCg.png)\n",
        "\n",
        "**Sub-task:** See if you can identify the individual terms in the code block below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NkGEWszvFb_A"
      },
      "source": [
        "from keras.losses import mse\n",
        "reconstruction_loss = mse(K.flatten(inputs), K.flatten(outputs))\n",
        "reconstruction_loss *= 28 * 28\n",
        "kl_loss = K.exp(z_log_var) + K.square(z_mean) - z_log_var/2 - 1\n",
        "kl_loss = K.sum(kl_loss, axis=-1)\n",
        "kl_loss *= 0.5\n",
        "vae_loss = K.mean(reconstruction_loss + kl_loss)\n",
        "vae.add_loss(vae_loss)\n",
        "vae.compile(optimizer='rmsprop')\n",
        "vae.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aMnJJSHMazvZ"
      },
      "source": [
        "###4.3 Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8B6HwxBBFs9v"
      },
      "source": [
        "num_samples = int(np.floor(x_train.shape[0] / batch_size) * batch_size)\n",
        "vae.fit(x_train[0:num_samples,:], epochs=30, batch_size=batch_size,\n",
        "        shuffle=True, verbose=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "__jYiiWaa4tp"
      },
      "source": [
        "###4.4 Plot the latent space representation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CSyV2LxfvYj_"
      },
      "source": [
        "z_mean, _, _ = encoder.predict(x_train[0:num_samples,:],\n",
        "                                batch_size=batch_size)\n",
        "import matplotlib as mpl\n",
        "mpl.rc('image', cmap='jet')\n",
        "plt.figure(figsize=(12, 10))\n",
        "plt.scatter(z_mean[:, 0], z_mean[:, 1], c=np.argmax(y_train[0:num_samples,:],axis=1))\n",
        "plt.colorbar()\n",
        "plt.xlabel(\"z[0]\")\n",
        "plt.ylabel(\"z[1]\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rmdhxSxk2LKw"
      },
      "source": [
        "n = 20\n",
        "digit_size = 28\n",
        "figure = np.zeros((digit_size * n, digit_size * n))\n",
        "# linearly spaced coordinates corresponding to the 2D plot\n",
        "# of digit classes in the latent space\n",
        "grid_x = np.linspace(-2, 2, n)\n",
        "grid_y = np.linspace(-2, 2, n)[::-1]\n",
        "\n",
        "for i, yi in enumerate(grid_y):\n",
        "    for j, xi in enumerate(grid_x):\n",
        "        z_sample = np.array([[xi, yi]])\n",
        "        x_decoded = decoder.predict(z_sample)\n",
        "        digit = x_decoded[0].reshape(digit_size, digit_size)\n",
        "        figure[i * digit_size: (i + 1) * digit_size,\n",
        "                j * digit_size: (j + 1) * digit_size] = digit\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "start_range = digit_size // 2\n",
        "end_range = n * digit_size + start_range + 1\n",
        "pixel_range = np.arange(start_range, end_range, digit_size)\n",
        "sample_range_x = np.round(grid_x, 1)\n",
        "sample_range_y = np.round(grid_y, 1)\n",
        "plt.xticks(pixel_range, sample_range_x)\n",
        "plt.yticks(pixel_range, sample_range_y)\n",
        "plt.xlabel(\"z[0]\")\n",
        "plt.ylabel(\"z[1]\")\n",
        "plt.imshow(figure, cmap='Greys_r')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lRrVCxqraIS5"
      },
      "source": [
        "###4.5 Questions\n",
        "1. What do you think of this latent representation? In terms of quality? In terms of smoothness? Compare to the same plot for the traditional convolutional autoencoder.\n",
        "2. Which digits does the model faithfully reconstruct? Which digits does it have trouble reconstructing? Why?\n",
        "3. What happens if you set the weight of the KL term to, say 5 (`kl_loss *= 5`), and re-train the model?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wR5jbEkPbYIa"
      },
      "source": [
        "###4.6 Encoding, decoding and latent space interpolation\n",
        "Like we did for the traditional conovoolutional autoencoder (see section 2.6), your task is to \n",
        "\n",
        "1. Encode an image of a 7 and an image of a 9.\n",
        "2. Decode the encodings (to generate reconstructed images)\n",
        "3. Interpolate between the two digits in latent space"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDTMCG_d21ue"
      },
      "source": [
        "The challenge here is that the encoder expects a fixed batch size (of 256 in our case). So you cannot just input one image to the model. The solution is to first process the training data set in batches of 256, and store the encodings:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wK3ocjvz20YT"
      },
      "source": [
        "# Get 2D encoding of all test examples (in batches of 100 images)\n",
        "num_samples = int(np.floor(x_test.shape[0] / batch_size) * batch_size)\n",
        "z_test = np.zeros((num_samples,latent_dim))\n",
        "num_batches = int(num_samples / batch_size)\n",
        "for i in range(num_batches):\n",
        "  batch = x_test[i*batch_size:(i+1)*batch_size,:,:,:]\n",
        "  [z_m, z_lv, z_val] = encoder.predict(batch,batch_size=batch_size)\n",
        "  z_test[i*batch_size:(i+1)*batch_size,:] = z_m"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rDZwMmg49oOF"
      },
      "source": [
        "# Draw two samples (a 7 and a 9)\n",
        "y_test_category = np.argmax(y_test,axis=1)\n",
        "ix7 = np.where(y_test_category==7)[0][1]\n",
        "ix9 = np.where(y_test_category==9)[0][1]\n",
        "\n",
        "print('Original input images')\n",
        "plt.subplot(221);plt.imshow(x_test[ix7,:].squeeze(),cmap='gray')\n",
        "plt.subplot(222);plt.imshow(x_test[ix9,:].squeeze(),cmap='gray')\n",
        "plt.show()\n",
        "\n",
        "# Subtask 1 (encoding): Calculate the latent representation of each sample\n",
        "z7 = z_test[ix7,:]\n",
        "z9 = z_test[ix9,:]\n",
        "\n",
        "# Subtask 2 (decoding): Reconstruct images from the two latent vectors\n",
        "x_hat_7 = decoder.predict(z7.reshape(1,???))\n",
        "x_hat_9 = decoder.predict(z9.reshape(1,???))\n",
        "\n",
        "# Show reconstruction\n",
        "plt.subplot(223);plt.imshow(x_hat_7.squeeze(),cmap='gray')\n",
        "plt.subplot(224);plt.imshow(x_hat_9.squeeze(),cmap='gray')\n",
        "\n",
        "N = 8\n",
        "interp_features = np.zeros((N,latent_dim))\n",
        "for i in range(latent_dim):\n",
        "  interp_features[:,i] = np.linspace(z7[i].squeeze(),z9[i].squeeze(),N)\n",
        "\n",
        "plt.figure(figsize=(20,6))\n",
        "for i in range(N):\n",
        "  x = interp_features[i,:].reshape(1,latent_dim)\n",
        "  out = decoder.predict(x)\n",
        "  plt.subplot(1,N,i+1)\n",
        "  plt.imshow(out.squeeze(),cmap='gray')  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pg5OtlGMQLYm"
      },
      "source": [
        "**Note:** If you want to make nicer reconstructions and better interpolations, increase the latent dimensionality (latent_dim) and re-train the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_5WzyIk9LFB"
      },
      "source": [
        "##5. Task 4: Generative Adversarial Networks\n",
        "Last year's version of this task was based on [this tutorial](https://github.com/eriklindernoren/Keras-GAN/blob/master/dcgan/dcgan.py). Its an implementation of a Deep Convolutional GAN (DCGAN) for MNIST.\n",
        "\n",
        "Unfortunately that doesn't work anymore and I haven't been able to figure out why. Therefore, I advice you to go through this tutorial instead: https://www.tensorflow.org/tutorials/generative/dcgan\n",
        "\n",
        "You can load it as a notebook in Colab using [this link](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/generative/dcgan.ipynb)\n",
        "\n",
        "###5.1 Recommended tasks (advanced and not mandatory)\n",
        "1. See if you can figure out how to decrease the dimensionality of the latent space from 100 to 2? Then retrain the model.\n",
        "\n",
        "2. Extend the code such that you can train the model to learn a 2D latent representation of MNIST and subsequently make it generate images based on some 2D latent vector that you specify. Use this to make a plot of the 2D latent space, like we did above. What do you observe?\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VM8KYT38Ekd3"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}