W_GANs_with_GPs.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyO/zKHolzxS+9fPlnzika43",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/raghulchandramouli/GAN_series/blob/main/W_GANs_with_GPs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "TRAINING OF DCGANs NETWORK WITH WGANS LOSS"
      ],
      "metadata": {
        "id": "tU3kQw17fe6n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "import torchvision.datasets as datasets\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader\n",
        "from tqdm import tqdm\n",
        "from torch.utils.tensorboard import SummaryWriter"
      ],
      "metadata": {
        "id": "HqfGjHwUfmE1"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "BUILDING DISCRIMINATOR"
      ],
      "metadata": {
        "id": "_AUOiIWzg04U"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Discriminator(nn.Module):\n",
        "  def __init__(self, channels_img, features_d):\n",
        "    super(Discriminator, self).__init__()\n",
        "    self.disc = nn.Sequential(\n",
        "        nn.Conv2d(\n",
        "            channels_img, features_d, kernel_size=4, stride=2, padding=1\n",
        "        ),\n",
        "        nn.LeakyReLU(0.2),\n",
        "        ### sequnce of Conv_Layers connected togather at the end:\n",
        "        self._block(features_d, features_d*2, 4, 2, 1),\n",
        "        self._block(features_d*2, features_d*4, 4, 2, 1),\n",
        "        self._block(features_d*4, features_d*8, 4, 2, 1),\n",
        "\n",
        "        ### After Spatial Reduction in latent vectors:\n",
        "        nn.Conv2d(features_d*8, 1, kernel_size=4, stride=2, padding=0)\n",
        "    )\n",
        "\n",
        "  def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "    return nn.Sequential(\n",
        "        nn.Conv2d(\n",
        "            in_channels,\n",
        "            out_channels,\n",
        "            kernel_size,\n",
        "            stride,\n",
        "            padding,\n",
        "            bias = False\n",
        "        ),\n",
        "        nn.BatchNorm2d(out_channels),\n",
        "        nn.LeakyReLU(0.2),\n",
        "    )\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.disc(x)"
      ],
      "metadata": {
        "id": "kSkmLXKNgIoA"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "BUILDING GENERATORS"
      ],
      "metadata": {
        "id": "9kKs2tKIjVuz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Generator(nn.Module):\n",
        "  def __init__(self, channels_noise, channels_img, features_g):\n",
        "    super(Generator, self).__init__()\n",
        "    self.net = nn.Sequential(\n",
        "        ### INPUTs:\n",
        "        self._block(channels_noise, features_g*16, 4, 1, 0),\n",
        "        self._block(features_g*16, features_g*8, 4, 2, 1),\n",
        "        self._block(features_g*8, features_g*4, 4, 2, 1),\n",
        "        self._block(features_g*4, features_g*2, 4, 2, 1),\n",
        "        nn.ConvTranspose2d(\n",
        "            features_g*2, channels_img, kernel_size=4, stride=2, padding=1\n",
        "        ),\n",
        "        ### OUTPUTs:\n",
        "        nn.Tanh(),\n",
        "\n",
        "    )\n",
        "\n",
        "  def _block(self, in_channels, out_channels, kernel_size, stride, padding):\n",
        "    return nn.Sequential(\n",
        "        nn.ConvTranspose2d(\n",
        "            in_channels,\n",
        "            out_channels,\n",
        "            kernel_size,\n",
        "            stride,\n",
        "            padding,\n",
        "            bias=False\n",
        "        ),\n",
        "        nn.BatchNorm2d(out_channels),\n",
        "        nn.ReLU(),\n",
        "    )\n",
        "\n",
        "  def forward(self, x):\n",
        "    return self.net(x)"
      ],
      "metadata": {
        "id": "XQxDcHEkjMVA"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def initialize_weights(model):\n",
        "\n",
        "    for m in model.modules():\n",
        "        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):\n",
        "            nn.init.normal_(m.weight.data, 0.0, 0.02)"
      ],
      "metadata": {
        "id": "33ZtD7fZmkYE"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "HYPERPARAMETERS:"
      ],
      "metadata": {
        "id": "akdSEc17lqF3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "LEARNING_RATE = 5e-5\n",
        "BATCH_SIZE = 64\n",
        "IMAGE_SIZE = 64\n",
        "CHANNELS_IMG = 1\n",
        "Z_DIM = 128\n",
        "NUM_EPOCHS = 5\n",
        "FEATURES_CRITIC = 64\n",
        "FEATURES_GEN = 64\n",
        "CRITIC_ITERATIONS = 5\n",
        "WEIGHT_CLIP = 0.01"
      ],
      "metadata": {
        "id": "tAYlf4ZOlujO"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torchvision.transforms as T\n",
        "\n",
        "transformations = T.Compose([\n",
        "    T.Resize(IMAGE_SIZE),\n",
        "    T.ToTensor(),\n",
        "    T.Normalize([0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]),\n",
        "])\n",
        "\n",
        "dataset = datasets.MNIST(root=\"dataset/\", transform=transformations, download=True)\n",
        "# or for CelebA dataset\n",
        "# dataset = datasets.ImageFolder(root=\"celeb_dataset\", transform=transformations)\n",
        "\n",
        "loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)\n"
      ],
      "metadata": {
        "id": "ocJ0IJozlwN4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "INITIALIZING OPTIMIZERS FOR GEN/DISC CRITIC"
      ],
      "metadata": {
        "id": "R0Q2e5qMmUhG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gen = Generator(Z_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)\n",
        "critic = Discriminator(CHANNELS_IMG, FEATURES_CRITIC).to(device)\n",
        "initialize_weights(gen)\n",
        "initialize_weights(critic)\n",
        "\n",
        "# initializate optimizer\n",
        "opt_gen = optim.RMSprop(gen.parameters(), lr=LEARNING_RATE)\n",
        "opt_critic = optim.RMSprop(critic.parameters(), lr=LEARNING_RATE)\n",
        "\n",
        "# for tensorboard plotting\n",
        "fixed_noise = torch.randn(32, Z_DIM, 1, 1).to(device)\n",
        "writer_real = SummaryWriter(f\"logs/real\")\n",
        "writer_fake = SummaryWriter(f\"logs/fake\")\n",
        "step = 0\n"
      ],
      "metadata": {
        "id": "cYsSqJysmcg2"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gen.train()\n",
        "critic.train()"
      ],
      "metadata": {
        "id": "mHc6pp_Bme45"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "TRAINING:"
      ],
      "metadata": {
        "id": "26B9MlXlnJhE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(NUM_EPOCHS):\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    for batch_idx, (data, _) in enumerate(tqdm(loader)):\n",
        "        data = data.to(device)\n",
        "        cur_batch_size = data.shape[0]\n",
        "\n",
        "\n",
        "        for _ in range(CRITIC_ITERATIONS):\n",
        "            noise = torch.randn(cur_batch_size, Z_DIM, 1, 1).to(device)\n",
        "            fake = gen(noise)\n",
        "            critic_real = critic(data).reshape(-1)\n",
        "            critic_fake = critic(fake).reshape(-1)\n",
        "            loss_critic = -(torch.mean(critic_real) - torch.mean(critic_fake))\n",
        "            critic.zero_grad()\n",
        "            loss_critic.backward(retain_graph=True)\n",
        "            opt_critic.step()\n",
        "\n",
        "\n",
        "            for p in critic.parameters():\n",
        "                p.data.clamp_(-WEIGHT_CLIP, WEIGHT_CLIP)\n",
        "\n",
        "\n",
        "        gen_fake = critic(fake).reshape(-1)\n",
        "        loss_gen = -torch.mean(gen_fake)\n",
        "        gen.zero_grad()\n",
        "        loss_gen.backward()\n",
        "        opt_gen.step()\n",
        "\n",
        "\n",
        "        if batch_idx % 100 == 0 and batch_idx > 0:\n",
        "            gen.eval()\n",
        "            critic.eval()\n",
        "            print(\n",
        "                f\"Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_idx}/{len(loader)} \\\n",
        "                  Loss D: {loss_critic:.4f}, loss G: {loss_gen:.4f}\"\n",
        "            )\n",
        "\n",
        "            with torch.no_grad():\n",
        "                fake = gen(noise)\n",
        "\n",
        "                img_grid_real = torchvision.utils.make_grid(\n",
        "                    data[:32], normalize=True\n",
        "                )\n",
        "                img_grid_fake = torchvision.utils.make_grid(\n",
        "                    fake[:32], normalize=True\n",
        "                )\n",
        "\n",
        "                writer_real.add_image(\"Real\", img_grid_real, global_step=step)\n",
        "                writer_fake.add_image(\"Fake\", img_grid_fake, global_step=step)\n",
        "\n",
        "            step += 1\n",
        "            gen.train()\n",
        "            critic.train()"
      ],
      "metadata": {
        "id": "4J_scRuxm7Ca"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "TESTING:"
      ],
      "metadata": {
        "id": "RfgUk5XtKC4p"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def test():\n",
        "    N, in_channels, H, W = 8, 3, 64, 64\n",
        "    noise_dim = 100\n",
        "    x = torch.randn((N, in_channels, H, W))\n",
        "    disc = Discriminator(in_channels, 8)\n",
        "    assert disc(x).shape == (N, 1, 1, 1), \"Discriminator test failed\"\n",
        "    gen = Generator(noise_dim, in_channels, 8)\n",
        "    z = torch.randn((N, noise_dim, 1, 1))\n",
        "    assert gen(z).shape == (N, in_channels, H, W), \"Generator test failed\"\n",
        "    print(\"Success, tests passed!\")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    test()"
      ],
      "metadata": {
        "id": "a0eojc2BnK31"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}