DCGAN walkthrough.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "joined-distributor",
   "metadata": {},
   "source": [
    "# Building Complex Computational Graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "robust-johns",
   "metadata": {},
   "source": [
    "## Machine Learning Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wound-closing",
   "metadata": {},
   "source": [
    "Machine learning models are often described as parameterized nonlinear mathematical functions that transform an input into a desired output. A less common view on ML models is from a computational perspective where a model is simply a sequence of computations performed on some data.\n",
    "\n",
    "<img src=\"imgs/model_theoretical_perspective.png\" style=\"float: left;\" width=\"200px\"/>\n",
    "<img src=\"imgs/model_computational_perspective.png\" width=\"200px\"/>\n",
    "\n",
    "Computations performed by a model are usually represented as a directed acyclic graph\n",
    "\n",
    "<img src=\"https://miro.medium.com/max/2434/1*_rCyzi7fQzc_Q1gCqSLM1g.png\" width=\"500px\"/>\n",
    "\n",
    "which can be quite complex for modern ML models. \n",
    "Data is usually represented as multidimensional arrays (tensors)\n",
    "\n",
    "<img src=\"imgs/multidimensional_tensors.png\" width=\"500px\"/>\n",
    "\n",
    "and it is quite common to work with 4D or 5D tensors.\n",
    "The highly abstract nature of these data structures makes the development of ML models a demanding problem where it's hard to find bugs in the code itself.\n",
    "\n",
    "<img src=\"imgs/karpathy_tweet.png\" width=\"500px\"/>\n",
    "\n",
    "## No program is bug-free\n",
    "\n",
    "The standard research and development cycle \n",
    "\n",
    "<br>\n",
    "<img src=\"imgs/research_cycle_ok.png\" width=\"500px\"/>\n",
    "<br>\n",
    "\n",
    "is implicitly dependent on the correcntess of your code\n",
    "\n",
    "<br>\n",
    "<img src=\"imgs/research_cycle_bug.png\" width=\"500px\"/>\n",
    "<br>\n",
    "\n",
    "so instead of a single loop there are two separate ones\n",
    "coupled through your code\n",
    "\n",
    "<br>\n",
    "<img src=\"imgs/research_cycle_chain.png\" width=\"500px\"/>\n",
    "<br>\n",
    "\n",
    "so being able to inspect and reason with the abstract data structures used for implementing an ML model is essential. \n",
    "\n",
    "## Best practices\n",
    "\n",
    "1. Start from the simplest possible model\n",
    "2. Train on a single mini-batch to ensure your code is correct and your model is sufficiently expressive\n",
    "3. Ensure your input data is normalized and encoded appropriately\n",
    "4. Confirm you are using the right loss for your problem\n",
    "5. Check intermediate outputs and connections\n",
    "\n",
    "<span style=\"color: #008cbb\">\n",
    "    <b>At each step visualize as much as you can!</b>\n",
    "</span>\n",
    "\n",
    "## Let's build a DCGAN\n",
    "\n",
    "[Paper](https://arxiv.org/abs/1511.06434)\n",
    "\n",
    "![DCGAN](imgs/dcgan.png)\n",
    "\n",
    "GANs are challenging to train since the optimization procedure is looking for a saddle point rather than a convex optima.\n",
    "\n",
    "![offconvex min max problem](http://www.offconvex.org/assets/GDA_spiral_2.gif)\n",
    "\n",
    "## Getting the code\n",
    "\n",
    "We will train a DCGAN to generate images from the CelebA aligned dataset http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html \n",
    "\n",
    "You can find all code from the following session and the trained models at https://github.com/efemarai/build_dcgan.\n",
    "\n",
    "## Efemarai\n",
    "We will be using [Efemarai](https://efemarai.com) to visualize all steps through the implmentation of the DCGAN. [Efemarai](https://efemarai.com) is a platform for testing and debugging ML code. If you want to try it out just sign up for the *free* Personal tier and start saving tons of hours fighting elusive bugs.\n",
    "\n",
    "## Let's get started"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "final-motivation",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.transforms as tf\n",
    "\n",
    "\n",
    "import efemarai as ef\n",
    "ef.notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "velvet-surgery",
   "metadata": {},
   "source": [
    "Do not forget to run the local `efemarai` daemon."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "understood-breast",
   "metadata": {},
   "source": [
    "## Generator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "requested-mexico",
   "metadata": {},
   "source": [
    "Generator takes in a random vector and should generate an image. For this we will use [transposed convolutions](https://towardsdatascience.com/transposed-convolution-demystified-84ca81b4baba). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sunset-scott",
   "metadata": {},
   "source": [
    "![Transposed Convolution](https://miro.medium.com/max/700/1*faRskFzI7GtvNCLNeCN8cg.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "moderate-territory",
   "metadata": {},
   "source": [
    "Let's see an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "naval-karaoke",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[ 0.0817,  0.1159,  0.1244, -0.0026],\n",
      "          [ 0.0735,  0.0646, -0.0672, -0.0087],\n",
      "          [-0.0300,  0.2063, -0.0606, -0.0689],\n",
      "          [ 0.0209, -0.0029, -0.0517, -0.0402]],\n",
      "\n",
      "         [[ 0.1119, -0.0153,  0.0446,  0.0726],\n",
      "          [-0.0798,  0.0302, -0.0007, -0.0745],\n",
      "          [ 0.1081, -0.0931, -0.0362,  0.0653],\n",
      "          [-0.0148, -0.0075, -0.0170,  0.0770]],\n",
      "\n",
      "         [[-0.0115, -0.1256,  0.0818,  0.0549],\n",
      "          [-0.0907, -0.1071,  0.0591, -0.0342],\n",
      "          [ 0.0119,  0.0514,  0.0248, -0.0155],\n",
      "          [-0.0887,  0.0594, -0.1385, -0.0716]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 0.0344,  0.1427, -0.0692, -0.0444],\n",
      "          [-0.1189,  0.0543, -0.0404,  0.0038],\n",
      "          [-0.0147, -0.0413, -0.0496, -0.0010],\n",
      "          [-0.1101,  0.0212, -0.0489,  0.1232]],\n",
      "\n",
      "         [[-0.0203, -0.0809,  0.0800, -0.0648],\n",
      "          [ 0.0002, -0.2017, -0.0872,  0.0648],\n",
      "          [-0.0676, -0.0369,  0.0148,  0.0029],\n",
      "          [-0.0779, -0.0248,  0.0022,  0.0272]],\n",
      "\n",
      "         [[ 0.0855,  0.0205,  0.0218,  0.1417],\n",
      "          [ 0.0136,  0.1443,  0.1355,  0.1151],\n",
      "          [-0.0241, -0.0235, -0.0216,  0.2125],\n",
      "          [ 0.1211,  0.0991,  0.0690, -0.0774]]]],\n",
      "       grad_fn=<SlowConvTranspose2DBackward>)\n"
     ]
    }
   ],
   "source": [
    "nz = 16\n",
    "noise = torch.randn(1, nz, 1, 1)\n",
    "\n",
    "convT = nn.ConvTranspose2d(\n",
    "    in_channels=nz, \n",
    "    out_channels=64, \n",
    "    kernel_size=4, \n",
    "    stride=1, \n",
    "    padding=0,\n",
    ")\n",
    "\n",
    "print(convT(noise))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "champion-rover",
   "metadata": {},
   "source": [
    "When working with tensors printing tensors rarely gives you useful information. How about something more visual?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "direct-weekly",
   "metadata": {},
   "outputs": [],
   "source": [
    "ef.print(convT(noise))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "african-dependence",
   "metadata": {},
   "source": [
    "Using Efemarai's `print()` function automatically generates a 3D visualizaton of the tensors where you can easily inspect any element or check out the values distribution with a few mouse clicks. Up to 6D tensors are supported."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "quality-california",
   "metadata": {},
   "source": [
    "![Print Tensor](imgs/ef_print.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "solved-death",
   "metadata": {},
   "source": [
    "Seeing the resulting tensor from a computation is useful, but what's even more useful is to explore the computation itself. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "controlled-warehouse",
   "metadata": {},
   "outputs": [],
   "source": [
    "with ef.scan():\n",
    "    output = convT(noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "operational-price",
   "metadata": {},
   "source": [
    "![Graph Scan](imgs/ef_scan.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "approximate-ceiling",
   "metadata": {},
   "source": [
    "Now lets create our generator module starting with the first layer of transformations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "noble-crash",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngf=64):\n",
    "        super(Generator, self).__init__()\n",
    "        self.network = nn.Sequential(\n",
    "            # input is going into a convolution\n",
    "            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 8),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.network(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "allied-redhead",
   "metadata": {},
   "source": [
    "and explore what's happening with the input noise passing through it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "generic-soundtrack",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = Generator(nz)\n",
    "with ef.scan(wait=False):\n",
    "    output = gen(noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "similar-november",
   "metadata": {},
   "source": [
    "The complete generator contains 5 transposed convolution layers and outputs an image of size `(3, 64, 64)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "known-permission",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, ngf=64):\n",
    "        super(Generator, self).__init__()\n",
    "        self.network = nn.Sequential(\n",
    "            # input is going into a convolution\n",
    "            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 8),\n",
    "            nn.ReLU(),\n",
    "            # state size. (ngf*8) x 4 x 4\n",
    "            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 4),\n",
    "            nn.ReLU(),\n",
    "            # state size. (ngf*4) x 8 x 8\n",
    "            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf * 2),\n",
    "            nn.ReLU(),\n",
    "            # state size. (ngf*2) x 16 x 16\n",
    "            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ngf),\n",
    "            nn.ReLU(),\n",
    "            # state size. (ngf) x 32 x 32\n",
    "            nn.ConvTranspose2d(ngf, 3, 4, 2, 1, bias=False),\n",
    "            nn.Tanh()\n",
    "            # state size. (nc) x 64 x 64\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, input):\n",
    "        return self.network(input)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stopped-combat",
   "metadata": {},
   "source": [
    "Here is what the final computational graph of the generator looks like"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "instant-imperial",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = Generator()\n",
    "with ef.scan(wait=False):\n",
    "    output = gen(noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "periodic-trademark",
   "metadata": {},
   "source": [
    "![Generator Graph](imgs/generator_graph.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "second-particular",
   "metadata": {},
   "source": [
    "You can easily confirm that\n",
    "* all the layers are connected as expected\n",
    "* all computations go as expected - there are no NaNs or Infs\n",
    "* the input vector is correctly transformed into a 3x64x64 image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "harmful-stadium",
   "metadata": {},
   "source": [
    "### Overfit to a small batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "congressional-minority",
   "metadata": {},
   "source": [
    "Load images from the CelebA dataset and create a small batch of `(noise, image)` that we are going to overfit to in order to make sure that our generator can generate images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "dress-apparatus",
   "metadata": {},
   "outputs": [],
   "source": [
    "nz = 100\n",
    "\n",
    "dataset = datasets.ImageFolder(\n",
    "    root=\"data\",\n",
    "    transform=tf.Compose([\n",
    "        tf.Resize(64),\n",
    "        tf.CenterCrop(64),\n",
    "        tf.ToTensor(),\n",
    "        tf.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "]))\n",
    "\n",
    "class GeneratorMiniBatch(torch.utils.data.Dataset):\n",
    "    def __init__(self, size): \n",
    "        self.size = size\n",
    "        self.noise = torch.randn(size, nz, 1, 1)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.noise[idx], dataset[idx][0]\n",
    "\n",
    "generator_minibatch = GeneratorMiniBatch(50)   \n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    generator_minibatch, batch_size=10, shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accredited-security",
   "metadata": {},
   "source": [
    "Loop over the minibatch to train and scan the execution periodically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "devoted-banking",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.465360552072525\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.12741057574748993\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.04939030110836029\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.05862098187208176\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.03839258477091789\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.022419793531298637\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.01586761884391308\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.024093030020594597\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.00936425942927599\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "Tensor of shape (512, 256, 4, 4) cannot be visualized with the current GPU.\n",
      "0.013083366677165031\n"
     ]
    }
   ],
   "source": [
    "generator = Generator()\n",
    "\n",
    "optimizer = optim.Adam(\n",
    "    generator.parameters(), lr=1e-3, \n",
    ")\n",
    "\n",
    "iteration = 0\n",
    "for epoch in range(100):\n",
    "    for noise, image in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        with ef.scan(\n",
    "            iteration, \n",
    "            enabled=iteration % 50 == 0, \n",
    "            wait=False,\n",
    "        ):\n",
    "            output = generator(noise)\n",
    "            loss = (output - image).square().mean()\n",
    "            loss.backward()\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        iteration += 1\n",
    "        \n",
    "    if epoch % 10 == 0:\n",
    "        print(loss.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addressed-uruguay",
   "metadata": {},
   "source": [
    "## Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "indirect-diamond",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, ndf=64):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.network = nn.Sequential(\n",
    "            # input is (nc) x 64 x 64\n",
    "            nn.Conv2d(3, ndf, 4, 2, 1, bias=False),\n",
    "            nn.LeakyReLU(0.2, inplace=False),\n",
    "            # state size. (ndf) x 32 x 32\n",
    "            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 2),\n",
    "            nn.LeakyReLU(0.2, inplace=False),\n",
    "            # state size. (ndf*2) x 16 x 16\n",
    "            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 4),\n",
    "            nn.LeakyReLU(0.2, inplace=False),\n",
    "            # state size. (ndf*4) x 8 x 8\n",
    "            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),\n",
    "            nn.BatchNorm2d(ndf * 8),\n",
    "            nn.LeakyReLU(0.2, inplace=False),\n",
    "            # state size. (ndf*8) x 4 x 4\n",
    "            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "\n",
    "    def forward(self, input):\n",
    "        assert torch.all(-1.0 <= input) and torch.all(input <= 1.0)\n",
    "        return self.network(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nasty-fourth",
   "metadata": {},
   "source": [
    "### Overfit to a small batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "wireless-hypothesis",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiscriminatorMiniBatch(torch.utils.data.Dataset): \n",
    "    def __init__(self):\n",
    "        self.size = 2 * len(generator_minibatch)\n",
    "        self.fakes = generator(generator_minibatch.noise).detach()\n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if idx < len(self.fakes):\n",
    "            return self.fakes[idx], torch.zeros(1, dtype=torch.float).squeeze()\n",
    "        else:        \n",
    "            _, image = generator_minibatch[idx - len(self.fakes)]\n",
    "            return image, torch.ones(1, dtype=torch.float).squeeze()\n",
    "\n",
    "\n",
    "discriminator_minibatch = DiscriminatorMiniBatch()   \n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    discriminator_minibatch, batch_size=10, shuffle=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "common-sauce",
   "metadata": {},
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-171-9175a492f485>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     10\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 12\u001b[0;31m         \u001b[0;32massert\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0.0\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m1.0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     13\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdiscriminator\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mview\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "discriminator = Discriminator()\n",
    "\n",
    "optimizer = optim.Adam(\n",
    "    discriminator.parameters(), lr=1e-4, \n",
    ")\n",
    "\n",
    "iteration = 0\n",
    "for epoch in range(50):\n",
    "    for image, label in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = discriminator(image).view(-1)\n",
    "        loss = F.binary_cross_entropy(output, label)\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "        \n",
    "        iteration += 1\n",
    "        \n",
    "    if epoch % 10 == 0:\n",
    "        print(loss.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "informal-artwork",
   "metadata": {},
   "source": [
    "# Train both networks\n",
    "\n",
    "The GAN training is dependent on balancing the min max criterion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "final-discount",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Training Loop\n",
    "ef.deregister_assertions(ef.assertions.NoNonZeroGradientsAssertion)\n",
    "print(\"Starting Training Loop...\")\n",
    "discriminator = Discriminator()\n",
    "generator = Generator()\n",
    "\n",
    "# From the current generator we assumed\n",
    "fake_label = 0.0\n",
    "real_label = 1.0\n",
    "\n",
    "batch_size = 10\n",
    "\n",
    "dataloader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "# Setup Adam optimizers for both G and D\n",
    "optimizerD = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "optimizerG = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
    "\n",
    "for epoch in range(5):\n",
    "    # For each batch in the dataloader\n",
    "    for i, data in enumerate(dataloader, 0):\n",
    "        ef.inspect(data[0])\n",
    "        with ef.scan(i):\n",
    "            ############################\n",
    "            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "            ###########################\n",
    "            ## Train with all-real batch\n",
    "\n",
    "            discriminator.zero_grad()\n",
    "            # Format batch\n",
    "            label = torch.full((batch_size,), real_label, dtype=torch.float)\n",
    "            # Forward pass real batch through D\n",
    "            output = discriminator(data[0]).view(-1)\n",
    "            # Calculate loss on all-real batch\n",
    "            errD_real = F.binary_cross_entropy(output, label)\n",
    "\n",
    "            ## Train with all-fake batch\n",
    "            # Generate batch of latent vectors\n",
    "            noise = torch.randn(batch_size, nz, 1, 1)\n",
    "            # Generate fake image batch with G\n",
    "            fake = generator(noise)\n",
    "            label = torch.full((batch_size,), fake_label, dtype=torch.float)\n",
    "            # Classify all fake batch with D\n",
    "            output = discriminator(fake.detach()).view(-1)\n",
    "            # Calculate D's loss on the all-fake batch\n",
    "            errD_fake = F.binary_cross_entropy(output, label)\n",
    "\n",
    "            errD = errD_real + errD_fake\n",
    "            errD.backward()\n",
    "\n",
    "        with ef.scan(i):\n",
    "            # Update D\n",
    "            optimizerD.step()\n",
    "\n",
    "            ############################\n",
    "            # (2) Update G network: maximize log(D(G(z)))\n",
    "            ###########################\n",
    "            generator.zero_grad()\n",
    "\n",
    "            # fake labels are real for generator cost\n",
    "            label = torch.full((batch_size,), real_label, dtype=torch.float)\n",
    "            # Since we just updated D, perform another forward pass of all-fake batch through D\n",
    "            output = discriminator(fake).view(-1)\n",
    "            # Calculate G's loss based on this output\n",
    "            errG = F.binary_cross_entropy(output, label)\n",
    "            # Calculate gradients for G\n",
    "            errG.backward()\n",
    "\n",
    "            # Update G\n",
    "            optimizerG.step()\n",
    "\n",
    "        # Output training stats\n",
    "        if i % 100 == 0:\n",
    "            print(\"[%d/5][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\"\n",
    "                % (epoch, i, len(dataloader), errD.item(), errG.item()))\n",
    "\n",
    "print('Done')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conceptual-genius",
   "metadata": {},
   "source": [
    "# Let's see the training loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "alike-geometry",
   "metadata": {},
   "source": [
    "![Training loss](models/training_loss.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "surface-salem",
   "metadata": {},
   "source": [
    "![Image Comparison - Real vs Fake](models/real_vs_gen_images.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "academic-pointer",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}