src/mnist-cgan.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MNIST-cGAN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate handwritten digits by training a cGAN on the MNIST dataset.\n",
    "\n",
    "This project is licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This project assumes a basic understanding of Generative Adversarial Networks (GANs). It builds on a basic MNIST-GAN. <br>\n",
    "For a comprehensive introduction to GANs and an example implementation in pytorch, check out [MNIST-GAN](https://github.com/umcconnell/mnist-gan)\n",
    "\n",
    "Conditional GANS (cGANs) are a specific type of GANs that can produce a user-specified class of images instead of only randomingly reproducing the input training set, as it is the case with plain GANs. For example, a cGAN can produce a specified handwritten digit instead of outputting any handwritten digit.\n",
    "\n",
    "To achieve this, a cGAN receives additional labels as inputs, specifying which class of images to reproduce."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "Like with a plain GAN, training is done in two steps."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Training the Discriminator\n",
    "\n",
    "In the first step, the discriminator is trained. The discriminator is fed some real samples from the training set, labelled as real (usually with a 1), and their corresponding digit label, such as `8`. Additionally, some fake samples generated by the discriminator, labelled as fake samples (usually with a 0), and their corresponding fake digit label, such as `3`, are fed to the discriminator. Then, standard backpropagation is applied onto the discriminator only.\n",
    "\n",
    "![Step 1: Training the discriminator](../static/cgan/figure1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Training the Generator\n",
    "\n",
    "Next, the generator is trained. First, the generator generates a batch of fake images with random noise and fake digit labels. The fake images and their corresponding digit labels 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/cgan/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",
    "HALF_BATCH = int(BATCH_SIZE/2)\n",
    "EPOCHS = 20\n",
    "NUM_CLASSES = 10\n",
    "\n",
    "# GAN configurations\n",
    "LATENT_DIM = 100\n",
    "REAL_LABEL = 1\n",
    "FAKE_LABEL = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading MNIST Dataset\n",
    "\n",
    "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=HALF_BATCH, shuffle=True, num_workers=2\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing\n",
    "\n",
    "Let's visualize our training set before actually using it:"
   ]
  },
  {
   "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": [
    "dataiter = iter(dataloader)\n",
    "imgs, labels = next(dataiter)\n",
    "\n",
    "show_img(imgs[0].squeeze())\n",
    "print(\"Label %i\" % labels[0].item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Latent space\n",
    "\n",
    "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\n",
    "\n",
    "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",
    "        with torch.no_grad():\n",
    "            return self.forward(*x).detach()\n",
    "        \n",
    "    def train_on(self, x, y, criterion, optimizer):\n",
    "        output = self.forward(*x)\n",
    "        loss = criterion(output, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also define a custom embed layer, that turns the MNIST digit labels, such as `6`, into one-hot vectors with a `1` at the digit label index. For the digit label `6` this would be `[0, 0, 0, 0, 0, 0, 1, 0, 0, 0]`. Check out the [pytorch docs](https://pytorch.org/docs/stable/nn.functional.html#one-hot) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embed(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Embed, self).__init__()\n",
    "        \n",
    "    def forward(self, label):\n",
    "        # you can also try:\n",
    "        # return F.embedding(NUM_CLASSES, NUM_CLASSES)\n",
    "        return F.one_hot(label, NUM_CLASSES).float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generator\n",
    "\n",
    "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 and a digit label as input and upsamples this seed with the embedded digit label to the desired image of the digit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(Model):\n",
    "    def __init__(self):\n",
    "        super(Generator, self).__init__()\n",
    "                \n",
    "        self.image = nn.Sequential(\n",
    "            nn.Linear(LATENT_DIM, 128*7*7),\n",
    "            nn.BatchNorm1d(128*7*7),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            Reshape((-1, 128, 7, 7)),\n",
    "        )\n",
    "        \n",
    "        self.label = nn.Sequential(\n",
    "            Embed(),\n",
    "            nn.Linear(10, 7*7),\n",
    "            Reshape((-1, 1, 7, 7))\n",
    "        )\n",
    "        \n",
    "        self.main = nn.Sequential(\n",
    "            # Upsample to 14x14\n",
    "            nn.ConvTranspose2d(129, 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",
    "        )\n",
    "        \n",
    "    def forward(self, img, label):\n",
    "        img = self.image(img)\n",
    "        label = self.label(label)\n",
    "        \n",
    "        x = torch.cat((img, label), dim=1)\n",
    "        \n",
    "        return self.main(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator Test\n",
    "\n",
    "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",
    "digit_label_example = torch.tensor(8)\n",
    "\n",
    "generated_image_example = generator_example.predict(latent_points_example, digit_label_example)\n",
    "show_img(generated_image_example.squeeze())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discriminator\n",
    "\n",
    "\n",
    "The discriminator takes in an image with a corresponding digit label and 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",
    "                \n",
    "        self.label = nn.Sequential(\n",
    "            Embed(),\n",
    "            nn.Linear(10, 28*28),\n",
    "            Reshape((-1, 1, 28, 28))\n",
    "        )\n",
    "        \n",
    "        self.main = nn.Sequential(\n",
    "            nn.Conv2d(2, 128, 3, stride=2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            nn.Conv2d(128, 128, 3, stride=2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            nn.Conv2d(128, 128, 3, stride=2),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            \n",
    "            nn.Flatten(),\n",
    "            nn.Dropout(0.4),\n",
    "                        \n",
    "            nn.Linear(512, 1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "        \n",
    "    def forward(self, img, label):\n",
    "        label = self.label(label)\n",
    "        \n",
    "        x = torch.cat((img, label), dim=1)\n",
    "        \n",
    "        return self.main(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator Test\n",
    "\n",
    "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.predict(generated_image_example, digit_label_example)\n",
    "print('Discriminator Prediction: %f' % prediction_example.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "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 cGAN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(EPOCHS):\n",
    "    for i, (imgs, digit_labels) in enumerate(dataloader, 0):\n",
    "        # ===========================\n",
    "        # STEP 1: Train Discriminator\n",
    "        # ===========================        \n",
    "        imgs, digit_labels = imgs.to(dev), digit_labels.to(dev)\n",
    "        \n",
    "        discriminator_real = torch.zeros(HALF_BATCH).fill_(REAL_LABEL)\n",
    "        discriminator_fake = torch.zeros(HALF_BATCH).fill_(FAKE_LABEL)\n",
    "        \n",
    "        # Fake training data\n",
    "        fake_digit_labels = torch.randint(NUM_CLASSES, (HALF_BATCH,)).to(dev)\n",
    "        \n",
    "        z = generate_latent_points(LATENT_DIM, HALF_BATCH).to(dev)\n",
    "        fake_imgs = generator.predict(z, fake_digit_labels)\n",
    "\n",
    "        discriminator_optimizer.zero_grad()\n",
    "\n",
    "        # Combine real and fake half-batches to one full batch \n",
    "        images_all = torch.cat((imgs, fake_imgs))\n",
    "        labels_all = torch.cat((digit_labels, fake_digit_labels)).to(dev)\n",
    "        discriminator_all = torch.cat((discriminator_real, discriminator_fake)).unsqueeze(1).to(dev)\n",
    "        \n",
    "        discriminator_loss = discriminator.train_on(\n",
    "            (images_all, labels_all), discriminator_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_digit_labels = torch.randint(NUM_CLASSES, (BATCH_SIZE,)).to(dev)\n",
    "        discriminator_fake = torch.zeros(BATCH_SIZE).fill_(REAL_LABEL).unsqueeze(1).to(dev)\n",
    "\n",
    "        generator_imgs = generator.forward(fake_batch, fake_digit_labels)\n",
    "        generator_loss = criterion(discriminator(generator_imgs, fake_digit_labels), discriminator_fake)\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",
    "            \n",
    "            generator.eval()\n",
    "            \n",
    "            z = generate_latent_points(LATENT_DIM, 30).to(dev)\n",
    "            # Generate digit labels 1 through 9 to arrange same digits in columns\n",
    "            y = torch.tensor([x % 10 for x in range(30)]).to(dev)\n",
    "            fake_imgs = generator.predict(z, y).to(\"cpu\").detach()\n",
    "            \n",
    "            fig = plt.figure()\n",
    "            for a in range(30):\n",
    "                ax = fig.add_subplot(3, 10, a + 1)\n",
    "                ax.axes.get_xaxis().set_ticks([])\n",
    "                ax.axes.get_yaxis().set_ticks([])\n",
    "                \n",
    "                plt.imshow(fake_imgs[a, 0, :, :], cmap='gray_r')\n",
    "            \n",
    "            plt.savefig('figs/plot  epoch ' + '%02d' % epoch + '  batch ' + '%05d' % i + '.png', dpi=600)\n",
    "            plt.close()\n",
    "            \n",
    "            generator.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}