8.5-introduction-to-gans.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras\n",
    "keras.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to generative adversarial networks\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A schematic GAN implementation\n",
    "\n",
    "\n",
    "In what follows, we explain how to implement a GAN in Keras, in its barest form -- since GANs are quite advanced, diving deeply into the \n",
    "technical details would be out of scope for us. Our specific implementation will be a deep convolutional GAN, or DCGAN: a GAN where the \n",
    "generator and discriminator are deep convnets. In particular, it leverages a `Conv2DTranspose` layer for image upsampling in the generator.\n",
    "\n",
    "We will train our GAN on images from CIFAR10, a dataset of 50,000 32x32 RGB images belong to 10 classes (5,000 images per class). To make \n",
    "things even easier, we will only use images belonging to the class \"frog\".\n",
    "\n",
    "Schematically, our GAN looks like this:\n",
    "\n",
    "* A `generator` network maps vectors of shape `(latent_dim,)` to images of shape `(32, 32, 3)`.\n",
    "* A `discriminator` network maps images of shape (32, 32, 3) to a binary score estimating the probability that the image is real.\n",
    "* A `gan` network chains the generator and the discriminator together: `gan(x) = discriminator(generator(x))`. Thus this `gan` network maps \n",
    "latent space vectors to the discriminator's assessment of the realism of these latent vectors as decoded by the generator.\n",
    "* We train the discriminator using examples of real and fake images along with \"real\"/\"fake\" labels, as we would train any regular image \n",
    "classification model.\n",
    "* To train the generator, we use the gradients of the generator's weights with regard to the loss of the `gan` model. This means that, at \n",
    "every step, we move the weights of the generator in a direction that will make the discriminator more likely to classify as \"real\" the \n",
    "images decoded by the generator. I.e. we train the generator to fool the discriminator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A bag of tricks\n",
    "\n",
    "\n",
    "Training GANs and tuning GAN implementations is notoriously difficult. There are a number of known \"tricks\" that one should keep in mind. \n",
    "Like most things in deep learning, it is more alchemy than science: these tricks are really just heuristics, not theory-backed guidelines. \n",
    "They are backed by some level of intuitive understanding of the phenomenon at hand, and they are known to work well empirically, albeit not \n",
    "necessarily in every context.\n",
    "\n",
    "Here are a few of the tricks that we leverage in our own implementation of a GAN generator and discriminator below. It is not an exhaustive \n",
    "list of GAN-related tricks; you will find many more across the GAN literature.\n",
    "\n",
    "* We use `tanh` as the last activation in the generator, instead of `sigmoid`, which would be more commonly found in other types of models.\n",
    "* We sample points from the latent space using a _normal distribution_ (Gaussian distribution), not a uniform distribution.\n",
    "* Stochasticity is good to induce robustness. Since GAN training results in a dynamic equilibrium, GANs are likely to get \"stuck\" in all sorts of ways. \n",
    "Introducing randomness during training helps prevent this. We introduce randomness in two ways: 1) we use dropout in the discriminator, 2) \n",
    "we add some random noise to the labels for the discriminator.\n",
    "* Sparse gradients can hinder GAN training. In deep learning, sparsity is often a desirable property, but not in GANs. There are two things \n",
    "that can induce gradient sparsity: 1) max pooling operations, 2) ReLU activations. Instead of max pooling, we recommend using strided \n",
    "convolutions for downsampling, and we recommend using a `LeakyReLU` layer instead of a ReLU activation. It is similar to ReLU but it \n",
    "relaxes sparsity constraints by allowing small negative activation values.\n",
    "* In generated images, it is common to see \"checkerboard artifacts\" caused by unequal coverage of the pixel space in the generator. To fix \n",
    "this, we use a kernel size that is divisible by the stride size, whenever we use a strided `Conv2DTranpose` or `Conv2D` in both the \n",
    "generator and discriminator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The generator\n",
    "\n",
    "\n",
    "First, we develop a `generator` model, which turns a vector (from the latent space -- during training it will sampled at random) into a \n",
    "candidate image. One of the many issues that commonly arise with GANs is that the generator gets stuck with generated images that look like \n",
    "noise. A possible solution is to use dropout on both the discriminator and generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras\n",
    "from keras import layers\n",
    "import numpy as np\n",
    "\n",
    "latent_dim = 32\n",
    "height = 32\n",
    "width = 32\n",
    "channels = 3\n",
    "\n",
    "generator_input = keras.Input(shape=(latent_dim,))\n",
    "\n",
    "# First, transform the input into a 16x16 128-channels feature map\n",
    "x = layers.Dense(128 * 16 * 16)(generator_input)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Reshape((16, 16, 128))(x)\n",
    "\n",
    "# Then, add a convolution layer\n",
    "x = layers.Conv2D(256, 5, padding='same')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "# Upsample to 32x32\n",
    "x = layers.Conv2DTranspose(256, 4, strides=2, padding='same')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "# Few more conv layers\n",
    "x = layers.Conv2D(256, 5, padding='same')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Conv2D(256, 5, padding='same')(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "\n",
    "# Produce a 32x32 1-channel feature map\n",
    "x = layers.Conv2D(channels, 7, activation='tanh', padding='same')(x)\n",
    "generator = keras.models.Model(generator_input, x)\n",
    "generator.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The discriminator\n",
    "\n",
    "\n",
    "Then, we develop a `discriminator` model, that takes as input a candidate image (real or synthetic) and classifies it into one of two \n",
    "classes, either \"generated image\" or \"real image that comes from the training set\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "discriminator_input = layers.Input(shape=(height, width, channels))\n",
    "x = layers.Conv2D(128, 3)(discriminator_input)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Conv2D(128, 4, strides=2)(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Conv2D(128, 4, strides=2)(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Conv2D(128, 4, strides=2)(x)\n",
    "x = layers.LeakyReLU()(x)\n",
    "x = layers.Flatten()(x)\n",
    "\n",
    "# One dropout layer - important trick!\n",
    "x = layers.Dropout(0.4)(x)\n",
    "\n",
    "# Classification layer\n",
    "x = layers.Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "discriminator = keras.models.Model(discriminator_input, x)\n",
    "discriminator.summary()\n",
    "\n",
    "# To stabilize training, we use learning rate decay\n",
    "# and gradient clipping (by value) in the optimizer.\n",
    "discriminator_optimizer = keras.optimizers.RMSprop(lr=0.0008, clipvalue=1.0, decay=1e-8)\n",
    "discriminator.compile(optimizer=discriminator_optimizer, loss='binary_crossentropy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The adversarial network\n",
    "\n",
    "Finally, we setup the GAN, which chains the generator and the discriminator. This is the model that, when trained, will move the generator \n",
    "in a direction that improves its ability to fool the discriminator. This model turns latent space points into a classification decision, \n",
    "\"fake\" or \"real\", and it is meant to be trained with labels that are always \"these are real images\". So training `gan` will updates the \n",
    "weights of `generator` in a way that makes `discriminator` more likely to predict \"real\" when looking at fake images. Very importantly, we \n",
    "set the discriminator to be frozen during training (non-trainable): its weights will not be updated when training `gan`. If the \n",
    "discriminator weights could be updated during this process, then we would be training the discriminator to always predict \"real\", which is \n",
    "not what we want!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set discriminator weights to non-trainable\n",
    "# (will only apply to the `gan` model)\n",
    "discriminator.trainable = False\n",
    "\n",
    "gan_input = keras.Input(shape=(latent_dim,))\n",
    "gan_output = discriminator(generator(gan_input))\n",
    "gan = keras.models.Model(gan_input, gan_output)\n",
    "\n",
    "gan_optimizer = keras.optimizers.RMSprop(lr=0.0004, clipvalue=1.0, decay=1e-8)\n",
    "gan.compile(optimizer=gan_optimizer, loss='binary_crossentropy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to train your DCGAN\n",
    "\n",
    "Now we can start training. To recapitulate, this is schematically what the training loop looks like:\n",
    "\n",
    "```\n",
    "for each epoch:\n",
    "    * Draw random points in the latent space (random noise).\n",
    "    * Generate images with `generator` using this random noise.\n",
    "    * Mix the generated images with real ones.\n",
    "    * Train `discriminator` using these mixed images, with corresponding targets, either \"real\" (for the real images) or \"fake\" (for the generated images).\n",
    "    * Draw new random points in the latent space.\n",
    "    * Train `gan` using these random vectors, with targets that all say \"these are real images\". This will update the weights of the generator (only, since discriminator is frozen inside `gan`) to move them towards getting the discriminator to predict \"these are real images\" for generated images, i.e. this trains the generator to fool the discriminator.\n",
    "```\n",
    "\n",
    "Let's implement it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from keras.preprocessing import image\n",
    "\n",
    "# Load CIFAR10 data\n",
    "(x_train, y_train), (_, _) = keras.datasets.cifar10.load_data()\n",
    "\n",
    "# Select frog images (class 6)\n",
    "x_train = x_train[y_train.flatten() == 6]\n",
    "\n",
    "# Normalize data\n",
    "x_train = x_train.reshape(\n",
    "    (x_train.shape[0],) + (height, width, channels)).astype('float32') / 255.\n",
    "\n",
    "iterations = 10000\n",
    "batch_size = 20\n",
    "save_dir = './data/gan_images/'\n",
    "\n",
    "# Start training loop\n",
    "start = 0\n",
    "for step in range(iterations):\n",
    "    # Sample random points in the latent space\n",
    "    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))\n",
    "\n",
    "    # Decode them to fake images\n",
    "    generated_images = generator.predict(random_latent_vectors)\n",
    "\n",
    "    # Combine them with real images\n",
    "    stop = start + batch_size\n",
    "    real_images = x_train[start: stop]\n",
    "    combined_images = np.concatenate([generated_images, real_images])\n",
    "\n",
    "    # Assemble labels discriminating real from fake images\n",
    "    labels = np.concatenate([np.ones((batch_size, 1)),\n",
    "                             np.zeros((batch_size, 1))])\n",
    "    # Add random noise to the labels - important trick!\n",
    "    labels += 0.05 * np.random.random(labels.shape)\n",
    "\n",
    "    # Train the discriminator\n",
    "    d_loss = discriminator.train_on_batch(combined_images, labels)\n",
    "\n",
    "    # sample random points in the latent space\n",
    "    random_latent_vectors = np.random.normal(size=(batch_size, latent_dim))\n",
    "\n",
    "    # Assemble labels that say \"all real images\"\n",
    "    misleading_targets = np.zeros((batch_size, 1))\n",
    "\n",
    "    # Train the generator (via the gan model,\n",
    "    # where the discriminator weights are frozen)\n",
    "    a_loss = gan.train_on_batch(random_latent_vectors, misleading_targets)\n",
    "    \n",
    "    start += batch_size\n",
    "    if start > len(x_train) - batch_size:\n",
    "      start = 0\n",
    "\n",
    "    # Occasionally save / plot\n",
    "    if step % 100 == 0:\n",
    "        # Save model weights\n",
    "        gan.save_weights('gan.h5')\n",
    "\n",
    "        # Print metrics\n",
    "        print('discriminator loss at step %s: %s' % (step, d_loss))\n",
    "        print('adversarial loss at step %s: %s' % (step, a_loss))\n",
    "\n",
    "        # Save one generated image\n",
    "        img = image.array_to_img(generated_images[0] * 255., scale=False)\n",
    "        img.save(os.path.join(save_dir, 'generated_frog' + str(step) + '.png'))\n",
    "\n",
    "        # Save one real image, for comparison\n",
    "        img = image.array_to_img(real_images[0] * 255., scale=False)\n",
    "        img.save(os.path.join(save_dir, 'real_frog' + str(step) + '.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Let's display a few of our fake images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Sample random points in the latent space\n",
    "random_latent_vectors = np.random.normal(size=(10, latent_dim))\n",
    "\n",
    "# Decode them to fake images\n",
    "generated_images = generator.predict(random_latent_vectors)\n",
    "\n",
    "for i in range(generated_images.shape[0]):\n",
    "    img = image.array_to_img(generated_images[i] * 255., scale=False)\n",
    "    plt.figure()\n",
    "    plt.imshow(img)\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Froggy with some pixellated artifacts."
   ]
  }
 ],
 "metadata": {
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}