chapter12_part05_gans.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "This is a companion notebook for the book [Deep Learning with Python, Second Edition](https://www.manning.com/books/deep-learning-with-python-second-edition?a_aid=keras&a_bid=76564dff). For readability, it only contains runnable code blocks and section titles, and omits everything else in the book: text paragraphs, figures, and pseudocode.\n\n**If you want to be able to follow what's going on, I recommend reading the notebook side by side with your copy of the book.**\n\nThis notebook was generated for TensorFlow 2.6."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Introduction to generative adversarial networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### A schematic GAN implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### A bag of tricks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Getting our hands on the CelebA dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Getting the CelebA data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "!mkdir celeba_gan\n",
    "!gdown --id 1O7m1010EJjLE5QxLZiM9Fpjs7Oj6e684 -O celeba_gan/data.zip\n",
    "!unzip -qq celeba_gan/data.zip -d celeba_gan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Creating a Dataset from a directory of images**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from tensorflow import keras\n",
    "dataset = keras.utils_dataset_from_directory(\n",
    "    \"celeba_gan\",\n",
    "    label_mode=None,\n",
    "    image_size=(64, 64),\n",
    "    batch_size=32,\n",
    "    smart_resize=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Rescaling the images**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "dataset = dataset.map(lambda x: x / 255.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Displaying the first image**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "for x in dataset:\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow((x.numpy() * 255).astype(\"int32\")[0])\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### The discriminator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**The GAN discriminator network**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "\n",
    "discriminator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(64, 64, 3)),\n",
    "        layers.Conv2D(64, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(128, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(128, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Flatten(),\n",
    "        layers.Dropout(0.2),\n",
    "        layers.Dense(1, activation=\"sigmoid\"),\n",
    "    ],\n",
    "    name=\"discriminator\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "discriminator.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### The generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**GAN generator network**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "latent_dim = 128\n",
    "\n",
    "generator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(latent_dim,)),\n",
    "        layers.Dense(8 * 8 * 128),\n",
    "        layers.Reshape((8, 8, 128)),\n",
    "        layers.Conv2DTranspose(128, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2DTranspose(256, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2DTranspose(512, kernel_size=4, strides=2, padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(3, kernel_size=5, padding=\"same\", activation=\"sigmoid\"),\n",
    "    ],\n",
    "    name=\"generator\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "generator.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### The adversarial network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**The GAN Model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "class GAN(keras.Model):\n",
    "    def __init__(self, discriminator, generator, latent_dim):\n",
    "        super().__init__()\n",
    "        self.discriminator = discriminator\n",
    "        self.generator = generator\n",
    "        self.latent_dim = latent_dim\n",
    "        self.d_loss_metric = keras.metrics.Mean(name=\"d_loss\")\n",
    "        self.g_loss_metric = keras.metrics.Mean(name=\"g_loss\")\n",
    "\n",
    "    def compile(self, d_optimizer, g_optimizer, loss_fn):\n",
    "        super(GAN, self).compile()\n",
    "        self.d_optimizer = d_optimizer\n",
    "        self.g_optimizer = g_optimizer\n",
    "        self.loss_fn = loss_fn\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [self.d_loss_metric, self.g_loss_metric]\n",
    "\n",
    "    def train_step(self, real_images):\n",
    "        batch_size = tf.shape(real_images)[0]\n",
    "        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n",
    "        generated_images = self.generator(random_latent_vectors)\n",
    "        combined_images = tf.concat([generated_images, real_images], axis=0)\n",
    "        labels = tf.concat(\n",
    "            [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0\n",
    "        )\n",
    "        labels += 0.05 * tf.random.uniform(tf.shape(labels))\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(combined_images)\n",
    "            d_loss = self.loss_fn(labels, predictions)\n",
    "        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)\n",
    "        self.d_optimizer.apply_gradients(\n",
    "            zip(grads, self.discriminator.trainable_weights)\n",
    "        )\n",
    "\n",
    "        random_latent_vectors = tf.random.normal(\n",
    "          shape=(batch_size, self.latent_dim))\n",
    "\n",
    "        misleading_labels = tf.zeros((batch_size, 1))\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(self.generator(random_latent_vectors))\n",
    "            g_loss = self.loss_fn(misleading_labels, predictions)\n",
    "        grads = tape.gradient(g_loss, self.generator.trainable_weights)\n",
    "        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))\n",
    "\n",
    "        self.d_loss_metric.update_state(d_loss)\n",
    "        self.g_loss_metric.update_state(g_loss)\n",
    "        return {\"d_loss\": self.d_loss_metric.result(), \"g_loss\": self.g_loss_metric.result()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**A callback to sample generated images during training**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "class GANMonitor(keras.callbacks.Callback):\n",
    "    def __init__(self, num_img=3, latent_dim=128):\n",
    "        self.num_img = num_img\n",
    "        self.latent_dim = latent_dim\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        random_latent_vectors = tf.random.normal(shape=(self.num_img, self.latent_dim))\n",
    "        generated_images = self.model.generator(random_latent_vectors)\n",
    "        generated_images *= 255\n",
    "        generated_images.numpy()\n",
    "        for i in range(self.num_img):\n",
    "            img = keras.utils.array_to_img(generated_images[i])\n",
    "            img.save(f\"generated_img_{epoch:03d}_{i}.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "**Compiling and training the GAN**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab_type": "code"
   },
   "outputs": [],
   "source": [
    "epochs = 100\n",
    "\n",
    "gan = GAN(discriminator=discriminator, generator=generator, latent_dim=latent_dim)\n",
    "gan.compile(\n",
    "    d_optimizer=keras.optimizers.Adam(learning_rate=0.0001),\n",
    "    g_optimizer=keras.optimizers.Adam(learning_rate=0.0001),\n",
    "    loss_fn=keras.losses.BinaryCrossentropy(),\n",
    ")\n",
    "\n",
    "gan.fit(\n",
    "    dataset, epochs=epochs, callbacks=[GANMonitor(num_img=10, latent_dim=latent_dim)]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "### Wrapping up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text"
   },
   "source": [
    "## Chapter summary"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "chapter12_part05_gans.i",
   "private_outputs": false,
   "provenance": [],
   "toc_visible": true
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}