src/mnist-gan.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST-GAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate handwritten digits by training a GAN on the MNIST dataset.\n",
    "\n",
    "This project is adapted from Boldizsar Zopcsak's great [GAN-Tutorial-Notebook](https://github.com/BoldizsarZopcsak/GAN-Tutorial-Notebook) (go check it out!). The original project is licensed under the [MIT License](https://github.com/BoldizsarZopcsak/GAN-Tutorial-Notebook/blob/master/LICENSE).<br>\n",
    "This project is licensed under the MIT License as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a quick primer on GANs. For a more thorough write-up check out these links:\n",
    "- [A Beginner's Guide to Generative Adversarial Networks](https://pathmind.com/wiki/generative-adversarial-network-gan)\n",
    "- [A Gentle Introduction to Generative Adversarial Networks](https://machinelearningmastery.com/what-are-generative-adversarial-networks-gans/)\n",
    "- [Generative Adversarial Nets - Original Paper by Ian Goodfellow](https://arxiv.org/pdf/1406.2661.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A GAN (Generative Adversarial Network) is a technique, in which two neural networks compete against each other, to improve their individual performance. GANs are mostly used for AI image generation, such as the generation of new faces.\n",
    "\n",
    "A GAN comprises two parts: a **generator model** and a **discriminator model**<br>\n",
    "The generator creates new samples resembling the training set (i.e. handwritten digits in our case), while the discriminator tries to discern between real samples from the training set and fake samples generated by the generator model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training a GAN is done in two steps.\n",
    "\n",
    "#### Step 1: Training the Discriminator\n",
    "\n",
    "In the first the discriminator is trained. The discriminator is fed some real samples from the training set, labelled as such (usually with a `1`), and some fake samples generated by the discriminator, labelled as fake samples (usually with a `0`). Then, standard backpropagation is applied onto the discriminator only.\n",
    "\n",
    "![Step 1: Training the discriminator](../static/gan/figure1.png)\n",
    "\n",
    "#### Step 2: Training the Generator\n",
    "\n",
    "Next, the generator is trained. First, the generator generates a batch of fake images. These are fed to the discriminator with real labels (usually a `1`). The error is then backpropagated to the generator, but only the generator weights are updated.\n",
    "\n",
    "![Step 2: Training the generator](../static/gan/figure2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 200\n",
    "EPOCHS = 20\n",
    "\n",
    "# GAN configurations\n",
    "LATENT_DIM = 100\n",
    "REAL_LABEL = 1\n",
    "FAKE_LABEL = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading MNIST Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll be using the MNIST dataset to train our GAN. It contains images of handwritten digits. Loading MNIST is trivial using `torchvision`.\n",
    "\n",
    "Before we can use the images to train the network, it's a best practice to normalize the images. The images are black-and-white, represented by values from [0, 1]. The transformation will bring the values in a range of [-1, 1]:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5), (0.5))\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torchvision.datasets.MNIST(\n",
    "    root='./data', download=True, transform=transform\n",
    ")\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=int(BATCH_SIZE/2), shuffle=True, num_workers=2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize our training set before actually using it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataiter = iter(dataloader)\n",
    "images, labels = dataiter.next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_img(img):\n",
    "    img = img / 2 + 0.5 # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(npimg[:, :], cmap='gray_r')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_img(images[0].squeeze())\n",
    "print('Label: %s' % labels[0].item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Latent Space"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generator creates new images by upsampling a random seed, the so-called latent space.\n",
    "\n",
    "We'll create a helper function that creates a minibatch of latent spaces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_latent_points(latent_dim, n_samples):\n",
    "    return torch.rand(n_samples, latent_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show random latent_space / seed\n",
    "img = generate_latent_points(28*28, 1).view(28, -1)\n",
    "show_img(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Model` class will be used as a base for the `Generator` and `Discriminator` class. It mainly offers the `predict` method, to run the model without accumulating gradients, and the `train_on` method, used to train a model on a batch and backpropagate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model, self).__init__()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "    \n",
    "    def predict(self, x):\n",
    "        return self.model(x).detach()\n",
    "        \n",
    "    def train_on(self, x, y, criterion, optimizer):\n",
    "        output = self.model(x)\n",
    "        loss = criterion(output, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we'll define a helper class to reshape tensors. This will be used as a layer in our network and reshape input tensors to a desired size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Reshape(nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super(Reshape, self).__init__()\n",
    "        self.shape = shape\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x.view(*self.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Generator takes a latent space as input and upsamples this seed to a random image of a digit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(Model):\n",
    "    def __init__(self):\n",
    "        super(Generator, self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            # Pass the latent space through a fully-connected layer\n",
    "            nn.Linear(LATENT_DIM, 128*7*7),\n",
    "            nn.BatchNorm1d(128*7*7),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            # Reshape 1D tensor to a 7x7 image\n",
    "            Reshape((-1, 128, 7, 7)),\n",
    "            \n",
    "            # Upsample to 14x14\n",
    "            nn.ConvTranspose2d(128, 128, 4, stride=2),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            # Upsample to 28x28\n",
    "            nn.ConvTranspose2d(128, 128, 4, stride=2),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            nn.Conv2d(128, 1, 7),\n",
    "            nn.Tanh()\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator Test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test the generator just to check that it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator_example = Generator()\n",
    "generator_example.train(False) # Set the generator to evaluation mode\n",
    "latent_points_example = generate_latent_points(LATENT_DIM, 1)\n",
    "\n",
    "generated_image_example = generator_example.predict(latent_points_example)\n",
    "show_img(generated_image_example.squeeze())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The discriminator takes in an image an assesses, whether the image is real or fake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(Model):\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(1, 64, 3, stride=2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Dropout(p=0.4),\n",
    "            \n",
    "            nn.Conv2d(64, 64, 3, stride=2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Dropout(p=0.4),\n",
    "            \n",
    "            nn.Flatten(),\n",
    "            nn.Linear(2304, 1),\n",
    "            nn.Sigmoid()\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator Test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now test discriminator just to check that it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "discriminator_example = Discriminator()\n",
    "\n",
    "prediction_example = discriminator_example(generated_image_example)\n",
    "print('Discriminator Prediction: %f' % prediction_example.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we'll try setting up pytorch to use a CUDA-capable GPU. If no GPU is detected, the GAN will be trained on CPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device \"%s\" for training' % dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll create the networks and move them to our selected device:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = Generator().to(dev)\n",
    "discriminator = Discriminator().to(dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also define the optimizers, used to train the networks, and our loss function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "criterion = nn.BCELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's finally train our GAN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "    for i, data in enumerate(dataloader, 0):\n",
    "        # ===========================\n",
    "        # STEP 1: Train Discriminator\n",
    "        # ===========================\n",
    "        \n",
    "        # MNIST training data\n",
    "        # data[0] are the images, data[1] are the labels (e.g. 1, 2, 3, 4, ...) that we don't need\n",
    "        real_imgs = data[0].to(dev)\n",
    "        real_labels = torch.zeros(int(BATCH_SIZE/2)).fill_(REAL_LABEL)\n",
    "        \n",
    "        # Fake training data\n",
    "        latent_points = generate_latent_points(LATENT_DIM, int(BATCH_SIZE/2)).to(dev)\n",
    "        fake_imgs = generator.predict(latent_points).detach()\n",
    "        fake_labels = torch.zeros(int(BATCH_SIZE/2)).fill_(FAKE_LABEL)\n",
    "\n",
    "        discriminator_optimizer.zero_grad()\n",
    "\n",
    "        # Combine real and fake half-batches to one full batch \n",
    "        images_all = torch.cat((real_imgs, fake_imgs))\n",
    "        labels_all = torch.cat((real_labels, fake_labels)).unsqueeze(1).to(dev)\n",
    "        \n",
    "        discriminator_loss = discriminator.train_on(\n",
    "            images_all, labels_all, criterion, discriminator_optimizer\n",
    "        )\n",
    "        \n",
    "        # =======================\n",
    "        # STEP 2: Train Generator\n",
    "        # =======================\n",
    "        generator_optimizer.zero_grad()\n",
    "\n",
    "        fake_batch = generate_latent_points(LATENT_DIM, BATCH_SIZE).to(dev)\n",
    "        fake_batch_labels = torch.zeros(BATCH_SIZE).fill_(REAL_LABEL).unsqueeze(1).to(dev)\n",
    "\n",
    "        generator_imgs = generator(fake_batch)\n",
    "        generator_loss = criterion(discriminator(generator_imgs), fake_batch_labels)\n",
    "        generator_loss.backward()\n",
    "        generator_optimizer.step()\n",
    "\n",
    "        # ===============================\n",
    "        # STEP 3: Logging and Visualizing\n",
    "        # ===============================\n",
    "        if i % 25 == 0:\n",
    "            print('Epoch: %.2i    Batch Number: %.3i / %.3i    Generator Loss: %.9f    Discriminator Loss: %.9f' %\n",
    "                 (epoch, i, len(dataloader), generator_loss.item(), discriminator_loss.item()))\n",
    "            # show_img(fake_imgs.to(\"cpu\")[0].squeeze().detach())\n",
    "            \n",
    "            fig = plt.figure()\n",
    "            for a in range(30):\n",
    "                ax = fig.add_subplot(6, 6, a + 1)\n",
    "                plt.imshow(fake_imgs.to(\"cpu\").detach()[a, 0, :, :], cmap='gray_r')\n",
    "            \n",
    "            plt.savefig('figs/gan/plot  epoch ' + '%02d' % epoch + '  batch ' + '%05d' % i + '.png', dpi=600)\n",
    "            plt.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}