05_generative_models/Auto Encoders.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2b329d8",
   "metadata": {},
   "source": [
    "# Auto Encoders\n",
    "\n",
    "At it's core, an autoencoder is a function that takes complex inputs and performs a dimensionality reduction of the data.  Principal Component Analysis is an example of a classical encoding algorithm.  Even a simple linear regression in 2D can serve as an encoder, since you map 2 dimensions (x and y) into just one dimension (position along a line).\n",
    "\n",
    "An important feature of autoencoders is the ability to reconstruct the input data from the compressed feature representation.  It's easy to see how this is done when you have used PCA or linear regression, but it is more complex with a neural network.\n",
    "\n",
    "One advantage of using neural networks as auto-encoders is that they have capacity to encode complex data, and the training regime for encoding is not overly complex and similar for many types of data.\n",
    "\n",
    "In this notebook, we will build an autoencoder for handwritten digits between 0 and 99, inclusive, generated via the mnist dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f23983f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load numpy and tensorflow:\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy\n",
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from IPython.display import display, clear_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22dd7e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use TF to get the dataset, will download if needed.\n",
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "x_train = x_train.astype(numpy.float32) * (1./256)\n",
    "x_test  = x_test.astype(numpy.float32) * (1./256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d08e9e58",
   "metadata": {},
   "source": [
    "Auto Encoders do a really good job at taking input, compressing it into a small latent vector, and then decoding that back into the original image.  But, it has some downsides.  Let's demonstrate this with a simple convolutional auto encoder on MNIST data.  This is based off of a resnet architecture, but a lot smaller:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5949d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResBlock(tf.keras.models.Model):\n",
    "    \n",
    "    '''\n",
    "    Residual block with configurable activation\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, n_filters, activation=tf.nn.relu):\n",
    "        tf.keras.models.Model.__init__(self)\n",
    "\n",
    "        self.convolution1 = tf.keras.layers.Convolution2D(\n",
    "            kernel_size = [3, 3], \n",
    "            filters     = n_filters,\n",
    "            padding     = \"same\",\n",
    "            activation  = None,\n",
    "        )\n",
    "        \n",
    "        self.convolution2 = tf.keras.layers.Convolution2D(\n",
    "            kernel_size = [3, 3], \n",
    "            filters     = n_filters,\n",
    "            padding     = \"same\",\n",
    "            activation  = None,\n",
    "        )\n",
    "        \n",
    "        self.batch_norm1 = tf.keras.layers.BatchNormalization()\n",
    "        self.batch_norm2 = tf.keras.layers.BatchNormalization()\n",
    "\n",
    "\n",
    "        self.activation = activation\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        \n",
    "        # Conv 1\n",
    "        x = self.convolution1(inputs)\n",
    "        x = self.batch_norm1(x)\n",
    "        x = self.activation(x)\n",
    "        \n",
    "        # Conv2\n",
    "        x = self.convolution2(x)\n",
    "        \n",
    "        x = self.batch_norm2(x)\n",
    "        \n",
    "        # Sum:\n",
    "        x = x + inputs \n",
    "        \n",
    "        x = self.activation(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "459140ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvolutionalEncoder(tf.keras.models.Model):\n",
    "    '''\n",
    "    Simple autoencoder forward layer\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, activation=tf.nn.relu, latent_size = 10):\n",
    "        tf.keras.models.Model.__init__(self)\n",
    "        self._latent_size = latent_size\n",
    "\n",
    "        # Apply a 5x5 kernel to the image:\n",
    "        self.encoder_layer_1 = ResBlock(activation=activation, n_filters=64)\n",
    "        \n",
    "        # Use a 2x2 kernel of stride 2x2 to downsample:\n",
    "        self.pool_1 = tf.keras.layers.MaxPool2D()\n",
    "        \n",
    "        self.encoder_layer_2 = ResBlock(activation=activation, n_filters=64)\n",
    "        self.pool_2 = tf.keras.layers.MaxPool2D()\n",
    "        \n",
    "        self.encoder_layer_3 = ResBlock(activation=activation, n_filters=64)\n",
    "        self.pool_3 = tf.keras.layers.MaxPool2D()\n",
    "\n",
    "        self.encoder_layer_final = tf.keras.layers.Dense(\n",
    "            units = latent_size,\n",
    "            activation = None,\n",
    "            )\n",
    "        \n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "    def call(self, inputs):\n",
    "        \n",
    "        batch_size = inputs.shape[0]\n",
    "        x = inputs\n",
    "        # Make sure the input is the right shape:\n",
    "        x = tf.reshape(x, [batch_size, 28, 28, 1])\n",
    "            \n",
    "        x = self.encoder_layer_1(x)\n",
    "        x = self.pool_1(x)\n",
    "        x = self.encoder_layer_2(x)\n",
    "        x = self.pool_2(x)\n",
    "        x = self.encoder_layer_3(x)\n",
    "        x = self.pool_3(x)\n",
    "        x = tf.keras.layers.Flatten()(x)\n",
    "        x = self.encoder_layer_final(x)\n",
    "        return tf.reshape(x, [batch_size, self._latent_size])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "094769b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvolutionalDecoder(tf.keras.models.Model):\n",
    "    \n",
    "    def __init__(self, activation=tf.nn.tanh):\n",
    "        tf.keras.models.Model.__init__(self)\n",
    "\n",
    "        # The decoder runs the encoder steps but in reverse.\n",
    "        \n",
    "        # The first step is a dense layer to get to the right number of units.\n",
    "        # It turns out to need to be of shape [3, 3]\n",
    "        self.dense = tf.keras.layers.Dense(\n",
    "            units = 3 * 3 * 64\n",
    "        )\n",
    "        \n",
    "        \n",
    "        self.decoder_layer_1 = ResBlock(n_filters=64, activation=activation)\n",
    "        \n",
    "        self.unpool_1 = tf.keras.layers.UpSampling2D(\n",
    "            size          = 2,\n",
    "            interpolation = \"nearest\",\n",
    "        )\n",
    "\n",
    "        self.decoder_layer_2 = ResBlock(n_filters=64, activation=activation)\n",
    "        \n",
    "        self.unpool_2 = tf.keras.layers.UpSampling2D(\n",
    "            size          = 2,\n",
    "            interpolation = \"nearest\",\n",
    "        )\n",
    "\n",
    "\n",
    "        self.decoder_layer_3 = ResBlock(n_filters=64, activation=activation)\n",
    "        \n",
    "        self.decoder_layer_final = tf.keras.layers.Convolution2D(\n",
    "            kernel_size = [1, 1],\n",
    "            filters     = 1,\n",
    "            padding     = \"same\",\n",
    "            use_bias    = True,\n",
    "            kernel_regularizer = tf.keras.regularizers.L2(),\n",
    "            activation  = None,\n",
    "        )\n",
    "        \n",
    "\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        ''' \n",
    "        Reshape at input and output: \n",
    "        '''\n",
    "            \n",
    "        \n",
    "        batch_size = inputs.shape[0]\n",
    "\n",
    "\n",
    "        x = self.dense(inputs)\n",
    "\n",
    "\n",
    "        # First Step is to to un-pool the encoded state into the right shape:\n",
    "        x = tf.reshape(x, [batch_size, 3, 3, 64])\n",
    "        x = tf.image.resize(\n",
    "            x,\n",
    "            size=[7,7]\n",
    "        )\n",
    "        \n",
    "        x = self.decoder_layer_1(x)\n",
    "        x = self.unpool_1(x)\n",
    "        x = self.decoder_layer_2(x)\n",
    "        x = self.unpool_2(x)\n",
    "        x = self.decoder_layer_3(x)\n",
    "        x = self.decoder_layer_final(x)\n",
    "\n",
    "        x = tf.nn.sigmoid(x)\n",
    "        \n",
    "        x = tf.reshape(x, [batch_size, 28, 28, 1])\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00db2bf4",
   "metadata": {},
   "source": [
    "Let's initialize an instance of each of these networks and do a forward pass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e902ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = ConvolutionalEncoder(latent_size=20)\n",
    "decoder = ConvolutionalDecoder()\n",
    "test_input = x_train[0].reshape([1,28,28,1])\n",
    "\n",
    "intermediate_state = encoder(test_input)\n",
    "decoded_image = decoder(intermediate_state)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e429d26",
   "metadata": {},
   "source": [
    "Let's inspect the layers of the models:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a2758eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649de0f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0a08476",
   "metadata": {},
   "source": [
    "What does the intermediate state and decoded images look like?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4d8ca12",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(intermediate_state)\n",
    "plt.imshow(intermediate_state);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb1817d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And here is the output image:\n",
    "decoded_images_reshape = decoded_image.numpy().reshape(28, 28)\n",
    "plt.imshow(decoded_images_reshape);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8924b2ca",
   "metadata": {},
   "source": [
    "Unsuprisingly, this is just empty noise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d7e7b6b",
   "metadata": {},
   "source": [
    "# Training an Autoencoder\n",
    "\n",
    "With a trained encoder/decoder pair, we want the input image to be as close as possible to the output image.  The loss for this network, then, can be constructed as just the mean-square-error of input to output.  We will train this with an AdamOptimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf40ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These parameters control the training behavior: batch size and how many iterations to use\n",
    "BATCH_SIZE=256\n",
    "N_EPOCHS=10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f762dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam()\n",
    "loss_function = tf.keras.losses.MeanSquaredError()\n",
    "\n",
    "# We will keep track of loss along the way:\n",
    "steps = []\n",
    "loss_history = []\n",
    "val_loss_history = []\n",
    "val_steps = []\n",
    "\n",
    "# Define a figure here which we can update along the way:\n",
    "fig       = plt.figure(figsize=(16,9))\n",
    "# Create 2x2 sub plots\n",
    "gs = gridspec.GridSpec(2, 2)\n",
    "ax_loss   = plt.subplot(gs[:,0])\n",
    "ax_input  = plt.subplot(gs[0,1])\n",
    "ax_output = plt.subplot(gs[1,1])\n",
    "\n",
    "\n",
    "for e in range(N_EPOCHS):\n",
    "    \n",
    "    # First, shuffle the training data:\n",
    "    numpy.random.shuffle(x_train)\n",
    "    \n",
    "    # Now, loop over the data in batches:\n",
    "    for batch in range(int(60000/BATCH_SIZE)):\n",
    "    \n",
    "        # Variable for plotting:\n",
    "        epoch_value = e + (batch*BATCH_SIZE)/60000\n",
    "    \n",
    "        # Load some data:\n",
    "        batch_images = x_train[batch*BATCH_SIZE : (batch+1)*BATCH_SIZE]\n",
    "        # Reshape the data:\n",
    "        batch_images = batch_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "        with tf.GradientTape() as tape:\n",
    "            intermediate_state = encoder(batch_images)\n",
    "            decoded_images = decoder(intermediate_state)\n",
    "            loss_value = loss_function(batch_images, decoded_images)\n",
    "        \n",
    "        # Collect the trainable variables as the union of the encoder and decoder variables:\n",
    "        # (This is a list concatenation, not an addition)\n",
    "        trainable_vars = encoder.trainable_variables + decoder.trainable_variables\n",
    "\n",
    "        # Update the loss history:\n",
    "        loss_history.append(loss_value.numpy())\n",
    "        steps.append(epoch_value)\n",
    "\n",
    "        # Apply the update to the model:\n",
    "        grads = tape.gradient(loss_value, trainable_vars)\n",
    "        optimizer.apply_gradients(zip(grads, trainable_vars))\n",
    "\n",
    "\n",
    "            \n",
    "        if batch % 25 == 0:\n",
    "            test_indexes = numpy.random.choice(range(10000), size=BATCH_SIZE)\n",
    "            test_images = x_test[test_indexes]\n",
    "            test_images = test_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "            val_intermediate_state = encoder(test_images)\n",
    "            val_decoded_images = decoder(val_intermediate_state)\n",
    "            val_loss_value = loss_function(test_images, val_decoded_images)\n",
    "            val_loss_history.append(val_loss_value.numpy())\n",
    "            val_steps.append(epoch_value)\n",
    "\n",
    "            # Here we update the plots.  Replot the loss, and show the latest input and output images\n",
    "            ax_loss.set_xlim(0, 1.1*epoch_value)\n",
    "            ax_loss.cla()\n",
    "            ax_loss.plot(steps, loss_history, label=\"Train Loss\")\n",
    "            ax_loss.plot(val_steps, val_loss_history, label=\"Test Loss\")\n",
    "            ax_loss.grid(True)\n",
    "            ax_loss.legend(fontsize=25)\n",
    "\n",
    "            # Display an input image:\n",
    "            ax_input.imshow(batch_images[0].reshape(28, 28),vmin=0, vmax=1)\n",
    "\n",
    "            # Display an output image:\n",
    "            ax_output.imshow(decoded_images.numpy()[0].reshape(28, 28),vmin=0, vmax=1)\n",
    "\n",
    "            clear_output(wait = True)\n",
    "            display(fig)\n",
    "            plt.pause(0.1)\n",
    "\n",
    "        \n",
    "# Let's save the models weights as well:\n",
    "encoder.save_weights(\"encoder_basic.h5\")\n",
    "decoder.save_weights(\"decoder_basic.h5\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe9249e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # If you don't want or don't have time to train, you can preload the models:\n",
    "if False:\n",
    "    encoder.load_weights(\"./saved_models/AE/encoder_basic.h5\")\n",
    "    decoder.load_weights(\"./saved_models/AE/decoder_basic.h5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d07c02d",
   "metadata": {},
   "source": [
    "This autoencoder is doing great, you should have pretty good results after just half an epoch.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ca8e0d",
   "metadata": {},
   "source": [
    "# What can you do with Auto Encoders?\n",
    "\n",
    "## Application: Data Compression\n",
    "\n",
    "In principle, auto encoders are good for data compression:  its like non-linear PCA.  But you can do more than just compression, as well.  In order to decode an image, the latent state ought to contain just about all of the information that is important to that image.  So, you could take the latent states directly and use them to classify digits.\n",
    "\n",
    "To demonstrate this, let's use a small classifier network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a8c4fbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleClassifier(tf.keras.models.Model):\n",
    "    \n",
    "    '''\n",
    "    Very simple classifier for latent space to mnist classification\n",
    "    '''\n",
    "    \n",
    "    def __init__(self):\n",
    "        tf.keras.models.Model.__init__(self)\n",
    "\n",
    "        self.dense = tf.keras.layers.Dense(\n",
    "            units = 10, \n",
    "        )\n",
    "\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        \n",
    "        return self.dense(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a907721",
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier = SimpleClassifier()\n",
    "\n",
    "classification = classifier(intermediate_state)\n",
    "classifier.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5384249",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reloading the dataset here to make sure it's shuffled properly\n",
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "x_train = x_train.astype(numpy.float32) * (1./256)\n",
    "x_test  = x_test.astype(numpy.float32) * (1./256)\n",
    "y_train = y_train.astype(numpy.int64)\n",
    "y_test  = y_test.astype(numpy.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584e6864",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_EPOCHS=5\n",
    "BATCH_SIZE=128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc8456f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam()\n",
    "loss_func = tf.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "\n",
    "# We will keep track of loss along the way:\n",
    "classifier_steps = []\n",
    "classifier_loss_history = []\n",
    "classifier_acc_history = []\n",
    "classifier_test_loss_history = []\n",
    "classifier_test_steps = []\n",
    "classifier_test_acc_history = []\n",
    "\n",
    "\n",
    "# Define a figure here which we can update along the way:\n",
    "fig       = plt.figure(figsize=(16,9))\n",
    "# Create 2x2 sub plot\n",
    "gs = gridspec.GridSpec(1,2)\n",
    "ax_loss   = plt.subplot(gs[0,0])\n",
    "ax_acc    = plt.subplot(gs[0,1])\n",
    "\n",
    "\n",
    "for e in range(N_EPOCHS):\n",
    "    \n",
    "    # First, shuffle the training data:\n",
    "    indexes = numpy.arange(len(x_train))\n",
    "    numpy.random.shuffle(indexes)\n",
    "    \n",
    "    # Now, loop over the data in batches:\n",
    "    for batch in range(int(60000/BATCH_SIZE)):\n",
    "    \n",
    "        # Variable for plotting:\n",
    "        epoch_value = e + (batch*BATCH_SIZE)/60000\n",
    "    \n",
    "        batch_index = indexes[batch*BATCH_SIZE:(batch+1)*BATCH_SIZE]\n",
    "        # Load some data:\n",
    "        batch_images = x_train[batch_index]\n",
    "        batch_labels = y_train[batch_index]\n",
    "        # Reshape the data:\n",
    "        batch_images = batch_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "        with tf.GradientTape() as tape:\n",
    "            intermediate_state = encoder(batch_images)\n",
    "            classification = classifier(intermediate_state)\n",
    "\n",
    "            loss_value = loss_func(y_true=batch_labels, y_pred=classification)\n",
    "            prediction = tf.argmax(classification, axis=-1)\n",
    "            accuracy = numpy.mean(prediction.numpy() == batch_labels)\n",
    "        # Collect the trainable variables of just the classifier:\n",
    "        trainable_vars = classifier.trainable_variables\n",
    "\n",
    "        # Update the loss history:\n",
    "        classifier_loss_history.append(loss_value.numpy())\n",
    "        classifier_acc_history.append(accuracy)\n",
    "        classifier_steps.append(epoch_value)\n",
    "\n",
    "        # Apply the update to the model:\n",
    "        grads = tape.gradient(loss_value, trainable_vars)\n",
    "        optimizer.apply_gradients(zip(grads, trainable_vars))\n",
    "\n",
    "\n",
    "            \n",
    "        if batch % 25 == 0:\n",
    "            test_indexes = numpy.random.choice(range(10000), size=BATCH_SIZE)\n",
    "            test_images = x_test[test_indexes]\n",
    "            test_label  = y_test[test_indexes]\n",
    "            test_images = test_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "            test_intermediate_state = encoder(test_images)\n",
    "            test_classification = classifier(test_intermediate_state)\n",
    "            test_loss_value = loss_func(y_true=test_label, y_pred=test_classification)\n",
    "            test_prediction = tf.argmax(test_classification, axis=-1)\n",
    "            test_accuracy = numpy.mean(test_prediction.numpy() == test_label)\n",
    "            classifier_test_loss_history.append(test_loss_value.numpy())\n",
    "            classifier_test_steps.append(epoch_value)\n",
    "            classifier_test_acc_history.append(accuracy)\n",
    "\n",
    "\n",
    "            # Here we update the plots.  Replot the loss and accuracy\n",
    "            ax_loss.set_xlim(0, 1.1*epoch_value)\n",
    "            ax_loss.cla()\n",
    "            ax_loss.plot(classifier_steps, classifier_loss_history, label=\"Train Loss\")\n",
    "            ax_loss.plot(classifier_test_steps, classifier_test_loss_history, label=\"Test Loss\")\n",
    "            ax_loss.grid(True)\n",
    "            ax_loss.legend(fontsize=25)\n",
    "            \n",
    "            ax_acc.set_xlim(0, 1.1*epoch_value)\n",
    "            ax_acc.set_ylim(0, 1.0)\n",
    "            ax_acc.cla()\n",
    "            ax_acc.plot(classifier_steps, classifier_acc_history, label=\"Train Acc.\")\n",
    "            ax_acc.plot(classifier_test_steps, classifier_test_acc_history, label=\"Test Acc.\")\n",
    "            ax_acc.grid(True)\n",
    "            ax_acc.legend(fontsize=25)\n",
    "            \n",
    "            clear_output(wait = True)\n",
    "            display(fig)\n",
    "            plt.pause(0.1)\n",
    "\n",
    "        \n",
    "# Let's save the models weights as well:\n",
    "classifier.save_weights(\"classifier.h5\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56b695d2",
   "metadata": {},
   "source": [
    "As you can see, the accuracy increased to decent levels (> 90%) on both the train and test sets.  This classifier we're training has only 210 parameters!  Clearly, the intermediate state has a lot of information for discriminating the class available.\n",
    "\n",
    "Ok, on MNIST 90% is not actually an impressive number.  But, random guessing would have an accuracy of just 10% on this dataset, so the intermediate state is hardly empty."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6162423",
   "metadata": {},
   "source": [
    "## Image interpolation\n",
    "\n",
    "The latent space is not \"continuous\" in real numbers.  That is, taking two latent vectors and decoding them, and all the numbers in between them, does not produce a suite of real and believable images.  Instead, it slowly transforms one number into the other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2c580d",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_vec = encoder(x_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb34e8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_transformations = 10\n",
    "n_steps = 15\n",
    "\n",
    "\n",
    "image_array = numpy.ndarray((28*(n_transformations), 28*(n_steps+1)))\n",
    "for j in range(n_transformations):\n",
    "    index_a, index_b = numpy.random.choice(range(len(latent_vec)), 2)\n",
    "    for i in range(0,n_steps+1):\n",
    "\n",
    "        latent_vector = (i / n_steps ) *latent_vec[index_a] + ((n_steps - i) / n_steps )*latent_vec[index_b]\n",
    "        latent_vector = tf.reshape(latent_vector, (1, 20))\n",
    "        image_array[28*j:28*(j+1),28*i:28*(i+1)] = decoder(latent_vector).numpy().reshape((28,28))\n",
    "\n",
    "fig = plt.figure(figsize=(16, 10*8))\n",
    "plt.imshow(image_array);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62f86c38",
   "metadata": {},
   "source": [
    "## Example Application: Denoising Auto-Encoders\n",
    "\n",
    "Auto Encoders - in principle - are meant to identify the most important features of an image and discard the least important.  For this reason, you could train an autoencoder on noisy images and decode them into denoised images.  A small modification of the training loop above showcases this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e26b983",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize new networks:\n",
    "encoder_denoising = ConvolutionalEncoder(latent_size=20)\n",
    "decoder_denoising = ConvolutionalDecoder()\n",
    "test_input = x_train[0].reshape([1,28,28,1])\n",
    "\n",
    "intermediate_state = encoder_denoising(test_input)\n",
    "decoded_image = decoder_denoising(intermediate_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7fbcb56",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam()\n",
    "loss_function = tf.keras.losses.MeanSquaredError()\n",
    "\n",
    "# We will keep track of loss along the way:\n",
    "steps = []\n",
    "loss_history = []\n",
    "val_loss_history = []\n",
    "val_steps = []\n",
    "\n",
    "# Define a figure here which we can update along the way:\n",
    "fig       = plt.figure(figsize=(16,9))\n",
    "# Create 2x2 sub plots\n",
    "gs = gridspec.GridSpec(2, 2)\n",
    "ax_loss   = plt.subplot(gs[:,0])\n",
    "ax_input  = plt.subplot(gs[0,1])\n",
    "ax_output = plt.subplot(gs[1,1])\n",
    "\n",
    "\n",
    "for e in range(N_EPOCHS):\n",
    "    \n",
    "    # First, shuffle the training data:\n",
    "    numpy.random.shuffle(x_train)\n",
    "    \n",
    "    # Now, loop over the data in batches:\n",
    "    for batch in range(int(60000/BATCH_SIZE)):\n",
    "    \n",
    "        # Variable for plotting:\n",
    "        epoch_value = e + (batch*BATCH_SIZE)/60000\n",
    "    \n",
    "        # Load some data:\n",
    "        batch_images = x_train[batch*BATCH_SIZE : (batch+1)*BATCH_SIZE]\n",
    "        \n",
    "        \n",
    "        # Reshape the data:\n",
    "        batch_images = batch_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "        noise = tf.random.uniform(shape=batch_images.shape, minval=0, maxval=0.3)\n",
    "        batch_images = batch_images + noise\n",
    "        with tf.GradientTape() as tape:\n",
    "            intermediate_state = encoder_denoising(batch_images)\n",
    "            decoded_images = decoder_denoising(intermediate_state)\n",
    "            loss_value = loss_function(batch_images, decoded_images)\n",
    "        \n",
    "        # Collect the trainable variables as the union of the encoder and decoder variables:\n",
    "        # (This is a list concatenation, not an addition)\n",
    "        trainable_vars = encoder_denoising.trainable_variables + decoder_denoising.trainable_variables\n",
    "\n",
    "        # Update the loss history:\n",
    "        loss_history.append(loss_value.numpy())\n",
    "        steps.append(epoch_value)\n",
    "\n",
    "        # Apply the update to the model:\n",
    "        grads = tape.gradient(loss_value, trainable_vars)\n",
    "        optimizer.apply_gradients(zip(grads, trainable_vars))\n",
    "\n",
    "\n",
    "            \n",
    "        if batch % 25 == 0:\n",
    "            test_indexes = numpy.random.choice(range(10000), size=BATCH_SIZE)\n",
    "            test_images = x_test[test_indexes]\n",
    "            test_images = test_images.reshape([BATCH_SIZE, 28,28,1])\n",
    "            noise = tf.random.uniform(shape=test_images.shape, minval=0, maxval=0.2)\n",
    "            test_images = test_images + noise\n",
    "            val_intermediate_state = encoder_denoising(test_images)\n",
    "            val_decoded_images = decoder_denoising(val_intermediate_state)\n",
    "            val_loss_value = loss_function(test_images, val_decoded_images)\n",
    "            val_loss_history.append(val_loss_value.numpy())\n",
    "            val_steps.append(epoch_value)\n",
    "\n",
    "            # Here we update the plots.  Replot the loss, and show the latest input and output images\n",
    "            ax_loss.set_xlim(0, 1.1*epoch_value)\n",
    "            ax_loss.cla()\n",
    "            ax_loss.plot(steps, loss_history, label=\"Train Loss\")\n",
    "            ax_loss.plot(val_steps, val_loss_history, label=\"Test Loss\")\n",
    "            ax_loss.grid(True)\n",
    "            ax_loss.legend(fontsize=25)\n",
    "\n",
    "            # Display an input image:\n",
    "            ax_input.imshow(batch_images[0].numpy().reshape(28, 28),vmin=0, vmax=1)\n",
    "\n",
    "            # Display an output image:\n",
    "            ax_output.imshow(decoded_images.numpy()[0].reshape(28, 28),vmin=0, vmax=1)\n",
    "\n",
    "            clear_output(wait = True)\n",
    "            display(fig)\n",
    "            plt.pause(0.1)\n",
    "\n",
    "        \n",
    "# Let's save the models weights as well:\n",
    "encoder.save_weights(\"denoising_encoder_basic.h5\")\n",
    "decoder.save_weights(\"denoising_decoder_basic.h5\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a46277c",
   "metadata": {},
   "source": [
    "As you might see above, after a few epochs the auto encoder is reproducing the core of the digits but without noise - it simply isn't something the encoder is designed to store, and therefore gets suppressed.  If you are encoding images or data objects that have a high intrinsic level of fast variations, this is an issue.  On the other hand, if you just want to remove noise from the data, this is a good technique!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda/2021-11-30",
   "language": "python",
   "name": "conda-2021-11-30"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}