C4_W2_Lab_2_MNIST_Autoencoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "C4_W2_Lab_2_MNIST_Autoencoder.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nBXbFax4YP34"
      },
      "source": [
        "## Ungraded Lab: MNIST Autoencoder\n",
        "\n",
        "You will now work on an autoencoder that works on the [MNIST dataset](https://www.tensorflow.org/datasets/catalog/mnist). This will encode the inputs to lower resolution images. The decoder should then be able to generate the original input from this compressed representation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ZYaLxnBYUKA"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3EXwoz-KHtWO"
      },
      "source": [
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "  \n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Is05FyRgYX0c"
      },
      "source": [
        "## Prepare the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xuhe2ksPI8A0"
      },
      "source": [
        "You will load the MNIST data from TFDS into train and test sets. Let's first define a preprocessing function for normalizing and flattening the images. Since we'll be training an autoencoder, this will return `image, image` because the input will also be the target or label while training.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t9F7YsCNIKSA"
      },
      "source": [
        "def map_image(image, label):\n",
        "  '''Normalizes and flattens the image. Returns image as input and label.'''\n",
        "  image = tf.cast(image, dtype=tf.float32)\n",
        "  image = image / 255.0\n",
        "  image = tf.reshape(image, shape=(784,))\n",
        "\n",
        "  return image, image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "08I1mbYSLbsC"
      },
      "source": [
        "# Load the train and test sets from TFDS\n",
        "\n",
        "BATCH_SIZE = 128\n",
        "SHUFFLE_BUFFER_SIZE = 1024\n",
        "\n",
        "train_dataset = tfds.load('mnist', as_supervised=True, split=\"train\")\n",
        "train_dataset = train_dataset.map(map_image)\n",
        "train_dataset = train_dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE).repeat()\n",
        "\n",
        "test_dataset = tfds.load('mnist', as_supervised=True, split=\"test\")\n",
        "test_dataset = test_dataset.map(map_image)\n",
        "test_dataset = test_dataset.batch(BATCH_SIZE).repeat()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for t in train_dataset:\n",
        "  print(t[0].shape, t[1].shape)\n",
        "  break"
      ],
      "metadata": {
        "id": "nS-jbazDijX4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z70w2KXjYk32"
      },
      "source": [
        "## Build the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1MKu48lBM2Qg"
      },
      "source": [
        "You will now build a simple autoencoder to ingest the data. Like before, the encoder will compress the input and reconstructs it in the decoder output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KRrE2BV4IpzR"
      },
      "source": [
        "def simple_autoencoder(inputs):\n",
        "  '''Builds the encoder and decoder using Dense layers.'''\n",
        "  encoder = tf.keras.layers.Dense(units=32, activation='relu')(inputs)\n",
        "  decoder = tf.keras.layers.Dense(units=784, activation='sigmoid')(encoder)\n",
        "  \n",
        "  return encoder, decoder\n",
        "\n",
        "# set the input shape\n",
        "inputs =  tf.keras.layers.Input(shape=(784,))\n",
        "\n",
        "# get the encoder and decoder output\n",
        "encoder_output, decoder_output = simple_autoencoder(inputs)\n",
        "\n",
        "# setup the encoder because you will visualize its output later\n",
        "encoder_model = tf.keras.Model(inputs=inputs, outputs=encoder_output)\n",
        "\n",
        "# setup the autoencoder\n",
        "autoencoder_model = tf.keras.Model(inputs=inputs, outputs=decoder_output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8k9OnSM4YxJd"
      },
      "source": [
        "## Compile the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DvvDqY_XQPyb"
      },
      "source": [
        "You will setup the model for training. You can use binary crossentropy to measure the loss between pixel values that range from 0 (black) to 1 (white)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cFwmAhWAYwcc"
      },
      "source": [
        "autoencoder_model.compile(\n",
        "    optimizer=tf.keras.optimizers.Adam(), \n",
        "    loss='binary_crossentropy')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zzLf0oQ1Y0cI"
      },
      "source": [
        "## Train the Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vsaSjlAgYz-7"
      },
      "source": [
        "train_steps = 60000 // BATCH_SIZE\n",
        "simple_auto_history = autoencoder_model.fit(train_dataset, \n",
        "                                            steps_per_epoch=train_steps, epochs=50)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1_sKLXnBZFSc"
      },
      "source": [
        "## Display sample results\n",
        "\n",
        "You can now visualize the results. The utility functions below will help in plotting the encoded and decoded values."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5tgFgilORr0M"
      },
      "source": [
        "def display_one_row(disp_images, offset, shape=(28, 28)):\n",
        "  '''Display sample outputs in one row.'''\n",
        "  for idx, test_image in enumerate(disp_images):\n",
        "    plt.subplot(3, 10, offset + idx + 1)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    test_image = np.reshape(test_image, shape)\n",
        "    plt.imshow(test_image, cmap='gray')\n",
        "\n",
        "\n",
        "def display_results(disp_input_images, disp_encoded, disp_predicted, enc_shape=(8,4)):\n",
        "  '''Displays the input, encoded, and decoded output values.'''\n",
        "  plt.figure(figsize=(15, 5))\n",
        "  display_one_row(disp_input_images, 0, shape=(28,28,))\n",
        "  display_one_row(disp_encoded, 10, shape=enc_shape)\n",
        "  display_one_row(disp_predicted, 20, shape=(28,28,))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qtQyQRxRN_hH"
      },
      "source": [
        "# take 1 batch of the dataset\n",
        "test_dataset = test_dataset.take(1)\n",
        "\n",
        "# take the input images and put them in a list\n",
        "output_samples = []\n",
        "for input_image, image in tfds.as_numpy(test_dataset):\n",
        "      output_samples = input_image\n",
        "\n",
        "# pick 10 random numbers to be used as indices to the list above\n",
        "idxs = np.random.choice(BATCH_SIZE, size=10)\n",
        "\n",
        "# get the encoder output\n",
        "encoded_predicted = encoder_model.predict(test_dataset)\n",
        "\n",
        "# get a prediction for the test batch\n",
        "simple_predicted = autoencoder_model.predict(test_dataset)\n",
        "\n",
        "# display the 10 samples, encodings and decoded values!\n",
        "display_results(output_samples[idxs], encoded_predicted[idxs], simple_predicted[idxs])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for input_image, image in tfds.as_numpy(test_dataset):\n",
        "  print(  np.all( np.equal(input_image, image)  ) )\n"
      ],
      "metadata": {
        "id": "M5UsS366gjkK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "OrvgEfGQhw9_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}