Face Synthesis with GANs/Face synthesis with GANs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Face synthesis with GANs in PyTorch (and Keras)\n",
    "**Nishant Prabhu, 30 July 2020**\n",
    "\n",
    "In this tutorial, we will build and train a simple Generative Adversarial Network (GAN) to synthesize faces of people. I'll begin with a brief introduction on GAN's: their architecture and the amazing idea that makes them work. Then, we'll look at some code to get this to work for us. I'll leave you with some ideas which can help you make them produce better results.\n",
    "\n",
    "**Note:** Since most of the tutorial involves PyTorch, all tensors will de represented in the NCHW format, i.e. (batch size, channels, height, width).\n",
    "\n",
    "### Dataset download\n",
    "This project uses a subset (11000 images) of the Flickr Faces dataset. The images can be downloaded from [this](https://drive.google.com/drive/folders/1tg-Ur7d4vk1T8Bn0pPpUSQPxlPGBlGfv) page. Each folder on this page downloads 1000 images. \n",
    "\n",
    "### System specifications\n",
    "I'm using the following version of these Python modules.\n",
    "\n",
    "    1. tensorflow-gpu         :   2.2.0\n",
    "    2. torch                  :   1.5.1\n",
    "    3. mtcnn                  :   0.1.0\n",
    "    4. Python                 :   3.6.9\n",
    "\n",
    "## Generative vs. Discriminative models\n",
    "\n",
    "Consider the example below.\n",
    "\n",
    "<div style=\"display:inline-block;\"><img src=\"https://i.imgur.com/BeJaJUD.jpg\" alt=\"alt text\" width=\"400\"></div>\n",
    "<div style=\"display:inline-block;\"><img src=\"https://i.imgur.com/cyGAQST.jpg\" alt=\"alt text\" width=\"400\"></div>\n",
    "\n",
    "Let's say the features you used to tell this man from that dog are the edges of the image. While you would successfully be able to differentiate between the two, will you (assume you have never seen humans or mirrors before) be able to reconstruct the image of the man or the dog using only those features? Chances are you cannot. This is because the image's edges aren't enough to describe the man or the dog in sufficient detail. Such features are called **discriminative features**.\n",
    "\n",
    "However, a major shift in interest took place in 2014 when Ian Goodfellow et. al. presented the **Generative Adversarial Networks** to the community (it would be unjust to say generative models were unknown at that time, however). This class of models called **generative networks** extract features from data that describe it as an individual. What's more interesting is that the converse holds true: given a description of the features, the models can reconstruct the data to represent those.\n",
    "\n",
    "## Generative Adversarial Networks\n",
    "\n",
    "A GAN consists of two networks playing a minimax game. The architecture of this model is shown in the figure below.\n",
    "\n",
    "<p align=\"center\">\n",
    "    <img src=\"https://i.imgur.com/afXgjbo.png\" alt=\"gan model\" width=\"800\" />\n",
    "</p>\n",
    "\n",
    "The way this model is trained clearly explains why it works. Let's say there is some form of data that we want to synthesize (MNIST digits, for example). We will assume that each image in the dataset can be near-completely described using a vector from a distribution $f(x)$ (called the **latent distribution)**.\n",
    "1. The generator $G$ is given a batch of **latent vectors** from $f(x)$. We expect each of these vectors to correspond to some image (doesn't matter which). $G$, being untrained, uses these vectors to generate some images which are just noise.\n",
    "2. A separate batch of good images is picked out from the dataset and combined with the noisy images the generator created. Now, we create class labels for these images: 0 if the image is real and 1 if the image is fake.\n",
    "3. This batch of real + fake images and their labels is used to train the discriminator $D$. It's sole job is to tell whether an image is fake or not. Since the images made by $G$ are quite bad, $D$ learns its job fairly quickly.\n",
    "4. The magic happends now. Another batch of latent vectors is drawn from $f(x)$ and $G$ generates their corresponding images. This time, we associate inverted labels with these images - we say that all of these images are real (class label 0). This batch is passed into $D$.\n",
    "5. Since $D$ is well trained, it will immediately predict that all of these images are fake. This results in a large classification loss, which we propagate backwards. But before that, **we render $D$ untrainable**. This does two things: \n",
    "    1. The generator receives the loss that its poor images caused. It modifies its weights so that next time, the images it makes look more realistic. \n",
    "    2. The discriminator does not learn the inverted weights. This is important because we need $D$ to bust $G$ as effectively as possible. \n",
    "\n",
    "Every time this cycle is repeated, $G$ becomes better and better at generating more realistic images, while $D$ has to keep track of $G$'s antics so it can still tell the real images from fake ones. Note that it is necessary for both networks to get better at each step: if the discriminator performs poorly, the generator will be satisfied with whatever silly images it is making and won't learn to the level we desire.\n",
    "\n",
    "**Why is this a minimax game?** The discriminator always tries to reduce the loss while the generator tries to raise it (by making the discriminator call its fake images real). The equilibrium reached by this duo determines the quality of output that we obtain from the generator. GAN seeks to optimize the loss function shown below.\n",
    "\n",
    "$$\n",
    "J(x,z) = E_{x}[\\log (1 - D(x))] + E_{z}[\\log(D(G(z)))]\n",
    "$$\n",
    "\n",
    "Here $x$ represents the data (like an image) and $z$ represents the vector drawn from the latent distribution $f$.\n",
    "1. The first term is the loss due to **real images**. The discriminator tries its best to predict that they're all real (0), driving that term to zero. \n",
    "2. The second term is the loss due to **fake images**. The generator tries its best to make the discriminator predict that its image is real, driving that term towards $\\infty$. The discriminator tries its best to tell that it's fake, driving that term to zero.\n",
    "\n",
    "## About this project\n",
    "In this project, we are going to build and train a GAN for generating synthetic faces of people. The specific type of GAN used to generate image data is called **DCGAN** (Deep Convolutional GAN). We are going to use a subset of the Flickr Faces dataset. This dataset consists of faces of random people of various age groups looking in various directions, clicked in varying lighting conditions. Since the original images contain a good amount of background, we will first use a pretrained model ([MTCNN for keras](https://github.com/ipazc/mtcnn)) to crop out the faces from these images. Also we resize the images to $(64 \\times 64)$ and grayscale it, so each image is finally of size $(1, 64, 64)$.  \n",
    "\n",
    "MTCNN can be installed as a Python package using pip. MTCNN runs on a Tensorflow backend.\n",
    "\n",
    "```    \n",
    "pip3 install mtcnn\n",
    "```\n",
    "\n",
    "## Face extraction\n",
    "Let's first prepare the data that we'll use for training the model. As mentioned earlier, we'll use MTCNN for this task, which can be installed as python package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dependencies\n",
    "import os\n",
    "import cv2\n",
    "from mtcnn import MTCNN\n",
    "from tqdm.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "```\n",
    "\n",
    "```python\n",
    "# Crop all images to face\n",
    "paths = os.listdir(\"../images\")\n",
    "save_root = \"../images_small/\"\n",
    "cnn = MTCNN()\n",
    "\n",
    "for name in tqdm(paths):\n",
    "    path = \"../images/\" + name\n",
    "    img = cv2.imread(path)\n",
    "    faces = cnn.detect_faces(img)\n",
    "    \n",
    "    # If no face is detected, ignore\n",
    "    if len(faces) == 0:\n",
    "        continue\n",
    "    \n",
    "    # Get the bounding box coordinates of the face\n",
    "    x1, y1, w, h = faces[0]['box']\n",
    "    x2, y2 = x1+w, y1+h\n",
    "    cropped = img[y1:y2, x1:x2]\n",
    "    cropped = cv2.resize(cropped, (64, 64))\n",
    "    \n",
    "    cv2.imwrite(save_root+name, cropped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will write the generated grayscale images to the directory you choose. Next, we will preprocess it to make a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dependencies\n",
    "import os \n",
    "import torch \n",
    "from tqdm import tqdm\n",
    "from PIL import Image, ImageOps\n",
    "from torchvision import transforms\n",
    "\n",
    "# Transformation for image\n",
    "img_transform = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "# Preprocessing function\n",
    "def prepare_input_data(img_dir, num_images=None):\n",
    "    \"\"\" Aggregates images into a Tensor \"\"\"\n",
    "\n",
    "    filenames = os.listdir(img_dir)\n",
    "    counter = 0\n",
    "    images = []\n",
    "    for name in tqdm(filenames):\n",
    "        path = img_dir + '/' + name\n",
    "        # Read the image\n",
    "        img = Image.open(path)\n",
    "        # Convert to grayscale\n",
    "        img = ImageOps.grayscale(img)\n",
    "        # Convert to tensor\n",
    "        img = img_transform(img)\n",
    "        # Convert pixel values ...\n",
    "        # ... from [0, 1] to [-1, 1]\n",
    "        # We'll see why later\n",
    "        img = 2.0 * img - 1.0\n",
    "        # Store as a numpy array\n",
    "        images.append(img.numpy())\n",
    "        counter += 1\n",
    "\n",
    "        # If you want to limit how many examples you'll use\n",
    "        # Else leave the argument as None\n",
    "        if num_images is not None and counter == num_images:\n",
    "            break\n",
    "\n",
    "    images = torch.FloatTensor(images)\n",
    "    return images\n",
    "\n",
    "# Call the function\n",
    "img_dir = \"../../images_small\"\n",
    "images = prepare_input_data(img_dir=img_dir, num_images=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the input data we need, let's start coding our DCGAN.\n",
    "\n",
    "## Model architecture\n",
    "Our DCGAN will consist of a generator that takes in a batch of latent vectors of 200 dimensions and outputs a batch of images of size (1, 64, 64) each. The discriminator takes in a batch of images of size (batch size, 1, height, width) and outputs a tensor of size (batch size, 2) which denotes the class probabilities for each image in the batch.\n",
    "\n",
    "The generator architecture is shown in the image below. We use **Transpose Convolutional layers** to upscale images. You can read more about them [here](https://towardsdatascience.com/transposed-convolution-demystified-84ca81b4baba). The outputs of the last convolutional layer are provided `tanh` activation. Sigmoidal activations in the output have been observed to provide better results. We use `tanh` since it has a larger active region (where gradient magnitudes are sufficiently large). This is the reason why we transformed our images' pixels to lie between (-1, 1) earlier. \n",
    "\n",
    "<p align=\"center\">\n",
    "    <img src=\"https://i.imgur.com/iAF6eFR.png\" alt=\"generator\" width=\"800\" />\n",
    "</p>\n",
    "\n",
    "The discriminator architecture is shown below. In the final parts, we add a fully connected layer which outputs 9 dimensional tensors. The idea is to make the model observe 9 regions in the image (say from a $(3 \\times 3)$ grid) and generate a \"goodness\" score for each. This is then collated into the class probabilities, output with Log softmax activation.\n",
    "\n",
    "<p align=\"center\">\n",
    "    <img src=\"https://i.imgur.com/OGUHMY7.png\" alt=\"discriminator\" width=\"800\" />\n",
    "</p>\n",
    "\n",
    "The GAN is a sequential model of the two above, with the discriminator following the generator. Let's start building the model now. We will construct the Generator, Discriminator and GAN as torch modules. We will call of them in the DCGAN object, which will have several other helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dependencies\n",
    "import numpy as np\n",
    "import torch\n",
    "import time\n",
    "import random\n",
    "import pickle\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "\n",
    "# Discriminator\n",
    "class Discriminator(torch.nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "\n",
    "        self.conv_1 = torch.nn.Conv2d(1, 4, kernel_size=(3, 3), stride=(2, 2))\n",
    "        self.dropout_1 = torch.nn.Dropout2d(0.2)\n",
    "        self.conv_2 = torch.nn.Conv2d(4, 8, kernel_size=(3, 3), stride=(1, 1))\n",
    "        self.dropout_2 = torch.nn.Dropout2d(0.2)\n",
    "        self.conv_3 = torch.nn.Conv2d(8, 16, kernel_size=(3, 3), stride=(1, 1))\n",
    "        self.dropout_3 = torch.nn.Dropout2d(0.2)\n",
    "        self.flatten = torch.nn.Flatten()\n",
    "        self.fc_reduce = torch.nn.Linear(11664, 100)\n",
    "        self.fc_map = torch.nn.Linear(100, 9)\n",
    "        self.fc_out = torch.nn.Linear(9, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # In shape : (batch_size, 1, 64, 64)\n",
    "\n",
    "        x = self.conv_1(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.dropout_1(x)\n",
    "        x = self.conv_2(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.dropout_2(x)\n",
    "        x = self.conv_3(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.dropout_3(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.fc_reduce(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.fc_map(x)\n",
    "        x = self.fc_out(x)\n",
    "        x = F.log_softmax(x, dim=-1)\n",
    "\n",
    "        return x\n",
    "\n",
    "# Generator\n",
    "class Generator(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, latent_dim):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        self.fc_1 = torch.nn.Linear(latent_dim, 16*4*4)\n",
    "        self.conv_trans_1 = torch.nn.ConvTranspose2d(\n",
    "            16, 32, stride=(2, 2), kernel_size=(4, 4), padding=(1, 1)\n",
    "        )\n",
    "        self.conv_trans_2 = torch.nn.ConvTranspose2d(\n",
    "            32, 64, stride=(2, 2), kernel_size=(4, 4), padding=(1, 1)\n",
    "        )\n",
    "        self.conv_trans_3 = torch.nn.ConvTranspose2d(\n",
    "            64, 128, stride=(2, 2), kernel_size=(4, 4), padding=(1, 1)\n",
    "        )\n",
    "        self.conv_trans_4 = torch.nn.ConvTranspose2d(\n",
    "            128, 128, stride=(2, 2), kernel_size=(4, 4), padding=(1, 1)\n",
    "        )\n",
    "        self.conv = torch.nn.Conv2d(\n",
    "            128, 1, stride=(1, 1), kernel_size=(1, 1)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Input shape : (batch_size, latent_dim)\n",
    "\n",
    "        x = self.fc_1(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = torch.reshape(x, (-1, 16, 4, 4))\n",
    "        x = self.conv_trans_1(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.conv_trans_2(x)\n",
    "        x = F.leaky_relu(x, 0.2)\n",
    "        x = self.conv_trans_3(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv_trans_4(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv(x)\n",
    "        x = torch.tanh(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "# GAN model\n",
    "class GAN(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, generator, discriminator):\n",
    "        super(GAN, self).__init__()\n",
    "\n",
    "        self.generator = generator\n",
    "        self.discriminator = discriminator\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Input shape : (batch_size, latent_dim)\n",
    "\n",
    "        x = self.generator(x)\n",
    "        x = self.discriminator(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The use of activations functions and dropouts the way I have has been empirically shown to produce better results. You can read more about these hacks [here](https://github.com/soumith/ganhacks), although I haven't followed it strictly. I will now show the `DCGAN` class definition in parts, since it has lot of methods and I want to discuss each of them separately.\n",
    "\n",
    "## DCGAN\n",
    "All functions defined this point forward are methods of the class below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DCGAN():\n",
    "\n",
    "    def __init__(self, data, latent_dim, learning_rate, n_critics, device):\n",
    "\n",
    "        self.data = data\n",
    "        self.latent_dim = latent_dim\n",
    "        self.lr = learning_rate\n",
    "        self.update_freq = update_freq\n",
    "        self.device = device\n",
    "\n",
    "        # Initialize generator, discriminator and GAN\n",
    "        self.generator = Generator(latent_dim).to(self.device)\n",
    "        self.discriminator = Discriminator().to(self.device)\n",
    "        self.gan_model = GAN(self.generator, self.discriminator).to(self.device)\n",
    "\n",
    "        # Create optimizers for discriminator and GAN\n",
    "        self.gan_optim = optim.RMSprop(self.gan_model.parameters(), lr=self.lr)\n",
    "        self.disc_optim = optim.RMSprop(self.discriminator.parameters(), lr=self.lr)\n",
    "\n",
    "        # Pretrain discriminator\n",
    "        self.pretrain_discriminator(data_size=1000, epochs=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the initialization function, we define a few variables:\n",
    "1. `data`, which is the tensor of images which we will use for training\n",
    "2. `latent_dim`, which is the dimensionality of the latent distribution from which we draw latent vectors\n",
    "3. `lr`, learning rates for the discriminator's and GAN's optimizers\n",
    "4. `update_freq`: To ensure that the discriminator sees through the generators fake images, it is trained more often than the generator. This variable determines how much more often.\n",
    "5. `device`, a vairable specific to PyTorch which determines whether computations happens on your CPU or another physical device like a GPU.\n",
    "\n",
    "Next, we initialize the generator, discriminator and GAN with appropriate arguments. We also initialize their optimizers, which I have chosen to be RMSprop. We'll talk about the method `pretrain_discriminator` a little later.\n",
    "\n",
    "First, we need a method to generate vectors from the latent space of given batch size. The latent disctribution could really be anything. Here, I'm using a normal distribution with 0 mean and unit variance. Feel free to try others like the uniform or poisson distributions. We will use NumPy's random number generator to do our bidding.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_latent_samples(self, size):\n",
    "    \n",
    "    mat = np.random.normal(loc=0, scale=1, size=(size, self.latent_dim))\n",
    "    return torch.FloatTensor(mat).to(self.device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's focus on training the discriminator now. We'll need batches of data which contain true and fake samples of images, plus their correct labels. The true images can be randomly sampled from the dataset. For the fake images, we'll generate latent samples using the function we just defined and then pass them through the generator. So, the fake images are bad initially but they get better along the way. \n",
    "\n",
    "We also implement another GAN hack here. Since the discriminator is trained much more often than the generator, it is possible that it may get very good at its job. The generator might now start looking for that one image which could cheat the discriminator most, and keep generating only that: resulting in loss of variance (and possibly quality) in generated images. This is known as **mode collapse**, but this is not the only process that results in it. To solve this, we randomly train the discriminator in a wrong way (with inverted labels) to ensure both the models are upskilling equally. This is done a small fraction of the time, which I have chosen to be 10% below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_disc_batch(self, size, pretraining=False):\n",
    "\n",
    "    # Extract size//2 random real samples from data\n",
    "    idx = np.random.choice(\n",
    "        np.arange(self.data.shape[0]), size=size//2, replace=True\n",
    "    )\n",
    "    true_data = torch.FloatTensor(self.data[idx]).to(self.device)\n",
    "\n",
    "    # Generate size//2 fake samples using generator\n",
    "    latent_ = self.generate_latent_samples(size=size//2)\n",
    "    fake_data = self.generator(latent_)\n",
    "\n",
    "    # Concatenate them on the batch size axis\n",
    "    data = torch.cat((true_data, fake_data), dim=0)\n",
    "    \n",
    "    # Labels corresponding to the images\n",
    "    # Key -> 0: real, 1: fake\n",
    "    # Flip labels randomly to train discriminator better\n",
    "    if random.random() < 0.9 or pretraining:\n",
    "        true_labels = torch.LongTensor([0]*(size//2)).to(self.device)\n",
    "        fake_labels = torch.LongTensor([1]*(size//2)).to(self.device)\n",
    "    else:\n",
    "        true_labels = torch.LongTensor([1]*(size//2)).to(self.device)\n",
    "        fake_labels = torch.LongTensor([0]*(size//2)).to(self.device)\n",
    "        \n",
    "    # Concatenate the labels on batch size axis\n",
    "    labels = torch.cat((true_labels, fake_labels), dim=0)\n",
    "\n",
    "    return data, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we'll write a similar function to generate data to train the GAN (but only the generator). We generate a batch of latent vectors from the latent space and **invert the corresponding labels**. That is, we lie to the GAN that all the images it gets out of these are real. This causes the generator to learn, as we discussed earlier. In a separate function, we will render the discriminator untrainable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_gan_batch(self, size):\n",
    "\n",
    "    # Generate size latent samples and generate inverted labels\n",
    "    data = self.generate_latent_samples(size)\n",
    "    labels = torch.LongTensor([0]*size).to(self.device)\n",
    "\n",
    "    return data, labels\n",
    "\n",
    "\n",
    "def discriminator_trainable(self, val):\n",
    "\n",
    "    for param in self.discriminator.parameters():\n",
    "        param.requires_grad = val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Passing `False` to the latter function makes the discriminator untrainable and vice versa. We'll now need a helper function to train these models given a batch of images/latent vectors. In Keras, you would have used `model.train_on_batch(x, y)` but here it'll have to be more elaborate. Also, we add **gradient clipping** here, which is another GAN hack to ensure stability. Basically, we just force the gradients to always lie between (-0.01, 0.01) so the parameter deltas don't become very large."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model_on_batch(self, model, optimizer, x, y):\n",
    "    \n",
    "    # Zero out optimizer gradients\n",
    "    optimizer.zero_grad()\n",
    "    # Generate predictions\n",
    "    probs = model(x)\n",
    "    # Compute nonlinear logloss\n",
    "    loss = F.nll_loss(probs, y, reduction='mean')\n",
    "    # Compute gradients by backpropagation\n",
    "    loss.backward()\n",
    "    # Clip gradients to lie between [-0.01, 0.01]\n",
    "    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.01)\n",
    "    # Update the optimizer and model weights\n",
    "    optimizer.step()\n",
    "\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now talk about the function we saw earlier, `pretrain_discriminator`. The idea of having it is simple: to make the generator start learning right away, we want the discriminator already trained on real and fake data. With this function, we do it in the initialization phase itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretrain_discriminator(self, data_size, epochs):\n",
    "\n",
    "    print(\"\\n[INFO] Pretraining discriminator...\\n\")\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        total_loss = 0\n",
    "\n",
    "        for _ in range(self.update_freq):\n",
    "            \n",
    "            # Generate a batch\n",
    "            x, y = self.generate_disc_batch(data_size, True)\n",
    "            # Train discriminator on this batch\n",
    "            loss = self.train_model_on_batch(\n",
    "                self.discriminator, self.disc_optim, x, y\n",
    "            )\n",
    "            # Record the loss\n",
    "            total_loss += loss\n",
    "\n",
    "        # Output status to console\n",
    "        print(\"Epoch {} - Loss {:.4f}\".format(\n",
    "            epoch+1, total_loss/self.update_freq\n",
    "        ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have all the helper functions (methods) we need. Let's now write the training function calling the above in correct sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(self, epochs, steps_per_epoch, batch_size, save_freq, save_path):\n",
    "\n",
    "    # Lists to store discriminator and GAN losses\n",
    "    disc_losses = [] \n",
    "    gan_losses = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        print(\"\\n\\nEpoch {}/{}\".format(epoch+1, epochs))\n",
    "        print(\"--------------------------------------\")\n",
    "\n",
    "        for step in range(steps_per_epoch):\n",
    "\n",
    "            disc_loss_hist = []\n",
    "            \n",
    "            # Train discriminator update_freq no. of times\n",
    "            for _ in range(self.update_freq):\n",
    "\n",
    "                # Generate batch for discriminator\n",
    "                x, y = self.generate_disc_batch(batch_size)\n",
    "\n",
    "                # Train discriminator on fake batch\n",
    "                disc_loss = self.train_model_on_batch(\n",
    "                    self.discriminator, self.disc_optim, x, y\n",
    "                )\n",
    "                disc_loss_hist.append(disc_loss)\n",
    "\n",
    "            # Render discriminator untrainable\n",
    "            self.discriminator_trainable(False)\n",
    "\n",
    "            # Generate batch for GAN\n",
    "            x_gan, y_gan = self.generate_gan_batch(batch_size)\n",
    "\n",
    "            # Train GAN on this batch\n",
    "            gan_loss = self.train_model_on_batch(\n",
    "                self.gan_model, self.gan_optim, x_gan, y_gan\n",
    "            )\n",
    "\n",
    "            # Render discriminator trainable\n",
    "            self.discriminator_trainable(True)\n",
    "            \n",
    "            # Output status every 100 steps\n",
    "            if step % 100 == 0:\n",
    "                print(\"Step {:3d} - Disc loss {:.4f} - GAN loss {:.4f}\".format(\n",
    "                    step, sum(disc_loss_hist)/self.update_freq, gan_loss\n",
    "                ))\n",
    "\n",
    "            # Update loss logs\n",
    "            disc_losses.append(sum(disc_loss_hist)/self.update_freq)\n",
    "            gan_losses.append(gan_loss)\n",
    "\n",
    "        # Save model every svae_freq epochs\n",
    "        if epoch % save_freq == 0:\n",
    "            torch.save(self.generator.state_dict(), save_path +\n",
    "                       \"/generator_{}\".format(epoch+1))\n",
    "            torch.save(self.gan_model.state_dict(),\n",
    "                       save_path + \"/gan_{}\".format(epoch+1))\n",
    "\n",
    "        # Test generator\n",
    "        self.test_generator(num_samples=50)\n",
    "\n",
    "        print(\"====================================\")\n",
    "\n",
    "    # Save critic and GAN losses\n",
    "    with open(save_path + \"/disc_loss.pkl\", \"wb\") as f:\n",
    "        pickle.dump(disc_losses, f)\n",
    "\n",
    "    with open(save_path + \"/gan_loss.pkl\", \"wb\") as f:\n",
    "        pickle.dump(gan_losses, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're not done with it yet. If you noticed, we have used a function called `test_generator` in there. It will be used to check the model's progress every epoch. We want to see the GAN's accuracy at the end of that epoch and display a few examples of the images it has generated. If the model is doing fine, we expect the GAN's accuracy to stay good throughout training and the generators images should slowly get better. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_generator(self, num_samples):\n",
    "\n",
    "    # Generate latent samples and corresponding correct labels\n",
    "    latent_ = self.generate_latent_samples(num_samples)\n",
    "    labels = torch.LongTensor([1]*num_samples).to(self.device)\n",
    "    \n",
    "    # Get predictions from model and calculate accuracy\n",
    "    probs = self.gan_model(latent_)\n",
    "    preds = probs.argmax(dim=-1)\n",
    "    correct = preds.eq(labels).sum().item()\n",
    "    print(\"\\n[TESTING] GAN accuracy: {:.4f}\".format(correct/num_samples))\n",
    "\n",
    "    # Show some images made by the generator\n",
    "    gen_out = self.generator(latent_)\n",
    "    images = gen_out.cpu().detach().numpy()[:5]\n",
    "    \n",
    "    # Scale output back to [0, 1] from [-1, 1]\n",
    "    # So the images can be displayed as grayscale\n",
    "    images = (images + 1.) / 2.0\n",
    "\n",
    "    # Plot the 5 images\n",
    "    fig = plt.figure(figsize=(5, 2))\n",
    "    for i in range(5):\n",
    "        fig.add_subplot(1, 5, i+1)\n",
    "        plt.axis(\"off\")\n",
    "        plt.imshow(images[i].squeeze(0), cmap='gray')\n",
    "    plt.tight_layout()\n",
    "    plt.show(block=False)\n",
    "    plt.pause(3)\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I suggest reproducing the DCGAN part of the code in a separate python script and running it through terminal. We can write another script `main.py` to do so. You may include the preprocessing functions here, which we wrote earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main script to run DCGAN\n",
    "\n",
    "import os \n",
    "import torch \n",
    "from tqdm import tqdm\n",
    "from PIL import Image, ImageOps\n",
    "from torchvision import transforms\n",
    "from dcgan import DCGAN\n",
    "\n",
    "\n",
    "# Transformation for image\n",
    "# Already the right size so convert to Tensor and normalize\n",
    "img_transform = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "\n",
    "def prepare_input_data(img_dir, num_images=None):\n",
    "    \"\"\" Aggregates images into a Tensor \"\"\"\n",
    "\n",
    "    filenames = os.listdir(img_dir)\n",
    "    counter = 0\n",
    "    images = []\n",
    "    for name in tqdm(filenames):\n",
    "        path = img_dir + '/' + name\n",
    "        img = Image.open(path)\n",
    "        img = ImageOps.grayscale(img)\n",
    "        img = img_transform(img)\n",
    "        img = 2.0 * img - 1.0\n",
    "        images.append(img.numpy())\n",
    "        counter += 1\n",
    "\n",
    "        if num_images is not None and counter == num_images:\n",
    "            break\n",
    "\n",
    "    images = torch.FloatTensor(images)\n",
    "    return images\n",
    "\n",
    "\n",
    "# Main script\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "\n",
    "    img_dir = \"../../images_small\"\n",
    "\n",
    "    images = prepare_input_data(img_dir=img_dir, num_images=None)\n",
    "\n",
    "    # Training params\n",
    "    epochs = 100\n",
    "    steps_per_epoch = 500\n",
    "    batch_size = 128\n",
    "    save_freq = 5\n",
    "    save_path = \"../../saved_data\"\n",
    "    latent_dim = 200\n",
    "    learning_rate = 5e-05\n",
    "    n_critics = 5 \n",
    "\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    print(\"\")\n",
    "\n",
    "    # Initialize GAN \n",
    "    gan = DCGAN(images, latent_dim, learning_rate, n_critics, device)\n",
    "\n",
    "    # Train GAN\n",
    "    gan.train(epochs, steps_per_epoch, batch_size, save_freq, save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feel free to change the parameters above to something your system can handle easily. You might have noticed that I have used a small learning rate. It is good to do so when training is unstable. This model was trained on an NVIDIA RTX 2060S in about 3 hours. \n",
    "\n",
    "## Results\n",
    "Now, we will load the trained model weights into a new generator object and see what images it generates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new generator object and load the saved weights\n",
    "\n",
    "model = Generator(latent_dim=200)\n",
    "model.load_state_dict(torch.load(\"../../saved_data/generator_96\", map_location=torch.device(\"cpu\")))\n",
    "\n",
    "# Generate latent samples (64)\n",
    "\n",
    "latent_ = np.random.normal(loc=0, scale=1, size=(64, 200))\n",
    "latent_ = torch.FloatTensor(latent_)\n",
    "\n",
    "# Generate corresponding images using generator\n",
    "\n",
    "images = model(latent_)\n",
    "images = images.cpu().detach().numpy()\n",
    "images = (images + 1.0) / 2.0\n",
    "\n",
    "# Show the images in a grid\n",
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "for i in range(images.shape[0]):\n",
    "    fig.add_subplot(8, 8, i+1)\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(images[i].squeeze(0), cmap='gray')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what it produced. Since our model is very simple, many images look slightly to very disfigured. Regardless, I think it has done a decent job.\n",
    "\n",
    "<p align=\"center\">\n",
    "    <img src=\"https://i.imgur.com/31J4KK7.png\" alt=\"result\" width=\"600\" />\n",
    "</p>\n",
    "\n",
    "## Exploring the latent space\n",
    "The generator has learn a function to map the latent space to the images that we see above. What information does the latent space give it? Let's find out. From the images above, I've picked out the latent vectors corresponding to three, which will help me demonstrate my point.\n",
    "1. Image of man at (4, 6), smiling to some extent\n",
    "2. Image of a person at (1, 4), who also seems to be smiling\n",
    "3. Image of another person at (5, 3), who looks slightly worried\n",
    "\n",
    "We now perform this vectorial operation, to get another vector in the same space.\n",
    "\n",
    "```\n",
    "final = person 1 - person 2 + person 3\n",
    "```\n",
    "\n",
    "Then, I'll generate the image corresponding to this latent vector. Let's see what we get."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "person_1 = latent_[29]\n",
    "person_2 = latent_[3]\n",
    "person_3 = latent_[34]\n",
    "\n",
    "final = person_1 - person_2 + person_3\n",
    "\n",
    "# Generate the corresponding image\n",
    "image = model(final.unsqueeze(0))\n",
    "image = image.cpu().detach().numpy()\n",
    "\n",
    "print(\"RESULT:\")\n",
    "plt.figure(figsize=(2, 2))\n",
    "plt.imshow(image[0].squeeze(0), cmap='gray')\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make it easy to visualize, I'll show all the three images and the result side by side (in the order person 1, person 2, person 3, result).\n",
    "\n",
    "<br>\n",
    "<div class=\"row\">\n",
    "    <div class=\"column\" style=\"float: left; width: 20%; padding: 0px;\">\n",
    "        <img src=\"https://i.imgur.com/3XxPfGK.png\" alt=\"person 1\" width=\"200\" />\n",
    "    </div>\n",
    "    <div class=\"column\" style=\"float: left; width: 20%; padding: 0px;\">\n",
    "        <img src=\"https://i.imgur.com/BI52VFX.png\" alt=\"person 2\" width=\"200\" />\n",
    "    </div>\n",
    "    <div class=\"column\" style=\"float: left; width: 20%; padding: 0px;\">\n",
    "        <img src=\"https://i.imgur.com/9YtVrWA.png\" alt=\"person 3\" width=\"200\" />\n",
    "    </div>\n",
    "    <div class=\"column\" style=\"float: left; width: 40%; padding: 0px;\">\n",
    "        <img src=\"https://i.imgur.com/UGVsIll.png\" alt=\"result\" width=\"200\" />\n",
    "    </div>\n",
    "</div>\n",
    "<br>\n",
    "\n",
    "**Did you see that?** Here's what has happened:\n",
    "1. The person in the new image isn't smiling anymore, since we subtracted person 2 from person 1. In fact, he has inherited the worried expression of person 3, which we added.\n",
    "2. This new person has some hair hanging on the top left, similar to how person 3 has their hair.\n",
    "3. The direction of his face has changes to match with person 3.\n",
    "\n",
    "This tells us that each dimension in the latent space carries information about some feature of the face: smile, gender, hair, skin tone, etc. When this is modified in the latent space, the generated image reflects those modifications too. Cool!\n",
    "\n",
    "## Endnote\n",
    "That was a lot! But satisfying, I hope. There's definitely much more and better things this model can do. Here's what you can explore to get this to work better.\n",
    "1. **Wasserstein GANs**. There use a completely different loss function to train the networks. It makes the training process much more stable, and both networks never stop getting updated (the loss function ensures the gradients are never zero).\n",
    "2. **Progressive DCGAN for super resolution**. [This](https://arxiv.org/pdf/1710.10196.pdf) publication by NVIDIA describes a model architecture for generation faces with very high resolution. In this model, new layers are (smoothly) introduced into the networks while they are training. So the model first learns to produce coarse images, then fine tunes itself to improve its resolution. \n",
    "\n",
    "A simple Google search for face synthesis using GANs reveals a plethora of work done in this field. Anyway, I hope this tutorial has been useful to you!"
   ]
  }
 ],
 "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}