notebooks/cis6930_week4_generative_adversarial_networks_student.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "cis6930-week4-generative-adversarial-networks-student.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oRgpqL0KVO43"
      },
      "source": [
        "# CIS6930 Week 4: Generative Adversarial Networks (Student version)\n",
        "\n",
        "---\n",
        "\n",
        "Preparation: Go to `Runtime > Change runtime type` and choose `GPU` for the hardware accelerator.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uayE3ioc3k_h"
      },
      "source": [
        "gpu_info = !nvidia-smi -L\n",
        "gpu_info = \"\\n\".join(gpu_info)\n",
        "if gpu_info.find(\"failed\") >= 0:\n",
        "    print(\"Not connected to a GPU\")\n",
        "else:\n",
        "    print(gpu_info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lTahfFKhDOjM"
      },
      "source": [
        "## Preparation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P9yzHwL53FDW"
      },
      "source": [
        "# Please create a Kaggle account and accept the term before downloading the dataset\n",
        "# https://www.kaggle.com/c/dogs-vs-cats\n",
        "# This is for educational purpose and do not distribute the link\n",
        "!gdown --id 1c_GFLUlW7nB1mwlhNT9t4a2jfilIKPNt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OdUhLYex4tTD"
      },
      "source": [
        "!unzip dogs-vs-cats.zip\n",
        "!unzip train.zip\n",
        "!unzip test1.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z7bJK-do5e9Q"
      },
      "source": [
        "!ls train"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "giQ6PmGV5EUh"
      },
      "source": [
        "from PIL import Image\n",
        "\n",
        "Image.open(\"train/dog.1.jpg\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q7v0eAMo7FOR"
      },
      "source": [
        "Image.open(\"train/cat.1.jpg\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a7h-6LpjVDHW"
      },
      "source": [
        "import copy\n",
        "import random\n",
        "from time import time\n",
        "from typing import Any, Dict\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch import optim\n",
        "from torch.utils.data import Dataset, TensorDataset, DataLoader\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "from matplotlib import pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t6bLN_iACve3"
      },
      "source": [
        "### Organizing data directory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nvXTxH0v7FTb"
      },
      "source": [
        "#!mkdir -p images/cat\n",
        "#!mv train/cat.*.jpg images/cat/\n",
        "!mkdir -p images/dog\n",
        "!mv train/dog.*.jpg images/dog/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LSuTVv-TCqRX"
      },
      "source": [
        "### Load dataset after normalization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "58t-iwNV86X4"
      },
      "source": [
        "image_size = 64 # 128, 224\n",
        "dataset = torchvision.datasets.ImageFolder(\n",
        "    root=\"images\", # it considers subdirectory as the class name\n",
        "    transform=transforms.Compose(\n",
        "        [transforms.Resize((image_size, image_size)),\n",
        "         transforms.CenterCrop(image_size),\n",
        "         transforms.ToTensor(),\n",
        "         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MygzMOdHAOV9"
      },
      "source": [
        "# Utility function\n",
        "def denormalize_image(x):\n",
        "    \"\"\"https://discuss.pytorch.org/t/what-is-the-most-simple-way-to-reobtain-a-pil-image-after-normalization/85355/3\n",
        "    \"\"\"\n",
        "    min_i = x.min(dim=(1), keepdim=True).values.min(dim=(2), keepdim=True).values\n",
        "    max_i = x.max(dim=(1), keepdim=True).values.max(dim=(2), keepdim=True).values\n",
        "    x = ((x - min_i) / (max_i - min_i)) * 255\n",
        "    return transforms.ToPILImage()(x.type(torch.uint8))\n",
        "\n",
        "plt.imshow(transforms.ToPILImage()(dataset[0][0]))\n",
        "plt.show()\n",
        "plt.imshow(denormalize_image(dataset[0][0]))\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUweHOA7DJk7"
      },
      "source": [
        "Alternatively, `torchvision.utils.make_grid()` is a go-to option."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zdY5wzPKC1Cb"
      },
      "source": [
        "dataloader = DataLoader(dataset, batch_size=16,\n",
        "                        shuffle=True)\n",
        "batch = next(iter(dataloader))\n",
        "plt.imshow(np.transpose(torchvision.utils.make_grid(batch[0][:64],\n",
        "                                         padding=2,\n",
        "                                         normalize=True).cpu(),\n",
        "                        (1,2,0)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fE8_6f8SEDZk"
      },
      "source": [
        "## DCGAN Model\n",
        "\n",
        "The following code is based on [this PyTorch tutorial](https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html) with some modification (adding more comments.)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mPlp0KGZEIIF"
      },
      "source": [
        "### Generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "prqZ7XU6D-oa"
      },
      "source": [
        "class Generator(nn.Module):\n",
        "    def __init__(self,\n",
        "                 n_z: int = 100,\n",
        "                 n_f: int = 64, # image_size\n",
        "                 n_c: int = 3,\n",
        "                 relu_slope: float = 0.2 # LeakyReLU slope\n",
        "                 ):\n",
        "        super().__init__()\n",
        "        self.network = nn.Sequential(\n",
        "            # (B, n_z, 1, 1) -> (B, n_f * 8, 4, 4)\n",
        "            nn.ConvTranspose2d(in_channels=n_z,\n",
        "                               out_channels=n_f * 8,\n",
        "                               kernel_size=4,\n",
        "                               stride=1,\n",
        "                               padding=0,\n",
        "                               bias=False),\n",
        "            nn.BatchNorm2d(n_f * 8),\n",
        "            #nn.LeakyReLU(negative_slope=relu_slope, inplace=True),\n",
        "            nn.ReLU(inplace=True),\n",
        "            # (B, n_f * 8, 4, 4) -> (B, n_f * 4, 8, 8)\n",
        "            nn.ConvTranspose2d(in_channels=n_f * 8,\n",
        "                               out_channels=n_f * 4,\n",
        "                               kernel_size=4,\n",
        "                               stride=2,\n",
        "                               padding=1,\n",
        "                               bias=False),\n",
        "            nn.BatchNorm2d(n_f * 4),\n",
        "            #nn.LeakyReLU(negative_slope=relu_slope, inplace=True),\n",
        "            nn.ReLU(True),\n",
        "            # (B, n_f * 4, 8, 8) -> (B, n_f * 2, 16, 16)\n",
        "            nn.ConvTranspose2d(in_channels=n_f * 4,\n",
        "                               out_channels=n_f * 2,\n",
        "                               kernel_size=4,\n",
        "                               stride=2,\n",
        "                               padding=1,\n",
        "                               bias=False),\n",
        "            nn.BatchNorm2d(n_f * 2),\n",
        "            # nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            nn.ReLU(True),\n",
        "            # (B, n_f * 2, 16, 16) -> (B, n_f, 32, 32)\n",
        "            nn.ConvTranspose2d(in_channels=n_f * 2,\n",
        "                               out_channels=n_f,\n",
        "                               kernel_size=4,\n",
        "                               stride=2,\n",
        "                               padding=1,\n",
        "                               bias=False),\n",
        "            nn.BatchNorm2d(n_f),\n",
        "            nn.ReLU(True),\n",
        "            # nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            # (B, n_f, 32, 32) -> (B, n_c, 64, 64)\n",
        "            nn.ConvTranspose2d(in_channels=n_f,\n",
        "                               out_channels=n_c,\n",
        "                               kernel_size=4,\n",
        "                               stride=2,\n",
        "                               padding=1,\n",
        "                               bias=False),\n",
        "            nn.Tanh()\n",
        "            # Output shape: (B, n_c, 64, 64)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.network(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "baGKuCgp3vSE"
      },
      "source": [
        "### Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "avd89JFKG6lk"
      },
      "source": [
        "class Discriminator(nn.Module):\n",
        "    def __init__(self,\n",
        "                 n_f: int = 64, # image_size\n",
        "                 n_c: int = 3,\n",
        "                 relu_slope: float = 0.2):\n",
        "        super().__init__()\n",
        "        self.network = nn.Sequential(\n",
        "            # (B, n_c, 64, 64) -> (B, n_f, 32, 32)\n",
        "            nn.Conv2d(in_channels=n_c,\n",
        "                      out_channels=n_f,\n",
        "                      kernel_size=4,\n",
        "                      stride=2,\n",
        "                      padding=1, bias=False),\n",
        "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            # (B, n_f, 32, 32) -> (B, n_f * 2, 16, 16)\n",
        "            nn.Conv2d(in_channels=n_f,\n",
        "                      out_channels=n_f * 2,\n",
        "                      kernel_size=4,\n",
        "                      stride=2,\n",
        "                      padding=1,\n",
        "                      bias=False),\n",
        "            nn.BatchNorm2d(n_f * 2),\n",
        "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            # (B, n_f * 2, 16, 16) -> (B, n_f * 4, 8, 8)\n",
        "            nn.Conv2d(in_channels=n_f * 2,\n",
        "                      out_channels=n_f * 4,\n",
        "                      kernel_size=4,\n",
        "                      stride=2,\n",
        "                      padding=1, bias=False),\n",
        "            nn.BatchNorm2d(n_f * 4),\n",
        "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            # (B, n_f * 4, 16, 16) -> (B, n_f * 8, 4, 4)\n",
        "            nn.Conv2d(in_channels=n_f * 4,\n",
        "                      out_channels=n_f * 8,\n",
        "                      kernel_size=4,\n",
        "                      stride=2,\n",
        "                      padding=1,\n",
        "                      bias=False),\n",
        "            nn.BatchNorm2d(n_f * 8),\n",
        "            nn.LeakyReLU(negative_slope=0.2, inplace=True),\n",
        "            # (B, n_f * 8, 4, 4) -> (B, 1, 1, 1)\n",
        "            nn.Conv2d(in_channels=n_f * 8,\n",
        "                      out_channels=1,\n",
        "                      kernel_size=4,\n",
        "                      stride=1,\n",
        "                      padding=0,\n",
        "                      bias=False),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.network(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sjgt4i-63x2d"
      },
      "source": [
        "### Custom weight initialization function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k2AmXSduD8CJ"
      },
      "source": [
        "# custom weights initialization for Generator and Discriminator\n",
        "def weights_init(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find(\"Conv\") != -1:\n",
        "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
        "    elif classname.find(\"BatchNorm\") != -1:\n",
        "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
        "        nn.init.constant_(m.bias.data, 0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rWViSNJO3quA"
      },
      "source": [
        "### Training script"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WgiMFP1bKWXV"
      },
      "source": [
        "## ===============\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "n_z = 100\n",
        "\n",
        "lr = 0.0002\n",
        "beta1 = 0.5  # <- 0.9\n",
        "beta2 = 0.999  # 0.999\n",
        "\n",
        "num_epochs = 20\n",
        "batch_size = 16 # 128 is the original configuration\n",
        "\n",
        "random_seed = 1\n",
        "## ===============\n",
        "\n",
        "# Random Seeds ===============\n",
        "torch.manual_seed(random_seed)\n",
        "random.seed(random_seed)\n",
        "np.random.seed(random_seed)\n",
        "# Random Seeds ===============\n",
        "\n",
        "\n",
        "# Create models\n",
        "model_D = Discriminator(n_f=image_size).to(device)\n",
        "model_D.apply(weights_init)\n",
        "\n",
        "model_G = Generator(n_f=image_size).to(device)\n",
        "model_G.apply(weights_init)\n",
        "\n",
        "# Initialize BCELoss function\n",
        "criterion = nn.BCELoss()\n",
        "\n",
        "# Create batch of latent vectors that we will use to visualize\n",
        "#  the progression of the generator\n",
        "fixed_noise = torch.randn(image_size,\n",
        "                          n_z,\n",
        "                          1,\n",
        "                          1, device=device)\n",
        "\n",
        "# Establish convention for real and fake labels during training\n",
        "real_label = 1.\n",
        "fake_label = 0.\n",
        "\n",
        "# Setup Adam optimizers for both G and D\n",
        "optimizer_D = optim.Adam(model_D.parameters(), lr=lr, betas=(beta1, beta2))\n",
        "optimizer_G = optim.Adam(model_G.parameters(), lr=lr, betas=(beta1, beta2))\n",
        "\n",
        "# Data loader\n",
        "dataloader = DataLoader(dataset,\n",
        "                        batch_size=batch_size,\n",
        "                        shuffle=True,\n",
        "                        drop_last=True)\n",
        "\n",
        "img_list = []\n",
        "G_losses = []\n",
        "D_losses = []\n",
        "iters = 0\n",
        "\n",
        "print(\"Starting Training Loop...\")\n",
        "# For each epoch\n",
        "for epoch in range(num_epochs):\n",
        "    for i, batch in enumerate(dataloader):\n",
        "        ############################\n",
        "        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
        "        ###########################\n",
        "        ## Train with all-real batch\n",
        "        model_D.train()\n",
        "        model_G.train()\n",
        "        X, _ = batch  # (B, 3, 64, 64)\n",
        "        X = X.to(device)\n",
        "        model_D.zero_grad()\n",
        "        # create labels (all 1) for Discriminator\n",
        "        label = torch.full((X.shape[0],),\n",
        "                           real_label,\n",
        "                           dtype=torch.float,\n",
        "                           device=device)\n",
        "        # Predictions for real images by D (i.e., [0, 1])\n",
        "        output = model_D(X).view(-1) # (B, 3, 64, 64) -> (B, 1)\n",
        "        # Calculate BCE loss\n",
        "        errD_real = criterion(output, label)\n",
        "        # Calculate gradients for D by backpropagation\n",
        "        errD_real.backward()\n",
        "        D_x = output.mean().item()\n",
        "\n",
        "        ## Train with all-fake batch\n",
        "        # Generate batch of latent vectors\n",
        "        noise = torch.randn(batch_size, n_z, 1, 1, device=device)\n",
        "        # Generate fake image batch with G\n",
        "        fake = model_G(noise)\n",
        "        label.fill_(fake_label)\n",
        "        # Classify all fake batch with D\n",
        "        output = model_D(fake.detach()).view(-1)\n",
        "        # Calculate D's loss on the all-fake batch\n",
        "        errD_fake = criterion(output, label)\n",
        "        # Calculate the gradients for this batch, accumulated (summed) with previous gradients\n",
        "        errD_fake.backward()\n",
        "        D_G_z1 = output.mean().item()\n",
        "        # Compute error of D as sum over the fake and the real batches\n",
        "        errD = errD_real + errD_fake\n",
        "        # Update D\n",
        "        optimizer_D.step()\n",
        "\n",
        "        ############################\n",
        "        # (2) Update G network: maximize log(D(G(z)))\n",
        "        ###########################\n",
        "        model_G.zero_grad()\n",
        "        label.fill_(real_label)  # fake labels are real for generator cost\n",
        "        # Since we just updated D, perform another forward pass of all-fake batch through D\n",
        "        output = model_D(fake).view(-1)\n",
        "        # Calculate G's loss based on this output\n",
        "        errG = criterion(output, label)\n",
        "        # Calculate gradients for G\n",
        "        errG.backward()\n",
        "        D_G_z2 = output.mean().item()\n",
        "        # Update G\n",
        "        optimizer_G.step()  # Note that this does not update D\n",
        "\n",
        "        # Output training stats\n",
        "        if i % 50 == 0:\n",
        "            print('[%d/%d][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\\tD(x): %.4f\\tD(G(z)): %.4f / %.4f'\n",
        "                  % (epoch, num_epochs, i, len(dataloader),\n",
        "                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
        "\n",
        "        # Save Losses for plotting later\n",
        "        G_losses.append(errG.item())\n",
        "        D_losses.append(errD.item())\n",
        "\n",
        "        # Check how the generator is doing by saving G's output on fixed_noise\n",
        "        model_D.eval()\n",
        "        model_G.eval()\n",
        "\n",
        "    # Save generated images for each epoch\n",
        "    with torch.no_grad():\n",
        "        fake = model_G(fixed_noise).detach().cpu()\n",
        "    img_list.append(torchvision.utils.make_grid(fake, padding=2, normalize=True))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fRZrijcF35V5"
      },
      "source": [
        "### Training losses for Generator/Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gvr-_c0sT1Ln"
      },
      "source": [
        "plt.figure(figsize=(10,5))\n",
        "plt.title(\"Generator and Discriminator Loss During Training\")\n",
        "plt.plot(G_losses,label=\"G\")\n",
        "plt.plot(D_losses,label=\"D\")\n",
        "plt.xlabel(\"iterations\")\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w_Mgz8U64AZM"
      },
      "source": [
        "### Visualize generated images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WCVI0UxqUaiE"
      },
      "source": [
        "from IPython.display import HTML\n",
        "import matplotlib.animation as animation\n",
        "fig = plt.figure(figsize=(8,8))\n",
        "plt.axis(\"off\")\n",
        "ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list]\n",
        "ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)\n",
        "\n",
        "HTML(ani.to_jshtml())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D6wFG-4OB_jC"
      },
      "source": [
        "### Download the animation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nTRZUbfwOL_m"
      },
      "source": [
        "filename = \"dcgan_dogs.mp4\" \n",
        "ani.save(filename)\n",
        "\n",
        "from google.colab import files\n",
        "files.download(filename)  # Note: this only works with Google Chrome"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}