Lab3_MNIST_cGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Lab3_MNIST_cGAN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/illhyhl1111/SNU_ML2019/blob/master/Lab3_MNIST_cGAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-sbkFbNUAKyK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from torchvision.datasets import MNIST\n",
        "from torchvision.transforms import Compose,ToTensor\n",
        "\n",
        "from google.colab import drive\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GXG_qmGXJjwV",
        "colab_type": "text"
      },
      "source": [
        "### GAN\n",
        "#### Generator\n",
        "- Generates fake images\n",
        "- Get labels to generates specific images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rmAkE7-J72hU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Generator(nn.Module):\n",
        "    def __init__(self,latent_size):\n",
        "        super(Generator,self).__init__()\n",
        "        self.latent_size = latent_size\n",
        "\n",
        "        self.gen_fc1 = nn.Linear(self.latent_size,128)\n",
        "        self.label_fc = nn.Linear(10,128)\n",
        "        self.gen_fc2 = nn.Linear(256,512)\n",
        "        self.gen_fc3 = nn.Linear(512,784)\n",
        "\n",
        "    def generator(self,label):\n",
        "        batch_size = label.shape[0]\n",
        "        z = torch.randn((batch_size,self.latent_size))\n",
        "        if torch.cuda.is_available():\n",
        "            z = z.cuda()\n",
        "\n",
        "        gen = self.gen_fc1(z)\n",
        "        gen_label = self.label_fc(label)\n",
        "        gen = torch.cat((gen,gen_label),-1)\n",
        "        gen = torch.relu(gen)\n",
        "        gen = self.gen_fc2(gen)\n",
        "        gen = torch.relu(gen)\n",
        "        gen = self.gen_fc3(gen)\n",
        "        img_ = torch.sigmoid(gen).view(-1,1,28,28)\n",
        "        return img_\n",
        "\n",
        "    def forward(self,label):\n",
        "        img = self.generator(label)\n",
        "        return img"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dJWBXpOyJ2oW",
        "colab_type": "text"
      },
      "source": [
        "#### Discriminator\n",
        "- Find fake images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iHG8DS1x2YO1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Discriminator(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Discriminator,self).__init__()\n",
        "\n",
        "        self.dis_fc1 = nn.Linear(784,512)\n",
        "        self.dis_fc2 = nn.Linear(512,128)\n",
        "        self.label_fc = nn.Linear(10,128)\n",
        "        self.dis_fc3 = nn.Linear(256,32)\n",
        "        self.dis_fc4 = nn.Linear(32,1)\n",
        "\n",
        "    def discriminator(self,img,label):\n",
        "        batch_size = img.shape[0]\n",
        "        img = img.view(batch_size,-1)\n",
        "\n",
        "        dis = self.dis_fc1(img)\n",
        "        dis = torch.relu(dis)\n",
        "        dis = self.dis_fc2(dis)\n",
        "        dis = torch.relu(dis)\n",
        "        dis_label = self.label_fc(label)\n",
        "        dis = torch.cat((dis,dis_label),-1)\n",
        "        dis = self.dis_fc3(dis)\n",
        "        dis = torch.relu(dis)\n",
        "        dis = self.dis_fc4(dis)\n",
        "        dis = torch.sigmoid(dis)\n",
        "        return dis\n",
        "\n",
        "    def forward(self,img,label):\n",
        "        criterion = self.discriminator(img,label)\n",
        "        return criterion"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Vyo66_AKAXC",
        "colab_type": "text"
      },
      "source": [
        "#### Util functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VuNTK2GI6nEn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def one_hot(x):\n",
        "    res = torch.zeros((x.shape[0],10))\n",
        "    for i in range(x.shape[0]):\n",
        "        res[i,x[i]]=1\n",
        "    return res"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ED6guRZjKB6M",
        "colab_type": "text"
      },
      "source": [
        "#### Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UcXaqkwV6ztF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "NUM_EPOCHS = 150\n",
        "BATCH_SIZE = 128\n",
        "LATENT_SIZE = 100\n",
        "LEARNING_RATE = 2e-4\n",
        "EPS = 1e-6\n",
        "LEARNING_RATIO = 1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NMYF_p3yKFgU",
        "colab_type": "text"
      },
      "source": [
        "#### Dataloader"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IpnCYlVHBuSW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "transforms = Compose([\n",
        "    ToTensor(),\n",
        "])\n",
        "\n",
        "trainset = MNIST('./',train=True,transform=transforms,download=True)\n",
        "testset = MNIST('./',train=False,transform=transforms,download=True)\n",
        "\n",
        "args = {\n",
        "    'num_workers' : 1,\n",
        "    'batch_size' : BATCH_SIZE,\n",
        "    'shuffle' : True,\n",
        "}\n",
        "\n",
        "train_loader = DataLoader(trainset,**args)\n",
        "test_loader = DataLoader(testset,**args)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jJEn6myqKItd",
        "colab_type": "text"
      },
      "source": [
        "#### Training part"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wDkcslck6wgQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "gen = Generator(LATENT_SIZE)\n",
        "dis = Discriminator()\n",
        "if torch.cuda.is_available():\n",
        "    gen = gen.cuda()\n",
        "    dis = dis.cuda()\n",
        "\n",
        "gen_parameters = filter(lambda p: p.requires_grad, gen.parameters())\n",
        "num_params = sum([np.prod(p.size()) for p in gen_parameters])\n",
        "print(\"number of parameters(generator) : {}\".format(num_params))\n",
        "dis_parameters = filter(lambda p: p.requires_grad, dis.parameters())\n",
        "num_params = sum([np.prod(p.size()) for p in dis_parameters])\n",
        "print(\"number of parameters(discriminator) : {}\".format(num_params))\n",
        "\n",
        "optimizer_G = Adam(gen.parameters(),lr=LEARNING_RATE)\n",
        "optimizer_D = Adam(dis.parameters(),lr=LEARNING_RATE)\n",
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "    tot_gen_loss = 0\n",
        "    tot_dis_loss = 0\n",
        "    cnt = 0\n",
        "    for img,label in train_loader:\n",
        "        onehot_label = one_hot(label)\n",
        "        if torch.cuda.is_available():\n",
        "            img = img.cuda()\n",
        "            onehot_label = onehot_label.cuda()\n",
        "        \n",
        "        if cnt % LEARNING_RATIO == 0:\n",
        "            optimizer_D.zero_grad()\n",
        "            real = dis(img,onehot_label)\n",
        "            img_ = gen(onehot_label)\n",
        "            fake = dis(img_,onehot_label)\n",
        "            dis_loss = -torch.mean(torch.log(real+EPS)+torch.log(1+EPS-fake))\n",
        "            dis_loss.backward()\n",
        "            optimizer_D.step()\n",
        "\n",
        "        optimizer_G.zero_grad()\n",
        "        img_ = gen(onehot_label)\n",
        "        fake = dis(img_,onehot_label)\n",
        "        gen_loss = -torch.mean(torch.log(fake+EPS))\n",
        "        gen_loss.backward()\n",
        "        optimizer_G.step()\n",
        "        tot_dis_loss +=dis_loss.item()\n",
        "        tot_gen_loss +=gen_loss.item()\n",
        "        cnt+=1\n",
        "    \n",
        "    print(\"epoch {}\\n gen : {:.5f}\\t dis : {:.5f}\".format(epoch+1,tot_gen_loss,tot_dis_loss))\n",
        "    img, label = next(iter(test_loader))\n",
        "    onehot_label = one_hot(label)\n",
        "    if torch.cuda.is_available():\n",
        "        onehot_label = onehot_label.cuda()\n",
        "    img_ = gen(onehot_label)\n",
        "    fig = plt.figure(figsize=(3,3))\n",
        "    fig.add_subplot(1,2,1)\n",
        "    plt.imshow(img[0].squeeze().numpy(),cmap='gray')\n",
        "    plt.axis('off')\n",
        "    img_out = img_[0].squeeze().cpu().detach().numpy()\n",
        "    fig.add_subplot(1,2,2)\n",
        "    plt.imshow(img_out,cmap='gray')\n",
        "    plt.axis('off')\n",
        "    plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}