GAN/pytorch_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "pytorch_GAN.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vKBPRI8vjxCS"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets\n",
        "from torchvision import transforms\n",
        "from torchvision.utils import save_image\n",
        "\n",
        "import numpy as np\n",
        "import datetime\n",
        "import os, sys\n",
        "\n",
        "from matplotlib.pyplot import imshow, imsave\n",
        "%matplotlib inline"
      ],
      "metadata": {
        "id": "PKQrtz1pj0S6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "MODEL_NAME = 'VanillaGAN'\n",
        "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
      ],
      "metadata": {
        "id": "X8KlUXsbj2LW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_sample_image(G, n_noise):\n",
        "    \"\"\"\n",
        "        save sample 100 images\n",
        "    \"\"\"\n",
        "    z = torch.randn(100, n_noise).to(DEVICE)\n",
        "    y_hat = G(z).view(100, 28, 28) # (100, 28, 28)\n",
        "    result = y_hat.cpu().data.numpy()\n",
        "    img = np.zeros([280, 280])\n",
        "    for j in range(10):\n",
        "        img[j*28:(j+1)*28] = np.concatenate([x for x in result[j*10:(j+1)*10]], axis=-1)\n",
        "    return img"
      ],
      "metadata": {
        "id": "iR8FI_-8kAF9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Discriminator(nn.Module):\n",
        "    \"\"\"\n",
        "        Simple Discriminator w/ MLP\n",
        "    \"\"\"\n",
        "    def __init__(self, input_size=784, num_classes=1):\n",
        "        super(Discriminator, self).__init__()\n",
        "        self.layer = nn.Sequential(\n",
        "            nn.Linear(input_size, 512),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(512, 256),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(256, num_classes),\n",
        "            nn.Sigmoid(),\n",
        "        )\n",
        "    \n",
        "    def forward(self, x):\n",
        "        y_ = x.view(x.size(0), -1)\n",
        "        y_ = self.layer(y_)\n",
        "        return y_"
      ],
      "metadata": {
        "id": "gZ_OdvAekACm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Generator(nn.Module):\n",
        "    \"\"\"\n",
        "        Simple Generator w/ MLP\n",
        "    \"\"\"\n",
        "    def __init__(self, input_size=100, num_classes=784):\n",
        "        super(Generator, self).__init__()\n",
        "        self.layer = nn.Sequential(\n",
        "            nn.Linear(input_size, 128),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(128, 256),\n",
        "            nn.BatchNorm1d(256),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(256, 512),\n",
        "            nn.BatchNorm1d(512),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(512, 1024),\n",
        "            nn.BatchNorm1d(1024),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            nn.Linear(1024, num_classes),\n",
        "            nn.Tanh()\n",
        "        )\n",
        "        \n",
        "    def forward(self, x):\n",
        "        y_ = self.layer(x)\n",
        "        y_ = y_.view(x.size(0), 1, 28, 28)\n",
        "        return y_"
      ],
      "metadata": {
        "id": "MYFIiq9ikAAX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "n_noise = 100"
      ],
      "metadata": {
        "id": "wOLoA-15j_9t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "D = Discriminator().to(DEVICE)\n",
        "G = Generator(n_noise).to(DEVICE)"
      ],
      "metadata": {
        "id": "HPUmnLAOj_7R"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "transform = transforms.Compose([transforms.ToTensor(),\n",
        "                                transforms.Normalize(mean=[0.5],\n",
        "                                std=[0.5])]\n",
        ")\n",
        "\n",
        "mnist = datasets.MNIST(root='../data/', train = True, transform=transform, download=True)\n",
        "batch_size = 64"
      ],
      "metadata": {
        "id": "CIcqI3DQj_44"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data_loader = DataLoader(dataset=mnist, batch_size = batch_size, shuffle=True, drop_last = True)"
      ],
      "metadata": {
        "id": "YGVEq4MNj_0d"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = nn.BCELoss()\n",
        "D_opt = torch.optim.Adam(D.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
        "G_opt = torch.optim.Adam(G.parameters(), lr=0.0002, betas=(0.5, 0.999))"
      ],
      "metadata": {
        "id": "JmGnna5hkbK3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "max_epoch = 50\n",
        "step = 0\n",
        "n_critic = 1\n",
        "\n",
        "D_labels = torch.ones(batch_size,1).to(DEVICE)\n",
        "D_fakes = torch.zeros(batch_size,1).to(DEVICE)\n",
        "\n",
        "if not os.path.exists('samples'):\n",
        "    os.makedirs('samples')"
      ],
      "metadata": {
        "id": "couCiARhkmOj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(max_epoch):\n",
        "    for idx, (images, _) in enumerate(data_loader):\n",
        "        # Training Discriminator\n",
        "        x = images.to(DEVICE)\n",
        "        x_outputs = D(x)\n",
        "        D_x_loss = criterion(x_outputs, D_labels)\n",
        "\n",
        "        z = torch.randn(batch_size, n_noise).to(DEVICE)\n",
        "        z_outputs = D(G(z))\n",
        "        D_z_loss = criterion(z_outputs, D_fakes)\n",
        "        D_loss = D_x_loss + D_z_loss\n",
        "        \n",
        "        D.zero_grad()\n",
        "        D_loss.backward()\n",
        "        D_opt.step()\n",
        "\n",
        "        if step % n_critic == 0:\n",
        "            # Training Generator\n",
        "            z = torch.randn(batch_size, n_noise).to(DEVICE)\n",
        "            z_outputs = D(G(z))\n",
        "            G_loss = criterion(z_outputs, D_labels)\n",
        "\n",
        "            G.zero_grad()\n",
        "            G_loss.backward()\n",
        "            G_opt.step()\n",
        "        \n",
        "        if step % 500 == 0:\n",
        "            print('Epoch: {}/{}, Step: {}, D Loss: {}, G Loss: {}'.format(epoch, max_epoch, step, D_loss.item(), G_loss.item()))\n",
        "        \n",
        "        if step % 1000 == 0:\n",
        "            G.eval()\n",
        "            img = get_sample_image(G, n_noise)\n",
        "            imsave('samples/{}_step{}.jpg'.format(MODEL_NAME, str(step).zfill(3)), img, cmap='gray')\n",
        "            G.train()\n",
        "        step += 1"
      ],
      "metadata": {
        "id": "aWV1RmhXk6os"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "G.eval()\n",
        "imshow(get_sample_image(G, n_noise), cmap='gray')"
      ],
      "metadata": {
        "id": "LwjAClvQlEyi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(D.state_dict(), 'D.pkl')\n",
        "torch.save(G.state_dict(), 'G.pkl')"
      ],
      "metadata": {
        "id": "CVcL5qu_lKLw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "eQcKbPqJlK5x"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}