WGAN-MNIST_safwankdb.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "WGAN.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "JypOqJliDIaY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "from torchvision import transforms\n",
        "from torchvision.datasets import MNIST\n",
        "from torchvision.utils import make_grid\n",
        "from torch.utils.data import DataLoader\n",
        "import imageio"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZB43sCZSEyHe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "transform = transforms.Compose([\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize((0.5,),(0.5,))\n",
        "                ])\n",
        "to_image = transforms.ToPILImage()\n",
        "trainset = MNIST(root='./data/', train=True, download=True, transform=transform)\n",
        "trainloader = DataLoader(trainset, batch_size=64, shuffle=True)\n",
        "\n",
        "device = 'cuda'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lynoz06jF3nJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Generator(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Generator, self).__init__()\n",
        "        self.n_features = 100\n",
        "        self.n_out = 784\n",
        "        self.model = nn.Sequential(\n",
        "                    nn.Linear(self.n_features, 256),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Linear(256, 512),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Linear(512, 1024),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Linear(1024, self.n_out),\n",
        "                    nn.Tanh()\n",
        "                    )\n",
        "    def forward(self, x):\n",
        "        x = self.model(x)\n",
        "        x = x.view(-1, 1, 28, 28)\n",
        "        return x\n",
        "\n",
        "class Critic(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Critic, self).__init__()\n",
        "        self.n_in = 784\n",
        "        self.n_out = 1\n",
        "        self.model = nn.Sequential(\n",
        "                    nn.Linear(self.n_in, 1024),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Dropout(0.3),\n",
        "                    nn.Linear(1024, 512),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Dropout(0.3),\n",
        "                    nn.Linear(512, 256),\n",
        "                    nn.LeakyReLU(0.2),\n",
        "                    nn.Dropout(0.3),\n",
        "                    nn.Linear(256, self.n_out)\n",
        "                    )\n",
        "    def forward(self, x):\n",
        "        x = x.view(-1, 784)\n",
        "        x = self.model(x)\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-LB03abLJUy3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "generator = Generator()\n",
        "critic = Critic()\n",
        "\n",
        "generator.to(device)\n",
        "critic.to(device)\n",
        "\n",
        "alpha = 5e-5\n",
        "g_optim = optim.RMSprop(generator.parameters(), lr=alpha)\n",
        "c_optim = optim.RMSprop(critic.parameters(), lr=alpha)\n",
        "\n",
        "g_losses = []\n",
        "c_losses = []\n",
        "images = []\n",
        "\n",
        "def noise(n, n_features=100):\n",
        "    return torch.randn(n, n_features).to(device)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n5qn3y44KE8q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train_critic(optimizer, real_data, fake_data, c=0.01):\n",
        "    optimizer.zero_grad()\n",
        "    error_real = critic(real_data).mean()\n",
        "    error_fake = critic(fake_data).mean()\n",
        "    total_error = -(error_real - error_fake)\n",
        "    total_error.backward()\n",
        "    optimizer.step()\n",
        "    for p in critic.parameters():\n",
        "        p.data.clamp_(-c, c)\n",
        "    return -total_error\n",
        "\n",
        "def train_generator(optimizer, fake_data):\n",
        "    optimizer.zero_grad()\n",
        "    error = -critic(fake_data).mean()\n",
        "    error.backward()\n",
        "    optimizer.step()\n",
        "    return error"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b_q1TRkaKmEk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "num_epochs = 250\n",
        "n_critic = 5\n",
        "test_noise = noise(64)\n",
        "\n",
        "generator.train()\n",
        "critic.train()\n",
        "for epoch in range(num_epochs):\n",
        "    g_error = 0.0\n",
        "    c_error = 0.0\n",
        "    for i, data in enumerate(trainloader):\n",
        "        imgs, _ = data\n",
        "        n = len(imgs)\n",
        "        fake_data = generator(noise(n)).detach()\n",
        "        real_data = imgs.to(device)\n",
        "        c_error += train_critic(c_optim, real_data, fake_data)\n",
        "        \n",
        "        if (i+1)%n_critic==0:\n",
        "            fake_data = generator(noise(n))\n",
        "            g_error += train_generator(g_optim, fake_data)\n",
        "    if epoch%5==0:\n",
        "        img = generator(test_noise).cpu().detach()\n",
        "        img = make_grid(img)\n",
        "        images.append(to_image(img))\n",
        "        g_losses.append(g_error)\n",
        "        c_losses.append(c_error)\n",
        "        plt.clf()\n",
        "        plt.plot(g_losses, label='Generator Losses')\n",
        "        plt.plot(c_losses, label='Critic Losses')\n",
        "        plt.legend()\n",
        "        plt.savefig('loss.png')\n",
        "        imageio.mimsave('progress.gif', [np.array(i) for i in images])\n",
        "    print('Epoch {}: G_loss: {:.4f} C_loss: {:.4f}'.format(epoch, g_error, c_error))\n",
        "    \n",
        "print('Training Finished')\n",
        "torch.save(generator.state_dict(), 'mnist_generator.pth')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}