path: root/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0TD5ZrvEMbhZ"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\").\n",
        "\n",
        "# DCGAN: An example with tf.keras and eager\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\u003ctd\u003e\n",
        "\u003ca target=\"_blank\"  href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb\"\u003e\n",
        "    \u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e  \n",
        "\u003c/td\u003e\u003ctd\u003e\n",
        "\u003ca target=\"_blank\"  href=\"https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb\"\u003e\u003cimg width=32px src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\u003c/td\u003e\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ITZuApL56Mny"
      },
      "source": [
        "This notebook demonstrates how to generate images of handwritten digits using [tf.keras](https://www.tensorflow.org/programmers_guide/keras) and [eager execution](https://www.tensorflow.org/programmers_guide/eager). To do so, we use Deep Convolutional Generative Adverserial Networks ([DCGAN](https://arxiv.org/pdf/1511.06434.pdf)).\n",
        "\n",
        "This model takes about ~30 seconds per epoch (using tf.contrib.eager.defun to create graph functions) to train on a single Tesla K80 on Colab, as of July 2018.\n",
        "\n",
        "Below is the output generated after training the generator and discriminator models for 150 epochs.\n",
        "\n",
        "![sample output](https://tensorflow.org/images/gan/dcgan.gif)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "u_2z-B3piVsw"
      },
      "outputs": [],
      "source": [
        "# to generate gifs\n",
        "!pip install imageio"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "e1_Y75QXJS6h"
      },
      "source": [
        "## Import TensorFlow and enable eager execution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "YfIk2es3hJEd"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "# Import TensorFlow \u003e= 1.9 and enable eager execution\n",
        "import tensorflow as tf\n",
        "tf.enable_eager_execution()\n",
        "\n",
        "import os\n",
        "import time\n",
        "import numpy as np\n",
        "import glob\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL\n",
        "import imageio\n",
        "from IPython import display"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "iYn4MdZnKCey"
      },
      "source": [
        "## Load the dataset\n",
        "\n",
        "We are going to use the MNIST dataset to train the generator and the discriminator. The generator will then generate handwritten digits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "a4fYMGxGhrna"
      },
      "outputs": [],
      "source": [
        "(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "NFC2ghIdiZYE"
      },
      "outputs": [],
      "source": [
        "train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')\n",
        "# We are normalizing the images to the range of [-1, 1]\n",
        "train_images = (train_images - 127.5) / 127.5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "S4PIDhoDLbsZ"
      },
      "outputs": [],
      "source": [
        "BUFFER_SIZE = 60000\n",
        "BATCH_SIZE = 256"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PIGN6ouoQxt3"
      },
      "source": [
        "## Use tf.data to create batches and shuffle the dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "-yKCCQOoJ7cn"
      },
      "outputs": [],
      "source": [
        "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "THY-sZMiQ4UV"
      },
      "source": [
        "## Write the generator and discriminator models\n",
        "\n",
        "* **Generator** \n",
        "  * It is responsible for **creating convincing images that are good enough to fool the discriminator**.\n",
        "  * It consists of Conv2DTranspose (Upsampling) layers. We start with a fully connected layer and upsample the image 2 times so as to reach the desired image size (mnist image size) which is (28, 28, 1). \n",
        "  * We use **leaky relu** activation except for the **last layer** which uses **tanh** activation.\n",
        "  \n",
        "* **Discriminator**\n",
        "  * **The discriminator is responsible for classifying the fake images from the real images.**\n",
        "  * In other words, the discriminator is given generated images (from the generator) and the real MNIST images. The job of the discriminator is to classify these images into fake (generated) and real (MNIST images).\n",
        "  * **Basically the generator should be good enough to fool the discriminator that the generated images are real**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "VGLbvBEmjK0a"
      },
      "outputs": [],
      "source": [
        "class Generator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Generator, self).__init__()\n",
        "    self.fc1 = tf.keras.layers.Dense(7*7*64, use_bias=False)\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization()\n",
        "    \n",
        "    self.conv1 = tf.keras.layers.Conv2DTranspose(64, (5, 5), strides=(1, 1), padding='same', use_bias=False)\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization()\n",
        "    \n",
        "    self.conv2 = tf.keras.layers.Conv2DTranspose(32, (5, 5), strides=(2, 2), padding='same', use_bias=False)\n",
        "    self.batchnorm3 = tf.keras.layers.BatchNormalization()\n",
        "    \n",
        "    self.conv3 = tf.keras.layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False)\n",
        "\n",
        "  def call(self, x, training=True):\n",
        "    x = self.fc1(x)\n",
        "    x = self.batchnorm1(x, training=training)\n",
        "    x = tf.nn.relu(x)\n",
        "\n",
        "    x = tf.reshape(x, shape=(-1, 7, 7, 64))\n",
        "\n",
        "    x = self.conv1(x)\n",
        "    x = self.batchnorm2(x, training=training)\n",
        "    x = tf.nn.relu(x)\n",
        "\n",
        "    x = self.conv2(x)\n",
        "    x = self.batchnorm3(x, training=training)\n",
        "    x = tf.nn.relu(x)\n",
        "\n",
        "    x = tf.nn.tanh(self.conv3(x))  \n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "bkOfJxk5j5Hi"
      },
      "outputs": [],
      "source": [
        "class Discriminator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Discriminator, self).__init__()\n",
        "    self.conv1 = tf.keras.layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same')\n",
        "    self.conv2 = tf.keras.layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same')\n",
        "    self.dropout = tf.keras.layers.Dropout(0.3)\n",
        "    self.flatten = tf.keras.layers.Flatten()\n",
        "    self.fc1 = tf.keras.layers.Dense(1)\n",
        "\n",
        "  def call(self, x, training=True):\n",
        "    x = tf.nn.leaky_relu(self.conv1(x))\n",
        "    x = self.dropout(x, training=training)\n",
        "    x = tf.nn.leaky_relu(self.conv2(x))\n",
        "    x = self.dropout(x, training=training)\n",
        "    x = self.flatten(x)\n",
        "    x = self.fc1(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "gDkA05NE6QMs"
      },
      "outputs": [],
      "source": [
        "generator = Generator()\n",
        "discriminator = Discriminator()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "k1HpMSLImuRi"
      },
      "outputs": [],
      "source": [
        "# Defun gives 10 secs/epoch performance boost\n",
        "generator.call = tf.contrib.eager.defun(generator.call)\n",
        "discriminator.call = tf.contrib.eager.defun(discriminator.call)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0FMYgY_mPfTi"
      },
      "source": [
        "## Define the loss functions and the optimizer\n",
        "\n",
        "* **Discriminator loss**\n",
        "  * The discriminator loss function takes 2 inputs; **real images, generated images**\n",
        "  * real_loss is a sigmoid cross entropy loss of the **real images** and an **array of ones (since these are the real images)**\n",
        "  * generated_loss is a sigmoid cross entropy loss of the **generated images** and an **array of zeros (since these are the fake images)**\n",
        "  * Then the total_loss is the sum of real_loss and the generated_loss\n",
        "  \n",
        "* **Generator loss**\n",
        "  * It is a sigmoid cross entropy loss of the generated images and an **array of ones**\n",
        "  \n",
        "\n",
        "* The discriminator and the generator optimizers are different since we will train them separately."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "wkMNfBWlT-PV"
      },
      "outputs": [],
      "source": [
        "def discriminator_loss(real_output, generated_output):\n",
        "    # [1,1,...,1] with real output since it is true and we want\n",
        "    # our generated examples to look like it\n",
        "    real_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=tf.ones_like(real_output), logits=real_output)\n",
        "\n",
        "    # [0,0,...,0] with generated images since they are fake\n",
        "    generated_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels=tf.zeros_like(generated_output), logits=generated_output)\n",
        "\n",
        "    total_loss = real_loss + generated_loss\n",
        "\n",
        "    return total_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "90BIcCKcDMxz"
      },
      "outputs": [],
      "source": [
        "def generator_loss(generated_output):\n",
        "    return tf.losses.sigmoid_cross_entropy(tf.ones_like(generated_output), generated_output)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "iWCn_PVdEJZ7"
      },
      "outputs": [],
      "source": [
        "discriminator_optimizer = tf.train.AdamOptimizer(1e-4)\n",
        "generator_optimizer = tf.train.AdamOptimizer(1e-4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Rw1fkAczTQYh"
      },
      "source": [
        "## Training\n",
        "\n",
        "* We start by iterating over the dataset\n",
        "* The generator is given **noise as an input** which when passed through the generator model will output a image looking like a handwritten digit\n",
        "* The discriminator is given the **real MNIST images as well as the generated images (from the generator)**.\n",
        "* Next, we calculate the generator and the discriminator loss.\n",
        "* Then, we calculate the gradients of loss with respect to both the generator and the discriminator variables (inputs) and apply those to the optimizer.\n",
        "\n",
        "## Generate Images\n",
        "\n",
        "* After training, its time to generate some images!\n",
        "* We start by creating noise array as an input to the generator\n",
        "* The generator will then convert the noise into handwritten images.\n",
        "* Last step is to plot the predictions and **voila!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "NS2GWywBbAWo"
      },
      "outputs": [],
      "source": [
        "EPOCHS = 150\n",
        "noise_dim = 100\n",
        "num_examples_to_generate = 100\n",
        "\n",
        "# keeping the random vector constant for generation (prediction) so\n",
        "# it will be easier to see the improvement of the gan.\n",
        "random_vector_for_generation = tf.random_normal([num_examples_to_generate,\n",
        "                                                 noise_dim])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "RmdVsmvhPxyy"
      },
      "outputs": [],
      "source": [
        "def generate_and_save_images(model, epoch, test_input):\n",
        "  # make sure the training parameter is set to False because we\n",
        "  # don't want to train the batchnorm layer when doing inference.\n",
        "  predictions = model(test_input, training=False)\n",
        "\n",
        "  fig = plt.figure(figsize=(10,10))\n",
        "  \n",
        "  for i in range(predictions.shape[0]):\n",
        "      plt.subplot(10, 10, i+1)\n",
        "      plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')\n",
        "      plt.axis('off')\n",
        "        \n",
        "  # tight_layout minimizes the overlap between 2 sub-plots\n",
        "  plt.tight_layout()\n",
        "  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "2M7LmLtGEMQJ"
      },
      "outputs": [],
      "source": [
        "def train(dataset, epochs, noise_dim):  \n",
        "  for epoch in range(epochs):\n",
        "    start = time.time()\n",
        "    \n",
        "    for images in dataset:\n",
        "      # generating noise from a uniform distribution\n",
        "      noise = tf.random_normal([BATCH_SIZE, noise_dim])\n",
        "      \n",
        "      with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
        "        generated_images = generator(noise, training=True)\n",
        "      \n",
        "        real_output = discriminator(images, training=True)\n",
        "        generated_output = discriminator(generated_images, training=True)\n",
        "        \n",
        "        gen_loss = generator_loss(generated_output)\n",
        "        disc_loss = discriminator_loss(real_output, generated_output)\n",
        "        \n",
        "      gradients_of_generator = gen_tape.gradient(gen_loss, generator.variables)\n",
        "      gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.variables)\n",
        "      \n",
        "      generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.variables))\n",
        "      discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.variables))\n",
        "\n",
        "      \n",
        "    if epoch % 10 == 0:\n",
        "      display.clear_output(wait=True)\n",
        "      generate_and_save_images(generator,\n",
        "                               epoch + 1,\n",
        "                               random_vector_for_generation)\n",
        "\n",
        "    print ('Time taken for epoch {} is {} sec'.format(epoch + 1,\n",
        "                                                      time.time()-start))\n",
        "  # generating after the final epoch\n",
        "  generate_and_save_images(generator,\n",
        "                           epochs,\n",
        "                           random_vector_for_generation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "Ly3UN0SLLY2l"
      },
      "outputs": [],
      "source": [
        "train(train_dataset, EPOCHS, noise_dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "P4M_vIbUi7c0"
      },
      "source": [
        "# Display an image using the epoch number"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "WfO5wCdclHGL"
      },
      "outputs": [],
      "source": [
        "def display_image(epoch_no):\n",
        "  plt.figure(figsize=(15,15))\n",
        "  plt.imshow(np.array(PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))))\n",
        "  plt.axis('off')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "5x3q9_Oe5q0A"
      },
      "outputs": [],
      "source": [
        "display_image(EPOCHS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NywiH3nL8guF"
      },
      "source": [
        "## Generate a GIF of all the saved images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xmO0Dmu2WICn"
      },
      "source": [
        "\u003c!-- TODO(markdaoust): Remove the hack when Ipython version is updated --\u003e\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "IGKQgENQ8lEI"
      },
      "outputs": [],
      "source": [
        "with imageio.get_writer('dcgan.gif', mode='I') as writer:\n",
        "  filenames = glob.glob('image*.png')\n",
        "  filenames = sorted(filenames)\n",
        "  for filename in filenames:\n",
        "    image = imageio.imread(filename)\n",
        "    writer.append_data(image)\n",
        "  # this is a hack to display the gif inside the notebook\n",
        "  os.system('mv dcgan.gif dcgan.gif.png')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "uV0yiKpzNP1b"
      },
      "outputs": [],
      "source": [
        "display.Image(filename=\"dcgan.gif.png\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "autoexec": {
            "startup": false,
            "wait_interval": 0
          }
        },
        "colab_type": "code",
        "id": "4UJjSnIMOzOJ"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "default_view": {},
      "name": "dcgan.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1eb0NOTQapkYs3X0v-zL1x5_LFKgDISnp",
          "timestamp": 1527173385672
        }
      ],
      "toc_visible": true,
      "version": "0.3.2",
      "views": {}
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}