GAN_DDP.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GAN_DDP.ipynb",
      "provenance": [],
      "mount_file_id": "1EtRinVecjrdZX5EBzRIQI4ENgphSz9cQ",
      "authorship_tag": "ABX9TyMwvNeORyRw1/JOqK4Ut08C",
      "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/abhistar/DDP_GAN/blob/main/GAN_DDP.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oDu-LpMW9Tvb"
      },
      "source": [
        "import os\n",
        "import random\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.nn.parallel\n",
        "import torch.backends.cudnn as cudnn\n",
        "import torch.optim as optim\n",
        "import torch.utils.data\n",
        "import torchvision.datasets as dset\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.utils as vutils\n",
        "import torchvision.models as models\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm\n",
        "import copy\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "__KFMaki9c0S",
        "outputId": "4c7b50a3-a748-4a18-e70c-de7cf4e6e0c3"
      },
      "source": [
        "manual_seed = 999\n",
        "print(\"Random Seed: \", manual_seed)\n",
        "random.seed(manual_seed)\n",
        "torch.manual_seed(manual_seed)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Random Seed:  999\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7fda23d62be8>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9AEplL2y9ek8",
        "outputId": "3819475c-e6f8-4884-ee81-f757d3287bb4"
      },
      "source": [
        "# location of folder\n",
        "dataroot = \"drive/My Drive/new_images\"\n",
        "# number of workers\n",
        "workers = 4\n",
        "# batch size for training\n",
        "batch_size = 30\n",
        "# image size for input\n",
        "image_size = 128\n",
        "# number of channel (1 for BW, 3 for RGB)\n",
        "nc = 3\n",
        "# Size of z latent vector (i.e. size of generator input)\n",
        "nz = 16\n",
        "# Size of feature maps in generator\n",
        "ngf = 128\n",
        "# Size of feature maps in discriminator\n",
        "ndf = 128\n",
        "# Number of training epochs\n",
        "num_epochs = 50\n",
        "# Learning rate for optimizers\n",
        "lr = 0.0005\n",
        "# Beta values hyperparam for Adam optimizers\n",
        "beta1 = 0.5\n",
        "beta2 = 0.99\n",
        "# Number of GPUs available, currently running on CPU\n",
        "ngpu = 1\n",
        "\n",
        "print(torch.cuda.get_device_name(torch.cuda.current_device()))"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Tesla T4\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GnieJNe5-Cki"
      },
      "source": [
        "# setup dataset and data loader\n",
        "dataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([transforms.Resize(image_size), \n",
        "                                                                        transforms.CenterCrop(image_size),\n",
        "                                                                        transforms.Grayscale(num_output_channels=3),\n",
        "                                                                        transforms.ToTensor(),\n",
        "                                                                        transforms.Normalize((0.5,),(0.5,))]))\n",
        "\n",
        "dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers)\n",
        "\n",
        "device = torch.device(\"cuda:0\" if(torch.cuda.is_available() and ngpu>0) else \"cpu\")"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C_0al6Ey9hQW"
      },
      "source": [
        "# define weights for layer and normalisation\n",
        "def weights_init(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find('Conv')!=-1:\n",
        "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
        "    elif classname.find('BatchNorm')!=-1:\n",
        "        nn.init.normal_(m.weight, 1.0, 0.02)\n",
        "        nn.init.constant_(m.bias.data, 0)"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u8fRKhh_tvQd",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a17996e9-9842-4ed1-ae7d-95ac82546eb4"
      },
      "source": [
        "class Generator(nn.Module):\n",
        "    def __init__(self, ngpu):\n",
        "        super(Generator, self).__init__()\n",
        "        self.ngpu = ngpu\n",
        "        self.main = nn.Sequential(\n",
        "            # input z into convolution\n",
        "            nn.ConvTranspose2d(nz, ngf*8, 8, 2, 0, bias=False),\n",
        "            nn.BatchNorm2d(ngf*8),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            # state ngf * 8 * 8\n",
        "            nn.ConvTranspose2d(ngf*8, ngf*4, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ngf*4),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            # state ngf/2 * 16*16\n",
        "            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ngf*2),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            # state ngf/4 * 32*32\n",
        "            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ngf),\n",
        "            nn.ReLU(True),\n",
        "            \n",
        "            # state ngf/8 *64*64\n",
        "            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),\n",
        "            nn.Tanh()\n",
        "            \n",
        "            #state 1 *128*128            \n",
        "        )\n",
        "        \n",
        "    def forward(self, input):\n",
        "        return self.main(input)\n",
        "\n",
        "\n",
        "# Create the generator\n",
        "netG = Generator(ngpu).to(device)\n",
        "\n",
        "# Handle multi-gpu if desired\n",
        "if (device.type == 'cuda') and (ngpu > 1):\n",
        "    netG = nn.DataParallel(netG, list(range(ngpu)))\n",
        "\n",
        "# Apply the weights_init function to randomly initialize all weights\n",
        "#  to mean=0, stdev=0.2.\n",
        "netG.apply(weights_init)\n",
        "\n",
        "# Print the model\n",
        "print(netG)"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Generator(\n",
            "  (main): Sequential(\n",
            "    (0): ConvTranspose2d(16, 1024, kernel_size=(8, 8), stride=(2, 2), bias=False)\n",
            "    (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (2): ReLU(inplace=True)\n",
            "    (3): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (5): ReLU(inplace=True)\n",
            "    (6): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (7): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (8): ReLU(inplace=True)\n",
            "    (9): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (10): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (11): ReLU(inplace=True)\n",
            "    (12): ConvTranspose2d(128, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (13): Tanh()\n",
            "  )\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n8-ouxTYtzE5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4862586d-c496-45c2-e20c-868d4a709644"
      },
      "source": [
        "class Discriminator(nn.Module):\n",
        "    def __init__(self, ngpu):\n",
        "        super(Discriminator, self).__init__()\n",
        "        self.ngpu = ngpu\n",
        "        self.main = nn.Sequential(\n",
        "            # input is (nc) *128*128\n",
        "            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            \n",
        "            # input is ndf/8 *64*64\n",
        "            nn.Conv2d(ndf, ndf*2, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ndf*2),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            \n",
        "            #input is ndf/4 *32*32\n",
        "            nn.Conv2d(ndf*2, ndf*4, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ndf*4),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            \n",
        "            #input is ndf/2 *16*16\n",
        "            nn.Conv2d(ndf*4, ndf*8, 4, 2, 1, bias=False),\n",
        "            nn.BatchNorm2d(ndf*8),\n",
        "            nn.LeakyReLU(0.2, inplace=True),\n",
        "            \n",
        "            #input is ndf *8*8\n",
        "            nn.Conv2d(ndf*8, 1, 8, 1, 0, bias=False),\n",
        "            nn.Sigmoid()\n",
        "        )\n",
        "    def forward(self, input):\n",
        "        return self.main(input)\n",
        "\n",
        "\n",
        "# Create the Discriminator\n",
        "netD = Discriminator(ngpu).to(device)\n",
        "\n",
        "# Handle multi-gpu if desired\n",
        "if (device.type == 'cuda') and (ngpu > 1):\n",
        "    netD = nn.DataParallel(netD, list(range(ngpu)))\n",
        "\n",
        "# Apply the weights_init function to randomly initialize all weights\n",
        "#  to mean=0, stdev=0.2.\n",
        "netD.apply(weights_init)\n",
        "\n",
        "# Print the model\n",
        "print(netD)"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Discriminator(\n",
            "  (main): Sequential(\n",
            "    (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (1): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (2): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (4): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (5): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (6): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (7): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (8): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
            "    (9): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
            "    (10): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (11): Conv2d(1024, 1, kernel_size=(8, 8), stride=(1, 1), bias=False)\n",
            "    (12): Sigmoid()\n",
            "  )\n",
            ")\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dp1U8fDi-bSD"
      },
      "source": [
        "# Initialize BCELoss function\n",
        "criterion = nn.BCELoss()\n",
        "\n",
        "# Create batch of latent vectors that we will use to visualize\n",
        "#  the progression of the generator\n",
        "fixed_noise = torch.randn(128, nz, 1, 1, device=device)\n",
        "\n",
        "# Establish convention for real and fake labels during training\n",
        "real_label = 1.\n",
        "fake_label = 0.\n",
        "\n",
        "# Setup Adam optimizers for both G and D\n",
        "optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, beta2))\n",
        "optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, beta2))"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Vt8VzxHw-dfT",
        "outputId": "9d1447dd-c8c1-4bcc-b8c3-8a2a0b0b9024"
      },
      "source": [
        "# Training Loop\n",
        "\n",
        "# Lists to keep track of progress\n",
        "img_list = []\n",
        "G_losses = []\n",
        "D_losses = []\n",
        "iters = 0\n",
        "alpha = 10\n",
        "\n",
        "print(\"Starting Training Loop...\")\n",
        "# For each epoch\n",
        "for epoch in tqdm(range(num_epochs)):\n",
        "    # For each batch in the dataloader\n",
        "    for i, data in enumerate(dataloader):\n",
        "\n",
        "        ############################\n",
        "        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
        "        ###########################\n",
        "        ## Train with all-real batch\n",
        "        netD.zero_grad()\n",
        "        # Format batch\n",
        "        real_cpu = data[0].to(device)\n",
        "        b_size = real_cpu.size(0)\n",
        "        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)\n",
        "        # Forward pass real batch through D\n",
        "        output = netD(real_cpu).view(-1)\n",
        "        # Calculate loss on all-real batch\n",
        "        errD_real = criterion(output, label)\n",
        "        # Calculate gradients for D in backward pass\n",
        "        errD_real.backward()\n",
        "        D_x = output.mean().item()\n",
        "\n",
        "        ## Train with all-fake batch\n",
        "        # Generate batch of latent vectors\n",
        "        noise = torch.randn(b_size, nz, 1, 1, device=device)\n",
        "        # Generate fake image batch with G\n",
        "        fake = netG(noise)\n",
        "        label.fill_(fake_label)\n",
        "        # Classify all fake batch with D\n",
        "        output = netD(fake.detach()).view(-1)\n",
        "        # Calculate D's loss on the all-fake batch\n",
        "        errD_fake = criterion(output, label)\n",
        "        # Calculate the gradients for this batch\n",
        "        errD_fake.backward()\n",
        "        D_G_z1 = output.mean().item()\n",
        "\n",
        "\n",
        "        # Add the gradients from the all-real and all-fake batches\n",
        "        errD = errD_real + errD_fake\n",
        "\n",
        "        # Update D\n",
        "        optimizerD.step()\n",
        "\n",
        "        ############################\n",
        "        # (2) Update G network: maximize log(D(G(z)))\n",
        "        ###########################\n",
        "        netG.zero_grad()\n",
        "        label.fill_(real_label)  # fake labels are real for generator cost\n",
        "        # Since we just updated D, perform another forward pass of all-fake batch through D\n",
        "        output = netD(fake).view(-1)\n",
        "        # Calculate G's loss based on this output\n",
        "        errG = criterion(output, label)\n",
        "        # Calculate gradients for G\n",
        "        errG.backward()\n",
        "        D_G_z2 = output.mean().item()\n",
        "        # Update G\n",
        "        optimizerG.step()\n",
        "\n",
        "        # Output training stats\n",
        "        if i % 60 == 0:\n",
        "            print('[%d/%d][%d/%d]\\tLoss_D: %.4f\\tLoss_G: %.4f\\tD(x): %.4f\\tD(G(z)): %.4f / %.4f'\n",
        "                  % (epoch, num_epochs, i, len(dataloader),\n",
        "                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
        "\n",
        "        # Save Losses for plotting later\n",
        "        G_losses.append(errG.item())\n",
        "        D_losses.append(errD.item())\n",
        "\n",
        "        # Check how the generator is doing by saving G's output on fixed_noise\n",
        "        if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):\n",
        "            with torch.no_grad():\n",
        "                fake = netG(fixed_noise).detach().cpu()\n",
        "            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "  0%|          | 0/50 [00:00<?, ?it/s]\u001b[A"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Starting Training Loop...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xoa7_mgU-fNQ"
      },
      "source": [
        "plt.figure(figsize=(10,5))\n",
        "plt.title(\"Generator and Discriminator Loss During Training\")\n",
        "plt.plot(G_losses,label=\"G\")\n",
        "plt.plot(D_losses,label=\"D\")\n",
        "plt.xlabel(\"iterations\")\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "52EPOu-6-izn"
      },
      "source": [
        "# Grab a batch of real images from the dataloader\n",
        "real_batch = next(iter(dataloader))\n",
        "\n",
        "# Plot the real images\n",
        "plt.figure(figsize=(15,15))\n",
        "plt.subplot(1,2,1)\n",
        "plt.axis(\"off\")\n",
        "plt.title(\"Real Images\")\n",
        "plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(),(1,2,0)))\n",
        "\n",
        "# Plot the fake images from the last epoch\n",
        "plt.subplot(1,2,2)\n",
        "plt.axis(\"off\")\n",
        "plt.title(\"Fake Images\")\n",
        "plt.imshow(np.transpose(img_list[-1],(1,2,0)))\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JldUdLEE-k-S"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}