WGAN_empty_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "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.5.2"
    },
    "colab": {
      "name": "WGAN_empty.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlelarge/dataflowr/blob/master/WGAN_empty_colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sdkhLNACHYkt",
        "colab_type": "text"
      },
      "source": [
        "# Wasserstein GAN\n",
        "\n",
        "Wasserstein GAN is a development on the famous GAN which avoids vanishing gradient issues by comparing the generated and real data distribution with Earth Mover distance.\n",
        "\n",
        "For a brief overview, refer to [this](https://paper.dropbox.com/doc/Wasserstein-GAN--AnSk4vkryFmJgICMb_fybpwHAg-GvU0p2V9ThzdwY3BbhoP7).\n",
        "\n",
        "For a more in-depth discussion, refer this [blogpost](https://www.alexirpan.com/2017/02/22/wasserstein-gan.html) or the [paper](https://arxiv.org/abs/1701.07875).\n",
        "\n",
        "Here we will implement a WGAN for MNIST.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VYG36si3ZNYy",
        "colab_type": "text"
      },
      "source": [
        "In practice, there are only really two things that change in WGAN compared to regular GAN :\n",
        "\n",
        "1.   Instead of having a discriminator, which outputs a probability, we have a critic, which outputs a score. Hence there is no need of sigmoid at the output, and no log in the loss\n",
        "2.   To ensure that the function represented by the critic is lipschitz, we clip the weights of the critic.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_oMgHF27HZoM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.autograd import Variable\n",
        "from torch.optim import RMSprop\n",
        "from torchvision import transforms\n",
        "from torchvision.utils import make_grid\n",
        "from torchvision.datasets import CIFAR10, MNIST\n",
        "from pylab import plt\n",
        "from tqdm import tqdm_notebook\n",
        "%matplotlib inline"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pkhZGrNiygHj",
        "colab_type": "code",
        "outputId": "d1369811-62b0-42d7-ba5f-d5e40cbd9c73",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 33
        }
      },
      "source": [
        "class Config:\n",
        "    lr = 0.00005\n",
        "    nz = 100 # noise dimension\n",
        "    image_size = 64 # we will resize all images to 64*64 using torch.transforms\n",
        "    nc = 1 # number of channels of the image (1 for MNIST, 3 for CIFAR10)\n",
        "    ngf = 64 # number of channels of generator\n",
        "    ndf = 64 # number of channels of discriminator\n",
        "    batch_size = 32\n",
        "    max_epoch = 50 # =1 when debug\n",
        "    gpu = torch.cuda.is_available() # use gpu or not\n",
        "    clamp_num=0.01 # WGAN gradient clipping parameter\n",
        "    \n",
        "opt=Config()\n",
        "print('Using gpu : ', opt.gpu)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using gpu :  True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eUqd-AigygHl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# data preprocess\n",
        "\n",
        "transform=transforms.Compose([\n",
        "                transforms.Resize(opt.image_size) ,\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize([0.5],[0.5])\n",
        "                ])\n",
        "\n",
        "dataset=MNIST(root='.',transform=transform,download=True)\n",
        "dataloader=torch.utils.data.DataLoader(dataset, opt.batch_size, shuffle = True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0tos5KoHygHo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "''' REMINDER : nn.Conv2d takes as arguments \n",
        "nn.Conv2D(input channels, output channels, kernel_size, stride, padding)\n",
        "\n",
        "spatial dimension after the convolution is given by  \n",
        "output_dim = (input_dim + 2*padding - kernel_size) / stride + 1\n",
        "\n",
        "nn.Conv2d does the exact opposite (upsampling rather than downsampling)\n",
        "'''\n",
        "\n",
        "netd = nn.Sequential(\n",
        "            # layer 1 : spatial dimension 64 -> 32\n",
        "            nn.Conv2d(?, opt.ndf, ?, ?, ?,bias=False), \n",
        "            nn.LeakyReLU(0.2,inplace=True), \n",
        "            \n",
        "            # layer 2 : spatial dimension 32 -> 16\n",
        "            nn.Conv2d(opt.ndf,opt.ndf*2, ?, ?, ?, bias=False), \n",
        "            nn.BatchNorm2d(opt.ndf*2),\n",
        "            nn.LeakyReLU(0.2,inplace=True),\n",
        "            \n",
        "            # layer 3 : spatial dimension 16 -> 8\n",
        "            nn.Conv2d(opt.ndf*2,opt.ndf*4, ?, ?, ?, bias=False), \n",
        "            nn.BatchNorm2d(opt.ndf*4),\n",
        "            nn.LeakyReLU(0.2,inplace=True),\n",
        "            \n",
        "            # layer 4 : spatial dimension 8 -> 4\n",
        "            nn.Conv2d(opt.ndf*4,opt.ndf*8, ?, ?, ?, bias=False), \n",
        "            nn.BatchNorm2d(opt.ndf*8),\n",
        "            nn.LeakyReLU(0.2,inplace=True),\n",
        "            \n",
        "            # layer 5 : spatial dimension 4 -> 1\n",
        "            nn.Conv2d(opt.ndf*8, ?, ?, ?, ?,bias=False), \n",
        "            # This is a \"critic\", not a discriminator, so no need of a sigmoid !\n",
        "        )\n",
        "\n",
        "netg = nn.Sequential(\n",
        "            nn.ConvTranspose2d(?,opt.ngf*8, ?, ?, ?,bias=False),\n",
        "            nn.BatchNorm2d(opt.ngf*8),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            nn.ConvTranspose2d(opt.ngf*8,opt.ngf*4, ?, ?, ?, bias=False),\n",
        "            nn.BatchNorm2d(opt.ngf*4),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            nn.ConvTranspose2d(opt.ngf*4,opt.ngf*2, ?, ?, ?, bias=False),\n",
        "            nn.BatchNorm2d(opt.ngf*2),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            nn.ConvTranspose2d(opt.ngf*2,opt.ngf, ?, ?, ?, bias=False),\n",
        "            nn.BatchNorm2d(opt.ngf),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            nn.ConvTranspose2d(opt.ngf, ?, ?, ?, ?, bias=False),\n",
        "            nn.Tanh() \n",
        "            # We are outputting images so the output should be in [-1,1] !\n",
        "        )\n",
        "\n",
        "\n",
        "def weight_init(m):\n",
        "    # weight_initialization: important for wgan\n",
        "    class_name=m.__class__.__name__\n",
        "    if class_name.find('Conv')!=-1:\n",
        "        m.weight.data.normal_(0,0.02)\n",
        "    elif class_name.find('Norm')!=-1:\n",
        "        m.weight.data.normal_(1.0,0.02)\n",
        "#     else:print(class_name)\n",
        "\n",
        "if opt.gpu:\n",
        "    netd.cuda()\n",
        "    netg.cuda()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vWnm92doygHq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# optimizer : Use torch.optim.RMSprop instead of torch.optim.Adam\n",
        "optimizerD = ?\n",
        "optimizerG = ?"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uG4hrcxPp_qQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# reset the weights\n",
        "netd.apply(weight_init)\n",
        "netg.apply(weight_init)\n",
        "lossD_history = []\n",
        "lossG_history = []"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "R8719q3lygHr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'''\n",
        "The training is very long, you can stop it and resume whenever you want !\n",
        "Just look at how the generated images evolve, and stop when you think it's not getting any better.\n",
        "You can resume training afterwards if you wish.\n",
        "''' begin training\n",
        "\n",
        "log_every = 100\n",
        "\n",
        "def train():\n",
        "\n",
        "  print('Start training')\n",
        "\n",
        "  fix_noise = torch.FloatTensor(opt.batch_size,opt.nz,1,1).normal_(0,1) \n",
        "  # we will see how the generator reconstructs this noise every 100 steps\n",
        "  if opt.gpu:\n",
        "      fix_noise = fix_noise.cuda()\n",
        "\n",
        "  for epoch in range(opt.max_epoch):\n",
        "      for ii, data in enumerate(dataloader):\n",
        "\n",
        "        # ----- train netd -----\n",
        "        # your code here\n",
        "        # remember to clip the parameters using param.clamp_(min, max) !\n",
        "        # \n",
        "\n",
        "        # ------ train netg -------\n",
        "        # your code here\n",
        "        # \n",
        "\n",
        "        # log every 100 steps\n",
        "        if ii%log_every==0:\n",
        "          print('LossD = {}, LossG = {}'.format(lossD, lossG))\n",
        "          fake = netg(fix_noise)\n",
        "          imgs = make_grid(fake.data*0.5+0.5).cpu() \n",
        "          plt.imshow(imgs.permute(1,2,0).numpy()) \n",
        "          plt.show()\n",
        "          \n",
        "  return\n",
        "          \n",
        "train()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gDrNPCBVwTQc",
        "colab_type": "text"
      },
      "source": [
        "### Visualization"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pn69kf1VwSC7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "noise = torch.randn(64,opt.nz,1,1)\n",
        "if opt.gpu:\n",
        "  noise = noise.cuda()\n",
        "fake=netg(noise)\n",
        "imgs = make_grid(fake.data*0.5+0.5).cpu()\n",
        "plt.figure(figsize=(10,10))\n",
        "plt.imshow(imgs.permute(1,2,0).numpy())\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDQiSsPsJF2P",
        "colab_type": "text"
      },
      "source": [
        "### Plot the losses of discriminator and generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UKHM3HbFbQCV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plt.plot(lossG_history, label='Generator')\n",
        "plt.plot(lossD_history, label='Discriminator')\n",
        "plt.legend()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VHZAf5ean2T0",
        "colab_type": "text"
      },
      "source": [
        "### Please comment the loss curves"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0_MsMPD6n8Zy",
        "colab_type": "text"
      },
      "source": [
        "Your answer here"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aZe5FTngn96E",
        "colab_type": "text"
      },
      "source": [
        "The training involves a competition between the critic and the generator.\n",
        "\n",
        "If one of the two is winning the fight too easily (loss curve going down much quicker), we can weaken it by updating its weights less often.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5v4v054jokkM",
        "colab_type": "text"
      },
      "source": [
        "### Suggest a modification to your previous code to fix this and see how results are affected"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jegSk6HKomYZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Remember to reset the weights !\n",
        "netd.apply(weight_init)\n",
        "netg.apply(weight_init)\n",
        "lossD_history = []\n",
        "lossG_history = []\n",
        "\n",
        "def train_modified():\n",
        "  # your code here\n",
        "  \n",
        "train_modified()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "joDB1YNLIguq",
        "colab_type": "text"
      },
      "source": [
        "### Now give the WGAN a quick try on CIFAR10 !"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q8MHvdUbInVs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "transform=transforms.Compose([\n",
        "                transforms.Resize(opt.image_size) ,\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize([0.5]*3,[0.5]*3)\n",
        "                ])\n",
        "\n",
        "dataset=CIFAR(root='.',transform=transform,download=True)\n",
        "dataloader=torch.utils.data.DataLoader(dataset, opt.batch_size, shuffle = True)\n",
        "\n",
        "opt.nc = ?\n",
        "netd = ?\n",
        "netg = ?\n",
        "\n",
        "netd.apply(weight_init)\n",
        "netg.apply(weight_init)\n",
        "lossD_history = []\n",
        "lossG_history = []\n",
        "\n",
        "train_modified()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0lNUw5gftHyx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}