fashion_mnist_gan_toronto_ai_workshop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "exNoSECDuU9E"
   },
   "source": [
    "# * Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8EDFk9tXrv_V"
   },
   "source": [
    "Notebook link: https://goo.gl/bB2m1T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yWCHiDFDuLrb"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.client import device_lib\n",
    "\n",
    "# enable eager execution\n",
    "tf.enable_eager_execution()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 485
    },
    "colab_type": "code",
    "id": "HinnV9Sdugax",
    "outputId": "e410a4e2-6636-43fc-f723-957fe29f08e0"
   },
   "outputs": [],
   "source": [
    "# check that TF can detect GPU\n",
    "device_lib.list_local_devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6twNlrZ4vFIs"
   },
   "source": [
    "# 1 Fashion MNIST Dataset\n",
    "\n",
    "Load train/test datasets from API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "colab_type": "code",
    "id": "jKD66jrGup9S",
    "outputId": "d0b16132-7bc5-4e49-fad8-40f9897e2603"
   },
   "outputs": [],
   "source": [
    "# download the fashion mnist dataset\n",
    "fashion_mnist = tf.keras.datasets.fashion_mnist\n",
    "(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 126
    },
    "colab_type": "code",
    "id": "d7Lqx7Znvhpe",
    "outputId": "6de05012-c2ee-4279-f267-0ac0d0b277d3"
   },
   "outputs": [],
   "source": [
    "# examine dataset format\n",
    "print(\"train images: type: {}; shape: {}\".format(type(train_images), train_images.shape))\n",
    "print(\"test images: type:  {}; shape: {}\".format(type(test_images), test_images.shape))\n",
    "\n",
    "# examine class frequency\n",
    "\n",
    "print(\"train labels: {}\".format([(a,b) for a, b in zip(np.unique(train_labels), np.bincount(train_labels))]))\n",
    "print(\"test labels:  {}\".format([(a,b) for a, b in zip(np.unique(test_labels), np.bincount(test_labels))]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 596
    },
    "colab_type": "code",
    "id": "dRlgXzyUvlei",
    "outputId": "cfbe3287-d9b2-471d-f0ed-9b21dda1a059"
   },
   "outputs": [],
   "source": [
    "# visualize some of the images by randomly sampling a few images from each class\n",
    "img_samples = []\n",
    "for cls in list(np.unique(train_labels)):\n",
    "    indices = np.where(train_labels==cls)[0]    \n",
    "    samples = np.random.choice(a=indices, size=10, replace=False)\n",
    "    img_samples.append(train_images[samples,:,:])\n",
    "img_samples = tf.constant(np.concatenate(img_samples))\n",
    "img_samples = tf.expand_dims(img_samples, 3)\n",
    "\n",
    "# generate a 10 x 10 image grid\n",
    "img_grid = tf.contrib.gan.eval.image_grid(\n",
    "                input_tensor=img_samples,\n",
    "                grid_shape=(10,10),\n",
    "                image_shape=(28,28),\n",
    "                num_channels=1\n",
    "            )\n",
    "\n",
    "# plot the image grid\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(img_grid[0,:,:,0], cmap='gray')\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-sg8Pu7C45T9"
   },
   "source": [
    "# 2 GAN Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YdhRhV9x-lNl"
   },
   "outputs": [],
   "source": [
    "# Import keras modules that we will be using to construct our models\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import (\n",
    "    Dense, ReLU, Activation, Flatten, Conv2D, LeakyReLU, \n",
    "    Conv2DTranspose, Input, Reshape, BatchNormalization\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "kn31vJh18lUP",
    "outputId": "303e120a-6d3d-449d-ac77-3c4856d64851"
   },
   "outputs": [],
   "source": [
    "# create an image batch for testing model functionalities\n",
    "images = train_images[:64,:,:]/128. - 1.\n",
    "images = tf.constant(images, dtype=tf.float32)\n",
    "images = tf.expand_dims(images, 3)\n",
    "images.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "JKXUd5cH_bS_"
   },
   "outputs": [],
   "source": [
    "# input image shape information\n",
    "H, W, C =28, 28, 1\n",
    "LATENT_DIM=16\n",
    "\n",
    "# alpha param for leaky ReLu\n",
    "ALPHA = 0.2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ALIt9LpW5zeO"
   },
   "source": [
    "### Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Q0lkSUEb2KwY"
   },
   "outputs": [],
   "source": [
    "def build_discriminator():\n",
    "    return Sequential([\n",
    "        # reshape input layer\n",
    "        Reshape((H, W, C), input_shape=(H, W)),\n",
    "        # 1st conv2D layer\n",
    "        Conv2D(filters=32, kernel_size=5, strides=2, padding='same'),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # 2nd conv2D layer\n",
    "        Conv2D(filters=64, kernel_size=5, strides=2, padding='same'),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # 3rd conv2D layer\n",
    "        Conv2D(filters=128, kernel_size=5, strides=2, padding='same'),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # Flatten conv2D output\n",
    "        Flatten(),\n",
    "        # Output layer\n",
    "        Dense(units=1),\n",
    "        Activation('tanh')\n",
    "    ])\n",
    "      \n",
    "      \n",
    "def compute_d_loss(d_logits_real, d_logits_fake):\n",
    "    \"\"\"\n",
    "    loss computation for the discriminator net\n",
    "    \"\"\"\n",
    "    # Hinge loss\n",
    "    real_loss = tf.reduce_mean(tf.nn.relu(1. - d_logits_real))\n",
    "    fake_loss = tf.reduce_mean(tf.nn.relu(1. + d_logits_fake))\n",
    "\n",
    "    return real_loss + fake_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 537
    },
    "colab_type": "code",
    "id": "80gORRNp9Uj2",
    "outputId": "6d2699f9-605f-431e-a161-367f2af847a7"
   },
   "outputs": [],
   "source": [
    "# unit test\n",
    "d_net = build_discriminator()\n",
    "print(d_net.summary())\n",
    "test = d_net(images)\n",
    "print('\\ndiscriminator loss: ', compute_d_loss(test, test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BSoEf6-k5jiC"
   },
   "source": [
    "### Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0s4e7BD6cRfZ"
   },
   "outputs": [],
   "source": [
    "def build_generator():\n",
    "    return Sequential([\n",
    "        # fully connected layer on latent vector\n",
    "        Dense(units=(H/4)*(W/4)*64, use_bias=False, input_shape=[LATENT_DIM]),\n",
    "        BatchNormalization(),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # reshape 1-D tensor into 3-D tensor\n",
    "        Reshape(((H/4), (W/4), 64)),\n",
    "        # 1st deconv operation\n",
    "        Conv2DTranspose(filters=64, kernel_size=5, strides=1, padding='same', use_bias=False),\n",
    "        BatchNormalization(),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # 2nd deconv operation\n",
    "        Conv2DTranspose(filters=32, kernel_size=5, strides=2, padding='same', use_bias=False),\n",
    "        BatchNormalization(),\n",
    "        LeakyReLU(ALPHA),\n",
    "        # 3rd deconv operation\n",
    "        Conv2DTranspose(filters=C, kernel_size=5, strides=2, padding='same', use_bias=False),\n",
    "        # output operation\n",
    "        Activation('tanh'),\n",
    "    ])\n",
    "\n",
    "      \n",
    "def compute_g_loss(d_logits_fake):\n",
    "    return - tf.reduce_mean(d_logits_fake)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 606
    },
    "colab_type": "code",
    "id": "e3gbdvQw5nAE",
    "outputId": "072d5ea0-8291-497a-8cee-61055b8ecbed"
   },
   "outputs": [],
   "source": [
    "# unit test\n",
    "g_net = build_generator()\n",
    "print(g_net.summary())\n",
    "z_test = tf.random_normal(shape=(64, LATENT_DIM), dtype='float32')\n",
    "test = g_net(z_test)\n",
    "test_logit = d_net.call(test)\n",
    "print('\\ndiscriminator loss on generated images: ', compute_g_loss(test_logit))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6WoSPJ4JfFI0"
   },
   "source": [
    "# 3 Model Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tSDNwqcvAfcv"
   },
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XSTBT10s_pYW"
   },
   "outputs": [],
   "source": [
    "# training configuration information\n",
    "SAMPLE_PER_N_STEPS=200\n",
    "\n",
    "# training batch size & epoch\n",
    "BUFFER_SIZE=1000\n",
    "EPOCHS = 10\n",
    "BATCH_SIZE = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9U9ggaN3CWp5"
   },
   "source": [
    "### Optimizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1VXOzuqRCZL0"
   },
   "outputs": [],
   "source": [
    "# learning rates\n",
    "G_LR = 0.0001\n",
    "D_LR = 0.0004\n",
    "\n",
    "# beta params for the Adam optimizer\n",
    "BETA1 = 0.0\n",
    "BETA2 = 0.999\n",
    "\n",
    "# generator optimizer\n",
    "g_optimizer = tf.train.AdamOptimizer(learning_rate=G_LR, \n",
    "                                     beta1=BETA1, \n",
    "                                     beta2=BETA2)\n",
    "\n",
    "# discriminator optimizer\n",
    "d_optimizer = tf.train.AdamOptimizer(learning_rate=D_LR, \n",
    "                                     beta1=BETA1, \n",
    "                                     beta2=BETA2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XycsWnnlN_iE"
   },
   "source": [
    "### Data Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vPpKDgJ7cXfd"
   },
   "outputs": [],
   "source": [
    "# combine the train/test datasets because we won't be needing the test dataset\n",
    "dataset = np.concatenate((train_images, test_images))\n",
    "\n",
    "# convert dataset to 4-D tensor object format (batch, height, width, channel)\n",
    "dataset = tf.constant(dataset, dtype=tf.float32)\n",
    "dataset = tf.expand_dims(dataset, 3)\n",
    "\n",
    "# create a tf.dataset object that will act as our input pipeline to feed data\n",
    "# to our models during training\n",
    "dataset = tf.data.Dataset.from_tensor_slices(dataset)\n",
    "\n",
    "# we will pre-program how the tf.dataset object will be feeding in data batches\n",
    "dataset = dataset.map(lambda x: (x / 128.) - 1., num_parallel_calls=4)\\\n",
    "                 .shuffle(BUFFER_SIZE)\\\n",
    "                 .repeat(EPOCHS)\\\n",
    "                 .batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "eoQuE-oiOPAy",
    "outputId": "49592def-f630-4c3f-f4dd-7e0d37a4ae1b"
   },
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PTtqINwTRIet"
   },
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 9217
    },
    "colab_type": "code",
    "id": "hCyi2dffRFlZ",
    "outputId": "c50818b6-22fa-4f27-eb23-5d6a4fcf7f34"
   },
   "outputs": [],
   "source": [
    "# create global training step tracker object\n",
    "global_step = tf.train.get_or_create_global_step()\n",
    "\n",
    "# keep track of losses\n",
    "d_losses, g_losses = [], []\n",
    "\n",
    "\n",
    "for real_img_batch in dataset:\n",
    "\n",
    "    # construct random normal z input to feed into generator\n",
    "    input_z = tf.random_normal(shape=(BATCH_SIZE, LATENT_DIM), dtype='float32')\n",
    "\n",
    "    # define gradient tapes to start recording computation operations\n",
    "    with tf.GradientTape() as g_tape, tf.GradientTape() as d_tape:\n",
    "      \n",
    "\n",
    "        # FORWARD RUN TO COMPUTE G/D-NET LOSSES\n",
    "        # ----------------------------------------------------------------------\n",
    "        # 1. run g_net with input_z to generate batch of fake images\n",
    "        g_fake_images = g_net(input_z, training=True)  \n",
    "\n",
    "        # 2. run d_net with the batch of generated fake images\n",
    "        d_logits_fake = d_net(g_fake_images, training=True)\n",
    "        \n",
    "        # 3. run d_net with a batch of real images from dataset\n",
    "        d_logits_real = d_net(real_img_batch, training=True)\n",
    "        \n",
    "        # 4. compute g_net losses with feedback from the d_net\n",
    "        g_loss = compute_g_loss(d_logits_fake)\n",
    "        g_losses.append(g_loss)\n",
    "        \n",
    "        # 5. compute d_net losses by revealing its logit values for both\n",
    "        #    real and fake images\n",
    "        d_loss = compute_d_loss(d_logits_real, d_logits_fake)\n",
    "        d_losses.append(d_loss)\n",
    "        \n",
    "                          \n",
    "        # UPDATE G/D-NET PARAMETERS\n",
    "        # ----------------------------------------------------------------------\n",
    "        # 1. get all learn-able G/D-net parameters (i.e. parameters to optimize)\n",
    "        d_variables = d_net.variables\n",
    "        g_variables = g_net.variables\n",
    "                \n",
    "        # 2. compute d(d_loss)/dx, d(g_loss)/dx\n",
    "        d_grads = d_tape.gradient(d_loss, d_variables)\n",
    "        g_grads = g_tape.gradient(g_loss, g_variables)\n",
    "\n",
    "        # 3. apply gradient updates to the parameters\n",
    "        d_optimizer.apply_gradients(zip(d_grads, d_variables),\n",
    "                                    global_step=global_step)\n",
    "        g_optimizer.apply_gradients(zip(g_grads, g_variables),\n",
    "                                    global_step=global_step)\n",
    "\n",
    "    \n",
    "    # EVERY NOW & THEN, DISPLAY OUTPUT TO TRACK TRAINING PROGRESS\n",
    "    # ----------------------------------------------------------------------\n",
    "    # get training step\n",
    "    step = global_step.numpy()\n",
    "    \n",
    "    # display losses every 100 steps\n",
    "    if step % 100==0:\n",
    "        print('training step {}: discriminator loss {}; generator loss {}'\\\n",
    "              .format(step, d_loss, g_loss))\n",
    "    \n",
    "    \n",
    "    # display sample images every SAMPLE_PER_N_STEPS\n",
    "    if step % SAMPLE_PER_N_STEPS==0:      \n",
    "        \n",
    "        # 1. create a batch of (100, LATENT_DIM) z-input tensor\n",
    "        eval_z = tf.random_normal(shape=(100, LATENT_DIM), dtype='float32')\n",
    "\n",
    "        # 2. generate images by using G-net\n",
    "        eval_img = g_net(eval_z, training=False)\n",
    "        \n",
    "        # 3. organize the images into a grid\n",
    "        img_grid = tf.contrib.gan.eval.image_grid(\n",
    "                        input_tensor=eval_img,\n",
    "                        grid_shape=(10,10),\n",
    "                        image_shape=(H, W),\n",
    "                        num_channels=C\n",
    "                    )\n",
    "                \n",
    "        # 4. plot the image grid\n",
    "        plt.figure(figsize=(10,10))\n",
    "        plt.imshow(img_grid[0,:,:,0], cmap='gray')\n",
    "        plt.axis('off')\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KiXUJQQdF59D"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "FMNIST GAN",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}