C4_W4_Lab_3_CelebA_GAN_Experiments.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "C4_W4_Lab_3_CelebA_GAN_Experiments.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmzwPwcssrVE"
      },
      "source": [
        "# Ungraded Lab: CelebA GAN Experiments\n",
        "\n",
        "This lab will demonstrate a GAN trained on the [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) dataset. This is a resource-intensive task so you will use a TPU and a distributed strategy to train the network. It will take 40 to 50 minutes to run the entire exercise. Afterwards, you will see a gif showing new faces generated by the trained model. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kUq5cRGbs17s"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lwhm4tTZZF21"
      },
      "source": [
        "# install tensorflow_addons\n",
        "!pip install -U tensorflow-addons"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "udWFxlOvYPgh"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow.keras.utils import plot_model\n",
        "from tensorflow.keras import layers\n",
        "import tensorflow_addons as tfa\n",
        "\n",
        "import os\n",
        "import zipfile\n",
        "import glob\n",
        "import urllib.request\n",
        "from enum import Enum\n",
        "from tqdm import tqdm\n",
        "from functools import partial\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "from PIL import Image\n",
        "from IPython.display import display\n",
        "from IPython.display import Image as IpyImage\n",
        "import imageio\n",
        "import cv2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zkpzuUi7vpsh"
      },
      "source": [
        "## Setup TPU\n",
        "\n",
        "You will use a TPU and its corresponding [distribution strategy](https://www.tensorflow.org/api_docs/python/tf/distribute/TPUStrategy) to speed up the training. We've provided the setup code and helper functions below. You might recognize some of these from taking Course 2 of this Specialization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rK_RTnYOYmKi"
      },
      "source": [
        "tpu_grpc_url = \"grpc://\" + os.environ[\"COLAB_TPU_ADDR\"]\n",
        "tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_grpc_url)\n",
        "tf.config.experimental_connect_to_cluster(tpu_cluster_resolver) \n",
        "tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)   \n",
        "strategy = tf.distribute.experimental.TPUStrategy(tpu_cluster_resolver)  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v9qrAyLDYpxJ"
      },
      "source": [
        "class Reduction(Enum):\n",
        "    NONE = 0\n",
        "    SUM = 1\n",
        "    MEAN = 2\n",
        "    CONCAT = 3\n",
        "\n",
        "def distributed(*reduction_flags):\n",
        "    def _decorator(fun):\n",
        "        def per_replica_reduction(z, flag):\n",
        "            if flag == Reduction.NONE:\n",
        "                return z\n",
        "            elif flag == Reduction.SUM:\n",
        "                return strategy.reduce(tf.distribute.ReduceOp.SUM, z, axis=None)\n",
        "            elif flag == Reduction.MEAN:\n",
        "                return strategy.reduce(tf.distribute.ReduceOp.MEAN, z, axis=None)\n",
        "            elif flag == Reduction.CONCAT:\n",
        "                z_list = strategy.experimental_local_results(z)\n",
        "                return tf.concat(z_list, axis=0)\n",
        "            else:\n",
        "                raise NotImplementedError()\n",
        "\n",
        "        @tf.function\n",
        "        def _decorated_fun(*args, **kwargs):\n",
        "            fun_result = strategy.run(fun, args=args, kwargs=kwargs)\n",
        "            if len(reduction_flags) == 0:\n",
        "                assert fun_result is None\n",
        "                return\n",
        "            elif len(reduction_flags) == 1:\n",
        "                assert type(fun_result) is not tuple and fun_redult is not None\n",
        "                return per_replica_reduction(fun_result, *reduction_flags)\n",
        "            else:\n",
        "                assert type(fun_result) is tuple\n",
        "                return tuple((per_replica_reduction(fr, rf) for fr, rf in zip(fun_result, reduction_flags)))\n",
        "        return _decorated_fun\n",
        "    return _decorator"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbhebu8EvsCz"
      },
      "source": [
        "## Download and Prepare the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oQ0-OdRiq10F"
      },
      "source": [
        "Next, you will fetch the celebrity faces dataset. We've hosted a copy of the data in a Google Drive but the filesize is around 1GB so it will take some time to download."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bqxdEl0wYa13"
      },
      "source": [
        "# make a data directory\n",
        "try:\n",
        "  os.mkdir('/tmp/celeb')\n",
        "except OSError:\n",
        "  pass\n",
        "\n",
        "# download the dataset archive\n",
        "data_url = \"https://storage.googleapis.com/laurencemoroney-blog.appspot.com/Resources/archive.zip\"\n",
        "data_file_name = \"archive.zip\"\n",
        "download_dir = '/tmp/celeb/'\n",
        "urllib.request.urlretrieve(data_url, data_file_name)\n",
        "\n",
        "# extract the zipped file\n",
        "zip_ref = zipfile.ZipFile(data_file_name, 'r')\n",
        "zip_ref.extractall(download_dir)\n",
        "zip_ref.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VAz0njNjq98x"
      },
      "source": [
        "You will then prepare the dataset. Preprocessing steps include cropping and transforming the pixel values to the range `[-1, 1]`. Training batches are then prepared so it can be fed into the model later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qb_FccO3urG8"
      },
      "source": [
        "def load_celeba(batch_size, resize=64, crop_size=128):\n",
        "  \"\"\"Creates batches of preprocessed images from the JPG files\n",
        "  Args:\n",
        "    batch_size - batch size\n",
        "    resize - size in pixels to resize the images\n",
        "    crop_size - size to crop from the image\n",
        "  \n",
        "  Returns:\n",
        "    prepared dataset\n",
        "  \"\"\"\n",
        "  # initialize zero-filled array equal to the size of the dataset\n",
        "  image_paths = sorted(glob.glob(\"/tmp/celeb/img_align_celeba/img_align_celeba/*.jpg\"))\n",
        "  images = np.zeros((len(image_paths), resize, resize, 3), np.uint8)\n",
        "  print(\"Creating Images\")\n",
        "\n",
        "  # crop and resize the raw images then put into the array\n",
        "  for i, path in tqdm(enumerate(image_paths)):\n",
        "    with Image.open(path) as img:\n",
        "      left = (img.size[0] - crop_size) // 2\n",
        "      top = (img.size[1] - crop_size) // 2\n",
        "      right = left + crop_size\n",
        "      bottom = top + crop_size\n",
        "      img = img.crop((left, top, right, bottom))\n",
        "      img = img.resize((resize, resize), Image.LANCZOS)\n",
        "      images[i] = np.asarray(img, np.uint8)\n",
        "  \n",
        "  # split the images array into two\n",
        "  split_n = images.shape[0] // 2\n",
        "  images1, images2 = images[:split_n], images[split_n:2 * split_n]\n",
        "  del images\n",
        "\n",
        "  # preprocessing function to convert the pixel values into the range [-1,1]\n",
        "  def preprocess(img):\n",
        "      x = tf.cast(img, tf.float32) / 127.5 - 1.0\n",
        "      return x\n",
        "  \n",
        "  # use the preprocessing function on the arrays and create batches\n",
        "  dataset = tf.data.Dataset.from_tensor_slices((images1, images2))\n",
        "  dataset = dataset.map(\n",
        "      lambda x1, x2: (preprocess(x1), preprocess(x2))\n",
        "  ).shuffle(4096).batch(batch_size, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE)\n",
        "  \n",
        "  return dataset\n",
        "\n",
        "# use the function above to load and prepare the dataset\n",
        "batch_size = 8\n",
        "batch_size = batch_size * strategy.num_replicas_in_sync\n",
        "dataset = load_celeba(batch_size)\n",
        "out_dir = \"celeba_out\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QpLfJuXfuVPk"
      },
      "source": [
        "## Build the Model\n",
        "\n",
        "Next, you will build the generator and discriminator. As mentioned in the lecture, the code in this notebook is generalized to make it easy to reconfigure (such as choosing the type of normalization). With that, you will notice a lot of extra code, mostly related to gradient penalty. You can ignore those and we've set the defaults to the reflect the architecture shown in class. \n",
        "\n",
        "You can try the other settings once you've gone through these defaults. Additional modes made available are based on DRAGAN and WGAN-GP and you can read about it [here](https://arxiv.org/abs/1705.07215) and [here](https://arxiv.org/abs/1704.00028v3). These settings are reconfigured using the utilities below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QvsgWnBEYvhO"
      },
      "source": [
        "# Utilities\n",
        "\n",
        "def _get_norm_layer(norm):\n",
        "    if norm == 'NA':\n",
        "        return lambda: lambda x: x\n",
        "    elif norm == 'batch_normalization':\n",
        "        return layers.BatchNormalization\n",
        "    elif norm == 'instance_normalization':\n",
        "        return tfa.layers.InstanceNormalization\n",
        "    elif norm == 'layer_normalization':\n",
        "        return layers.LayerNormalization\n",
        "\n",
        "\n",
        "def get_initializers():\n",
        "    return (tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.02), # conv initializer\n",
        "            tf.keras.initializers.RandomNormal(mean=1.0, stddev=0.02)) # bn gamma initializer\n",
        "\n",
        "\n",
        "def gradient_penalty(f, real, fake, mode):\n",
        "    def _gradient_penalty(f, real, fake=None):\n",
        "        def _interpolate(a, b=None):\n",
        "            if b is None:   # interpolation in DRAGAN\n",
        "                beta = tf.random.uniform(shape=tf.shape(a), minval=0., maxval=1.)\n",
        "                b = a + 0.5 * tf.math.reduce_std(a) * beta\n",
        "            shape = [tf.shape(a)[0]] + [1] * (a.shape.ndims - 1)\n",
        "            alpha = tf.random.uniform(shape=shape, minval=0., maxval=1.)\n",
        "            inter = a + alpha * (b - a)\n",
        "            inter.set_shape(a.shape)\n",
        "            return inter\n",
        "\n",
        "        x = _interpolate(real, fake)\n",
        "        with tf.GradientTape() as t:\n",
        "            t.watch(x)\n",
        "            pred = f(x)\n",
        "        grad = t.gradient(pred, x)\n",
        "        norm = tf.norm(tf.reshape(grad, [tf.shape(grad)[0], -1]), axis=1)\n",
        "        gp = tf.reduce_mean((norm - 1.)**2)\n",
        "\n",
        "        return gp\n",
        "\n",
        "    if mode == 'none':\n",
        "        gp = tf.constant(0, dtype=real.dtype)\n",
        "    elif mode == 'dragan':\n",
        "        gp = _gradient_penalty(f, real)\n",
        "    elif mode == 'wgan-gp':\n",
        "        gp = _gradient_penalty(f, real, fake)\n",
        "\n",
        "    return gp"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gn2D6EuKzasq"
      },
      "source": [
        "### Generator\n",
        "\n",
        "You will first define the generator layers. Again, you will notice some extra code but the default will follow the architecture in class. Like the DCGAN you previously built, the model here primarily uses blocks containing Conv2D, BatchNormalization, and ReLU layers. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7noZmy02ysCS"
      },
      "source": [
        "def create_generator(input_shape=(1, 1, 128),\n",
        "                    output_channels=3,\n",
        "                    dim=64,\n",
        "                    n_upsamplings=4,\n",
        "                    norm='batch_normalization',\n",
        "                    name='generator'):\n",
        "  \n",
        "    Normalization = _get_norm_layer(norm)\n",
        "    conv_initializer, bn_gamma_initializer = get_initializers()\n",
        "\n",
        "    # 0\n",
        "    x = inputs = tf.keras.Input(shape=input_shape)\n",
        "\n",
        "    # 1: 1x1 -> 4x4\n",
        "    dimensions = min(dim * 2 ** (n_upsamplings - 1), dim * 8)\n",
        "    x = layers.Conv2DTranspose(\n",
        "        dimensions, 4, strides=1, padding='valid', use_bias=False,         \n",
        "        # kernel_initializer=conv_initializer\n",
        "    )(x)\n",
        "    x = Normalization(\n",
        "        # gamma_initializer=bn_gamma_initializer\n",
        "        )(x)\n",
        "    x = layers.ReLU()(x)\n",
        "\n",
        "    # 2: upsamplings, 4x4 -> 8x8 -> 16x16 -> ...\n",
        "    for i in range(n_upsamplings - 1):\n",
        "        dimensions = min(dim * 2 ** (n_upsamplings - 2 - i), dim * 8)\n",
        "        x = layers.Conv2DTranspose(\n",
        "            dimensions, 4, strides=2, padding='same', use_bias=False,\n",
        "            # kernel_initializer=conv_initializer\n",
        "            )(x)\n",
        "        x = Normalization(\n",
        "            # gamma_initializer=bn_gamma_initializer\n",
        "            )(x)\n",
        "        x = layers.ReLU()(x)\n",
        "\n",
        "    x = layers.Conv2DTranspose(\n",
        "        output_channels, 4, strides=2, padding='same',\n",
        "        # kernel_initializer=conv_initializer\n",
        "    )(x)\n",
        "\n",
        "    outputs = layers.Activation('tanh')(x)\n",
        "\n",
        "    return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3T07PeND0vny"
      },
      "source": [
        "### Discriminator\n",
        "\n",
        "The discriminator will use strided convolutions to reduce the dimensionality of the features. These will be connected to a [LeakyReLU](https://keras.io/api/layers/activation_layers/leaky_relu/) activation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dMrJ1-JYnOEF"
      },
      "source": [
        "def create_discriminator(input_shape=(64, 64, 3),\n",
        "                        dim=64,\n",
        "                        n_downsamplings=4,\n",
        "                        norm='batch_normalization',\n",
        "                        name='discriminator'):\n",
        "    Normalization = _get_norm_layer(norm)\n",
        "    conv_initializer, bn_gamma_initializer = get_initializers()\n",
        "\n",
        "    # 0\n",
        "    x = inputs = tf.keras.Input(shape=input_shape)\n",
        "\n",
        "    # 1: downsamplings, ... -> 16x16 -> 8x8 -> 4x4\n",
        "    x = layers.Conv2D(dim, 4, strides=2, padding='same',\n",
        "                      # kernel_initializer=conv_initializer\n",
        "                      )(x)\n",
        "    x = layers.LeakyReLU(alpha=0.2)(x)\n",
        "\n",
        "    for i in range(n_downsamplings - 1):\n",
        "        dimensions = min(dim * 2 ** (i + 1), dim * 8)\n",
        "        x = layers.Conv2D(dimensions, 4, strides=2, padding='same', use_bias=False,\n",
        "                          # kernel_initializer=conv_initializer\n",
        "                          )(x)\n",
        "        x = Normalization(\n",
        "            # gamma_initializer=bn_gamma_initializer\n",
        "            )(x)\n",
        "        x = layers.LeakyReLU(alpha=0.2)(x)\n",
        "\n",
        "    # 2: logit\n",
        "    outputs = layers.Conv2D(1, 4, strides=1, padding='valid',\n",
        "                            # kernel_initializer=conv_initializer\n",
        "                            )(x)\n",
        "\n",
        "    return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QT7qi7j21VR3"
      },
      "source": [
        "With the layers for the generator and discriminator defined, you can now create the models and set it up for training."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9javUztztr9t"
      },
      "source": [
        "# Settings\n",
        "resize = 64\n",
        "shape = (resize, resize, 3)\n",
        "z_dim = 128\n",
        "n_G_upsamplings = n_D_downsamplings = 4\n",
        "gradient_penalty_mode = 'none'\n",
        "\n",
        "if gradient_penalty_mode == 'none':\n",
        "  d_norm = 'batch_normalization'\n",
        "elif gradient_penalty_mode in ['dragan', 'wgan-gp']:  \n",
        "  # Avoid using BN with GP\n",
        "  d_norm = 'layer_normalization'\n",
        "gradient_penalty_weight = 10.0\n",
        "\n",
        "\n",
        "# Build the GAN\n",
        "with strategy.scope():\n",
        "    # create the generator model\n",
        "    model_G = create_generator(input_shape=(1, 1, z_dim), output_channels=shape[-1], n_upsamplings=n_G_upsamplings)\n",
        "    \n",
        "    # create the discriminator model\n",
        "    model_D = create_discriminator(input_shape=shape, n_downsamplings=n_D_downsamplings, norm=d_norm)\n",
        "    \n",
        "    # print summaries\n",
        "    model_G.summary()\n",
        "    model_D.summary()\n",
        "\n",
        "    # set optimizers\n",
        "    param_G = tf.keras.optimizers.Adam(lr=0.0002, beta_1=0.5)\n",
        "    param_D = tf.keras.optimizers.Adam(lr=0.0002, beta_1=0.5)\n",
        "\n",
        "    # create distributed dataset\n",
        "    dataset = strategy.experimental_distribute_dataset(dataset)\n",
        "    \n",
        "    # set the loss function\n",
        "    loss_func = tf.keras.losses.BinaryCrossentropy(\n",
        "        from_logits=True, \n",
        "        reduction=tf.keras.losses.Reduction.NONE\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "98EZkGvKwZaP"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RLU54McU2g5d"
      },
      "source": [
        "Finally, you can now train the model. We've provided some helper functions for visualizing and saving the images per epoch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nm7tc4s2uWn_"
      },
      "source": [
        "# Utilities\n",
        "\n",
        "def make_grid(imgs, nrow, padding=0):\n",
        "    assert imgs.ndim == 4 and nrow > 0\n",
        "\n",
        "    batch, height, width, ch = imgs.shape\n",
        "    n = nrow * (batch // nrow + np.sign(batch % nrow))\n",
        "    ncol = n // nrow\n",
        "    pad = np.zeros((n - batch, height, width, ch), imgs.dtype)\n",
        "    x = np.concatenate([imgs, pad], axis=0)\n",
        "\n",
        "    # border padding if required\n",
        "    if padding > 0:\n",
        "        x = np.pad(x, ((0, 0), (0, padding), (0, padding), (0, 0)),\n",
        "                   \"constant\", constant_values=(0, 0))\n",
        "        height += padding\n",
        "        width += padding\n",
        "\n",
        "    x = x.reshape(ncol, nrow, height, width, ch)\n",
        "    x = x.transpose([0, 2, 1, 3, 4])  # (ncol, height, nrow, width, ch)\n",
        "    x = x.reshape(height * ncol, width * nrow, ch)\n",
        "    \n",
        "    if padding > 0:\n",
        "        x = x[:(height * ncol - padding),:(width * nrow - padding),:]\n",
        "    return x\n",
        "\n",
        "def save_img(imgs, filepath, nrow, padding=0):\n",
        "    grid_img = make_grid(imgs, nrow, padding=padding)\n",
        "    grid_img = ((grid_img + 1.0) * 127.5).astype(np.uint8)\n",
        "    with Image.fromarray(grid_img) as img:\n",
        "        img.save(filepath)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8hphQ5P53PQu"
      },
      "source": [
        "This function defines the training on a given batch. It does the two-phase training discussed in class.\n",
        "* First, you train the discriminator to distinguish between fake and real images.\n",
        "* Next, you train the generator to create fake images that will fool the discriminator."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gKNwOru-2_Vb"
      },
      "source": [
        "@distributed(Reduction.SUM, Reduction.SUM, Reduction.CONCAT)\n",
        "def train_on_batch(real_img1, real_img2):\n",
        "    '''trains the GAN on a given batch'''\n",
        "    # concatenate the real image inputs\n",
        "    real_img = tf.concat([real_img1, real_img2], axis=0)\n",
        "\n",
        "    # PHASE ONE - train the discriminator\n",
        "    with tf.GradientTape() as d_tape:\n",
        "\n",
        "        # create noise input\n",
        "        z = tf.random.normal(shape=(real_img.shape[0], 1, 1, z_dim))\n",
        "\n",
        "        # generate fake images\n",
        "        fake_img = model_G(z)\n",
        "\n",
        "        # feed the fake images to the discriminator\n",
        "        fake_out = model_D(fake_img)\n",
        "\n",
        "        # feed the real images to the discriminator\n",
        "        real_out = model_D(real_img)\n",
        "\n",
        "        # use the loss function to measure how well the discriminator\n",
        "        # labels fake or real images\n",
        "        d_fake_loss = loss_func(tf.zeros_like(fake_out), fake_out)\n",
        "        d_real_loss = loss_func(tf.ones_like(real_out), real_out)\n",
        "\n",
        "        # get the total loss\n",
        "        d_loss = (d_fake_loss + d_real_loss) \n",
        "        d_loss = tf.reduce_sum(d_loss) / (batch_size * 2)\n",
        "\n",
        "        # Gradient Penalty (ignore if you set mode to `none`)\n",
        "        gp = gradient_penalty(partial(model_D, training=True), real_img, fake_img, mode=gradient_penalty_mode)\n",
        "        gp = gp  / (batch_size * 2)\n",
        "        d_loss = d_loss + gp * gradient_penalty_weight\n",
        "\n",
        "    # get the gradients\n",
        "    gradients = d_tape.gradient(d_loss, model_D.trainable_variables)\n",
        "    \n",
        "    # update the weights of the discriminator\n",
        "    param_D.apply_gradients(zip(gradients, model_D.trainable_variables))\n",
        "    \n",
        "\n",
        "    # PHASE TWO - train the generator\n",
        "    with tf.GradientTape() as g_tape:\n",
        "        # create noise input\n",
        "        z = tf.random.normal(shape=(real_img.shape[0], 1, 1, z_dim))\n",
        "        \n",
        "        # generate fake images\n",
        "        fake_img = model_G(z)\n",
        "\n",
        "        # feed fake images to the discriminator\n",
        "        fake_out = model_D(fake_img)\n",
        "        \n",
        "        # use loss function to measure how well the generator\n",
        "        # is able to trick the discriminator (i.e. model_D should output 1's)\n",
        "        g_loss = loss_func(tf.ones_like(fake_out), fake_out)\n",
        "        g_loss = tf.reduce_sum(g_loss) / (batch_size * 2)\n",
        "    \n",
        "    # get the gradients\n",
        "    gradients = g_tape.gradient(g_loss, model_G.trainable_variables)\n",
        "\n",
        "    # update the weights of the generator\n",
        "    param_G.apply_gradients(zip(gradients, model_G.trainable_variables))\n",
        "    \n",
        "    # return the losses and fake images for monitoring\n",
        "    return d_loss, g_loss, fake_img "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZTn0jGin6ldz"
      },
      "source": [
        "This will start the training loop. We set the number of epochs but feel free to revise it. From initial runs, it takes around 50 seconds to complete 1 epoch. We've setup a progress bar to display the losses per epoch and there is code as well to print the fake images generated."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "snusiGSBtrlU"
      },
      "source": [
        "# generate a batch of noisy input\n",
        "test_z = tf.random.normal(shape=(64, 1, 1, z_dim))\n",
        "\n",
        "# start loop\n",
        "for epoch in range(30): \n",
        "    with tqdm(dataset) as pbar:\n",
        "        pbar.set_description(f\"[Epoch {epoch}]\")\n",
        "        for step, (X1, X2) in enumerate(pbar):\n",
        "            # train on the current batch\n",
        "            d_loss, g_loss, fake = train_on_batch(X1, X2)\n",
        "\n",
        "            # display the losses\n",
        "            pbar.set_postfix({\"g_loss\": g_loss.numpy(), \"d_loss\": d_loss.numpy()})\n",
        "    \n",
        "        # generate fake images\n",
        "        fake_img = model_G(test_z)\n",
        "\n",
        "    # save output\n",
        "    if not os.path.exists(out_dir):\n",
        "        os.makedirs(out_dir)\n",
        "    file_path = out_dir+f\"/epoch_{epoch:04}.png\"\n",
        "    save_img(fake_img.numpy()[:64], file_path, 8)\n",
        "    \n",
        "    # display gallery of fake faces\n",
        "    if epoch % 1 == 0:\n",
        "        with Image.open(file_path) as img:\n",
        "            plt.imshow(np.asarray(img))\n",
        "            plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NnIhFWClwbmq"
      },
      "source": [
        "## Display GIF sample results\n",
        "\n",
        "You can run the cells below to display the galleries as an animation. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6DiVvDz7d845"
      },
      "source": [
        "imgs = os.listdir('celeba_out')\n",
        "imgs.sort()\n",
        "imgs = [cv2.imread('celeba_out/' + i) for i in imgs]\n",
        "imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]\n",
        "imageio.mimsave('anim.gif', imgs, fps=2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JmllPCyfeZsN"
      },
      "source": [
        "path=\"anim.gif\"\n",
        "\n",
        "with open(path,'rb') as f:\n",
        "    display(IpyImage(data=f.read(), format='png'))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RouHcCpJ8bTv"
      },
      "source": [
        "**Congratulations on completing the final ungraded lab for this course!**"
      ]
    }
  ]
}