BiCycleGAN_Toward_Multimodal_Image-to-Image_Translation.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## BiCycleGAN - Paper - Toward Multimodal Image-to-Image Translation - PyTorch Implementation from Scratch\n",
        "\n",
        "# [Link to my Youtube Video Explaining this whole Notebook](https://www.youtube.com/watch?v=iCXruj3slIk&list=PLxqBkZuBynVRyOJs4RWmB_fKlOVe5S8CR&index=17)\n",
        "\n",
        "[![Imgur](https://imgur.com/n7xTxVm.png)](https://www.youtube.com/watch?v=iCXruj3slIk&list=PLxqBkZuBynVRyOJs4RWmB_fKlOVe5S8CR&index=17)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 👉 The Maps - Dataset link - http://efrosgans.eecs.berkeley.edu/pix2pix/datasets/\n",
        "\n",
        "## About the Satellite to Map Image Translation Dataset\n",
        "\n",
        "This is a dataset comprised of satellite images of New York and their corresponding Google maps pages. The image translation problem involves converting satellite photos to Google maps format, or the reverse, Google maps images to Satellite photos.\n",
        "\n",
        "The dataset is provided on the pix2pix website and can be downloaded as a 255-megabyte zip file.\n",
        "\n",
        "maps\n",
        "├── train\n",
        "└── val\n",
        "\n",
        "The train folder contains 1,097 images, whereas the validation dataset contains 1,099 images.\n",
        "\n",
        "Images have a digit filename and are in JPEG format. Each image is 1,200 pixels wide and 600 pixels tall and contains both the satellite image on the left and the Google maps image on the right."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jjIvEuipGqxu",
        "outputId": "fa2c5913-6374-4a01-d292-e4c85d1ca3e9"
      },
      "outputs": [],
      "source": [
        "!pip install icecream\n",
        "import os\n",
        "import numpy as np\n",
        "import math\n",
        "import itertools\n",
        "import scipy\n",
        "import sys\n",
        "import time\n",
        "import datetime\n",
        "\n",
        "import torchvision.transforms as transforms\n",
        "from torchvision.utils import save_image\n",
        "\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets\n",
        "from torch.autograd import Variable\n",
        "import torch.autograd as autograd\n",
        "from torchvision.utils import make_grid\n",
        "\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.pyplot import figure\n",
        "from IPython.display import clear_output\n",
        "from icecream import ic"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UokJxbOYHuAi",
        "outputId": "d98cf4ea-3cba-411d-a329-e89ef1f48bf0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "metadata": {
        "id": "rVBTY_4wwfbq"
      },
      "outputs": [],
      "source": [
        "class Hyperparameters(object):\n",
        "  def __init__(self, **kwargs):\n",
        "    self.__dict__.update(kwargs)\n",
        "\n",
        "hp = Hyperparameters(\n",
        "    epoch=0,\n",
        "    n_epochs=200,\n",
        "    batch_size=8,        \n",
        "    dataset_train_mode=\"train\",\n",
        "    dataset_test_mode=\"val\",    \n",
        "    lr=.0002,    \n",
        "    b1=.5,\n",
        "    b2=0.999,\n",
        "    n_cpu=8,\n",
        "    img_size=128,\n",
        "    channels=3,\n",
        "    latent_dim=8,\n",
        "    n_critic=5,\n",
        "    sample_interval=400,\n",
        "    lambda_pixel=10,\n",
        "    lambda_latent=.5,\n",
        "    lambda_kl=.01)\n",
        "\n",
        "img_root_folder = '/content/drive/MyDrive/All_Datasets/Maps-UCBerkeley-CycleGAN/maps'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 53,
      "metadata": {
        "id": "XKrlGFNzgjyr"
      },
      "outputs": [],
      "source": [
        "# Just print to check that the full file paths of the images are printed indeed\n",
        "# sorted(glob.glob(os.path.join(img_root_folder, 'train') + \"/*.*\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 54,
      "metadata": {
        "cellView": "form",
        "id": "yZaSsdJa_xqH"
      },
      "outputs": [],
      "source": [
        "import glob\n",
        "import random\n",
        "import os\n",
        "from torch.utils.data import Dataset\n",
        "from PIL import Image\n",
        "\n",
        "class ImageDataset(Dataset):\n",
        "    def __init__(self, root, transforms_=None, mode=\"train\"):\n",
        "        self.transform = transforms.Compose(transforms_)\n",
        "\n",
        "        # Using the glob and sorted functions to load all the images and sort them.\n",
        "        self.files = sorted(glob.glob(os.path.join(root, mode) + \"/*.*\"))\n",
        "        if mode == \"train\":\n",
        "            self.files.extend(sorted(glob.glob(os.path.join(root, \"test\") + \"/*.*\")))\n",
        "            # `extends()` method adds the specified list elements to the end of the current list.\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "\n",
        "        img = Image.open(self.files[index % len(self.files)])\n",
        "        ''' In above line, my target is to find index item in a list based on the length of a variable in the list\n",
        "        So, in case such indexes do not exist, None is returned. \n",
        "        a % b => a is divided by b, and the remainder of that division is returned.\n",
        "        5 % 100 => 5 '''\n",
        "        w, h = img.size # Pillow.Image.open() returns width, height\n",
        "        img_A = img.crop((0, 0, w / 2, h)) # (left, upper, right, lower)\n",
        "        img_B = img.crop((w / 2, 0, w, h)) # (left, upper, right, lower)        \n",
        "\n",
        "        ''' Below code block implements Horizontal Flipping or Mirroring Image based on randomly generated probability '''\n",
        "        if np.random.random() < 0.5:\n",
        "            # Converting from numpy arrays to a RGB image\n",
        "            # Call PIL.Image.fromarray(obj, mode) with obj as a 3-D array and mode as \"RGB\" to convert obj into an image.\n",
        "            # Image.fromarray creates an image memory from an object exporting the array interface\n",
        "            img_A = Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\") # Mirror in x direction (flip horizontally)\n",
        "            img_B = Image.fromarray(np.array(img_B)[:, ::-1, :], \"RGB\") # Mirror in x direction (flip horizontally)\n",
        "            ''' a[::-1]    # all items in the array, reversed. \n",
        "            Pillow image returns tuple of (width, height)    \n",
        "            The data has 3 dimensions: height, width and color. Numpy shape of the image is a tuple of (row (height), column (width), color(3) )\n",
        "            So ::-1 effectively reverses the order of the width. The height and color are not affected.\n",
        "            '''\n",
        "\n",
        "        img_A = self.transform(img_A)\n",
        "        img_B = self.transform(img_B)\n",
        "\n",
        "        return {\"A\": img_A, \"B\": img_B}\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.files)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Explanation of Horizontal Image Flipping with `Image.fromarray(np.array(img_A)[:, ::-1, :], \"RGB\")`\n",
        "\n",
        "Basically, its implementing, the official code's mechanism\n",
        "\n",
        "https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/data/base_dataset.py#L144\n",
        "\n",
        "```py\n",
        "def __flip(img, flip):\n",
        "    if flip:\n",
        "        return img.transpose(Image.FLIP_LEFT_RIGHT)\n",
        "    return \n",
        "\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Note - numpy arrays and PIL images have different shape, in Numpy its (H,W) and in PIL and (W,H)\n",
        "\n",
        "===========================================================================\n",
        "\n",
        "#### Why `img = Image.open(self.files[index % len(self.files)])`\n",
        "\n",
        "Because, here my target is to find index item in a list based on the length of a variable in the list\n",
        "\n",
        "So, in case such indexes do not exist, None is returned. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 56,
      "metadata": {
        "cellView": "form",
        "id": "VWcpYFFnWAbv"
      },
      "outputs": [],
      "source": [
        "#  IMAGE VISUALIZER HELPERS\n",
        "def imshow(img,size=10):\n",
        "  img = img / 2 + 0.5  # de-normalizing\n",
        "  npimg = img.numpy()\n",
        "  plt.figure(figsize=(size, size))\n",
        "  plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
        "  plt.show()\n",
        "\n",
        "\n",
        "import matplotlib.image as mpimg\n",
        "\n",
        "def visualise_output(path, x, y):\n",
        "    img = mpimg.imread(path)\n",
        "    plt.figure(figsize=(x,y))\n",
        "    plt.imshow(img)  \n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Explanation of line `img = img / 2 + 0.5`  # de-normalizing\n",
        "\n",
        "Since the normalization process is actually z = (x - mean) / sigma \n",
        "\n",
        "Where both mean and sigma is 0.5\n",
        "\n",
        "The inverse normalization should be x = z*sigma + mean\n",
        "\n",
        "\n",
        "https://discuss.pytorch.org/t/simple-way-to-inverse-transform-normalization/4821/7"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 57,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pR5LjEFQARFm",
        "outputId": "3bedecbb-57bf-43ec-ad89-22337d77f320"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torchvision/transforms/transforms.py:288: UserWarning: Argument interpolation should be of type InterpolationMode instead of int. Please, use InterpolationMode enum.\n",
            "  \"Argument interpolation should be of type InterpolationMode instead of int. \"\n"
          ]
        }
      ],
      "source": [
        "#  CONFIGURE DATALOADERS\n",
        "transforms_ = [\n",
        "    transforms.Resize((hp.img_size, hp.img_size), Image.BICUBIC),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
        "]\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    ImageDataset(img_root_folder, mode=hp.dataset_train_mode, transforms_=transforms_),\n",
        "    batch_size=hp.batch_size,\n",
        "    shuffle=True,\n",
        "    num_workers=1,\n",
        ")\n",
        "val_dataloader = DataLoader(\n",
        "    ImageDataset(img_root_folder, mode=hp.dataset_test_mode, transforms_=transforms_),\n",
        "    batch_size=16,\n",
        "    shuffle=True,\n",
        "    num_workers=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "7vp2cbB5okgD",
        "outputId": "1afdc064-3be4-4174-f895-503a44657601"
      },
      "outputs": [],
      "source": [
        "# VISUALING SAMPLE DATA { run: \"auto\" }\n",
        "pic_size = 16\n",
        "\n",
        "dataiter = iter(train_dataloader)\n",
        "images = dataiter.next()\n",
        "\n",
        "for i in range(len(images[\"A\"])):\n",
        "  imshow(make_grid([images[\"A\"][i],images[\"B\"][i]]), size=pic_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Unet\n",
        "\n",
        "Unet is an end-to-end fully convolutional network (FCN), i.e. it only contains Convolutional layers and does not contain any Dense layer because of which it can accept image of any size.\n",
        "\n",
        "The left hand side is the contraction path (Encoder) where we apply regular convolutions and max pooling layers.\n",
        "\n",
        "In the Encoder, the size of the image gradually reduces while the depth gradually increases. (e.g. Starting from 128x128x3 to 8x8x256 )\n",
        "\n",
        "\n",
        "The right hand side is the expansion path (Decoder) where we apply transposed convolutions along with regular convolutions\n",
        "\n",
        "In the decoder, the size of the image gradually increases and the depth gradually decreases."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 60,
      "metadata": {
        "id": "EEb5TdBmIy7l"
      },
      "outputs": [],
      "source": [
        "from torchvision.models import resnet18\n",
        "\n",
        "########################################################\n",
        "# Initialize convolution layer weights to N(0,0.02)\n",
        "########################################################\n",
        "def weights_init_normal(m):\n",
        "    classname = m.__class__.__name__\n",
        "    if classname.find(\"Conv\") != -1:\n",
        "        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
        "    elif classname.find(\"BatchNorm2d\") != -1:\n",
        "        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
        "        torch.nn.init.constant_(m.bias.data, 0.0)\n",
        "\n",
        "\n",
        "##############################\n",
        "#           U-NET\n",
        "##############################\n",
        "'''  As per the Paper - \"For generator G, we use the U-Net, which contains an encoder-decoder\n",
        "architecture, with symmetric skip connections.\" \n",
        "In the Encoder, the size of the image gradually reduces while the depth gradually increases. \n",
        "'''\n",
        "\n",
        "class UNetDown(nn.Module):\n",
        "    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):\n",
        "        super(UNetDown, self).__init__()\n",
        "        layers = [nn.Conv2d(in_size, out_size, 3, stride=2, padding=1, bias=False)]\n",
        "        if normalize:\n",
        "            layers.append(nn.BatchNorm2d(out_size, 0.8))\n",
        "        layers.append(nn.LeakyReLU(0.2))\n",
        "        self.model = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)\n",
        "\n",
        "''' In the decoder, the size of the image gradually increases and the depth gradually decreases. '''\n",
        "class UNetUp(nn.Module):\n",
        "    def __init__(self, in_size, out_size):\n",
        "        super(UNetUp, self).__init__()\n",
        "        self.model = nn.Sequential(\n",
        "            nn.Upsample(scale_factor=2), # Upsampling by a scaling factor of 2\n",
        "            nn.Conv2d(in_size, out_size, 3, stride=1, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(out_size, 0.8),\n",
        "            nn.ReLU(inplace=True),\n",
        "        )\n",
        "\n",
        "    def forward(self, x, skip_input):\n",
        "        x = self.model(x)\n",
        "        x = torch.cat((x, skip_input), 1)\n",
        "        return x\n",
        "\n",
        "\n",
        "class Generator(nn.Module):\n",
        "    def __init__(self, latent_dim, img_shape):\n",
        "        super(Generator, self).__init__()\n",
        "        channels, self.h, self.w = img_shape\n",
        "\n",
        "        self.fc = nn.Linear(latent_dim, self.h * self.w)\n",
        "\n",
        "        self.down1 = UNetDown(channels + 1, 64, normalize=False)\n",
        "        self.down2 = UNetDown(64, 128)\n",
        "        self.down3 = UNetDown(128, 256)\n",
        "        self.down4 = UNetDown(256, 512)\n",
        "        self.down5 = UNetDown(512, 512)\n",
        "        self.down6 = UNetDown(512, 512)\n",
        "        self.down7 = UNetDown(512, 512, normalize=False)\n",
        "        self.up1 = UNetUp(512, 512)\n",
        "        self.up2 = UNetUp(1024, 512)\n",
        "        self.up3 = UNetUp(1024, 512)\n",
        "        self.up4 = UNetUp(1024, 256)\n",
        "        self.up5 = UNetUp(512, 128)\n",
        "        self.up6 = UNetUp(256, 64)\n",
        "\n",
        "        self.final = nn.Sequential(\n",
        "            nn.Upsample(scale_factor=2), \n",
        "            nn.Conv2d(128, channels, 3, stride=1, padding=1), nn.Tanh()\n",
        "        )\n",
        "\n",
        "    def forward(self, x, z):\n",
        "        # Propagate noise through fc layer and reshape to img shape\n",
        "        z = self.fc(z).view(z.size(0), 1, self.h, self.w)\n",
        "        d1 = self.down1(torch.cat((x, z), 1))\n",
        "        d2 = self.down2(d1)\n",
        "        d3 = self.down3(d2)\n",
        "        d4 = self.down4(d3)\n",
        "        d5 = self.down5(d4)\n",
        "        d6 = self.down6(d5)\n",
        "        d7 = self.down7(d6)\n",
        "        u1 = self.up1(d7, d6)\n",
        "        u2 = self.up2(u1, d5)\n",
        "        u3 = self.up3(u2, d4)\n",
        "        u4 = self.up4(u3, d3)\n",
        "        u5 = self.up5(u4, d2)\n",
        "        u6 = self.up6(u5, d1)\n",
        "\n",
        "        return self.final(u6)\n",
        "\n",
        "\n",
        "##############################\n",
        "#        Encoder\n",
        "##############################\n",
        "\n",
        "class Encoder(nn.Module):\n",
        "    # 1. Use this encoder and get mu and log_var\n",
        "    # 2. std = exp(log_var / 2)\n",
        "    # 3. random_z = N(0, 1)\n",
        "    # 4. encoded_z = random_z * std + mu (Reparameterization trick)\n",
        "    def __init__(self, latent_dim, input_shape):\n",
        "        super(Encoder, self).__init__()\n",
        "        resnet18_model = resnet18(pretrained=False)\n",
        "        self.feature_extractor = nn.Sequential(*list(resnet18_model.children())[:-3])\n",
        "        # [:-3] => Everything except the last 3 items\n",
        "        self.pooling = nn.AvgPool2d(kernel_size=8, stride=8, padding=0)\n",
        "        # Output is mu and log(var) for reparameterization trick used in VAEs\n",
        "        # mu and logvar assigned the same value (the encoder’s last layer output)\n",
        "        self.fc_mu = nn.Linear(256, latent_dim)\n",
        "        self.fc_logvar = nn.Linear(256, latent_dim)\n",
        "\n",
        "    def forward(self, img):\n",
        "        out = self.feature_extractor(img)\n",
        "        out = self.pooling(out)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        mu = self.fc_mu(out)\n",
        "        logvar = self.fc_logvar(out)\n",
        "        return mu, logvar\n",
        "\n",
        "\n",
        "##############################\n",
        "#        Discriminator\n",
        "##############################\n",
        "\n",
        "class MultiDiscriminator(nn.Module):\n",
        "    def __init__(self, input_shape):\n",
        "        super(MultiDiscriminator, self).__init__()\n",
        "\n",
        "        def discriminator_block(in_filters, out_filters, normalize=True):\n",
        "            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n",
        "            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n",
        "            if normalize:\n",
        "                layers.append(nn.BatchNorm2d(out_filters, 0.8))\n",
        "            layers.append(nn.LeakyReLU(0.2))\n",
        "            return layers\n",
        "\n",
        "        channels, _, _ = input_shape\n",
        "        # Extracts discriminator models\n",
        "        self.models = nn.ModuleList()\n",
        "        for i in range(3):\n",
        "            self.models.add_module(\n",
        "                \"disc_%d\" % i,\n",
        "                nn.Sequential(\n",
        "                    *discriminator_block(channels, 64, normalize=False),\n",
        "                    *discriminator_block(64, 128),\n",
        "                    *discriminator_block(128, 256),\n",
        "                    *discriminator_block(256, 512),\n",
        "                    nn.Conv2d(512, 1, 3, padding=1)\n",
        "                ),\n",
        "            )\n",
        "\n",
        "        self.downsample = nn.AvgPool2d(channels, stride=2, padding=[1, 1], count_include_pad=False)\n",
        "\n",
        "    def compute_loss(self, x, ground_truth):\n",
        "        \"\"\"Computes the MSE between model output and scalar ground_truth\"\"\"\n",
        "        loss = sum([torch.mean((out - ground_truth) ** 2) for out in self.forward(x)])\n",
        "        return loss\n",
        "\n",
        "    def forward(self, x):\n",
        "        outputs = []\n",
        "        for m in self.models:\n",
        "            outputs.append(m(x))\n",
        "            x = self.downsample(x)\n",
        "        return outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 61,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YxFuX3PCKOVW",
        "outputId": "c0a5bccd-572c-444d-d805-226d29a71b99"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using CUDA\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "MultiDiscriminator(\n",
              "  (models): ModuleList(\n",
              "    (disc_0): Sequential(\n",
              "      (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (1): LeakyReLU(negative_slope=0.2)\n",
              "      (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (3): BatchNorm2d(128, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (4): LeakyReLU(negative_slope=0.2)\n",
              "      (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (6): BatchNorm2d(256, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (7): LeakyReLU(negative_slope=0.2)\n",
              "      (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (9): BatchNorm2d(512, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (10): LeakyReLU(negative_slope=0.2)\n",
              "      (11): Conv2d(512, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    )\n",
              "    (disc_1): Sequential(\n",
              "      (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (1): LeakyReLU(negative_slope=0.2)\n",
              "      (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (3): BatchNorm2d(128, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (4): LeakyReLU(negative_slope=0.2)\n",
              "      (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (6): BatchNorm2d(256, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (7): LeakyReLU(negative_slope=0.2)\n",
              "      (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (9): BatchNorm2d(512, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (10): LeakyReLU(negative_slope=0.2)\n",
              "      (11): Conv2d(512, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    )\n",
              "    (disc_2): Sequential(\n",
              "      (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (1): LeakyReLU(negative_slope=0.2)\n",
              "      (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (3): BatchNorm2d(128, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (4): LeakyReLU(negative_slope=0.2)\n",
              "      (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (6): BatchNorm2d(256, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (7): LeakyReLU(negative_slope=0.2)\n",
              "      (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
              "      (9): BatchNorm2d(512, eps=0.8, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (10): LeakyReLU(negative_slope=0.2)\n",
              "      (11): Conv2d(512, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    )\n",
              "  )\n",
              "  (downsample): AvgPool2d(kernel_size=3, stride=2, padding=[1, 1])\n",
              ")"
            ]
          },
          "execution_count": 61,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "##############################################\n",
        "#  SETUP, LOSS, INITIALIZE MODELS and BUFFERS\n",
        "##############################################\n",
        "cuda = True if torch.cuda.is_available() else False\n",
        "print(\"Using CUDA\" if cuda else \"Not using CUDA\")\n",
        "\n",
        "# Loss functions\n",
        "mae_loss = torch.nn.L1Loss()\n",
        "input_shape = (hp.channels, hp.img_size, hp.img_size)\n",
        "\n",
        "# Initialize generator, encoder and discriminators\n",
        "generator = Generator(hp.latent_dim, input_shape)\n",
        "encoder = Encoder(hp.latent_dim, input_shape)\n",
        "\n",
        "D_VAE = MultiDiscriminator(input_shape)\n",
        "D_LR = MultiDiscriminator(input_shape)\n",
        "\n",
        "if cuda:\n",
        "    generator = generator.cuda()\n",
        "    encoder.cuda()\n",
        "    D_VAE = D_VAE.cuda()\n",
        "    D_LR = D_LR.cuda()\n",
        "    mae_loss.cuda()\n",
        "\n",
        "# Initialize weights\n",
        "generator.apply(weights_init_normal)\n",
        "D_VAE.apply(weights_init_normal)\n",
        "D_LR.apply(weights_init_normal)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {
        "id": "TpEIFH7BIRwA"
      },
      "outputs": [],
      "source": [
        "#  SAMPLING IMAGES\n",
        "def sample_images(batches_done):\n",
        "    \"\"\"From the validation set this method will create images and \n",
        "    save those Generated samples in a path  \"\"\"\n",
        "    generator.eval()\n",
        "    imgs = next(iter(val_dataloader))\n",
        "    # next() will supply each subsequent element from the iterable\n",
        "    # So in this case each subsequent set of images from val_dataloader\n",
        "    img_samples = None\n",
        "    # For below line to work, I need to create a folder named 'maps' in the root_path\n",
        "    path = \"/content/%s/%s.png\" % ('maps', batches_done)\n",
        "    for img_A, img_B in zip(imgs[\"A\"], imgs[\"B\"]):\n",
        "        # Repeat input image by number of desired columns\n",
        "        real_A = img_A.view(1, *img_A.shape).repeat(hp.latent_dim, 1, 1, 1)\n",
        "        real_A = Variable(real_A.type(Tensor))\n",
        "        # Sample latent representations\n",
        "        sampled_z = Variable(Tensor(np.random.normal(0, 1, (hp.latent_dim, hp.latent_dim))))\n",
        "        # Generate samples\n",
        "        fake_B = generator(real_A, sampled_z)\n",
        "        # Concatenate samples horizontally\n",
        "        fake_B = torch.cat([x for x in fake_B.data.cpu()], -1)\n",
        "        img_sample = torch.cat((img_A, fake_B), -1)\n",
        "        img_sample = img_sample.view(1, *img_sample.shape)\n",
        "        # Concatenate with previous samples vertically\n",
        "        img_samples = img_sample if img_samples is None else torch.cat((img_samples, img_sample), -2)\n",
        "    save_image(img_samples, path, nrow=8, normalize=True)\n",
        "    generator.train() \n",
        "    return path\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "id": "fgRoX2DbHSwk"
      },
      "outputs": [],
      "source": [
        "#  OPTIMIZERS\n",
        "optimizer_E = torch.optim.Adam(encoder.parameters(), lr=hp.lr, betas=(hp.b1, hp.b2))\n",
        "optimizer_G = torch.optim.Adam(generator.parameters(), lr=hp.lr, betas=(hp.b1, hp.b2))\n",
        "\n",
        "optimizer_D_VAE = torch.optim.Adam(D_VAE.parameters(), lr=hp.lr, betas=(hp.b1, hp.b2))\n",
        "optimizer_D_LR = torch.optim.Adam(D_LR.parameters(), lr=hp.lr, betas=(hp.b1, hp.b2))\n",
        "\n",
        "Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The reparameterization trick\n",
        "\n",
        "#### The reparameterization trick is to learn two vectors σ and μ, \n",
        "\n",
        "First, sample ϵ from N(0,1) and then your latent vector Z would be (where ⊙ symbol or notation is the element-wise product.):\n",
        "\n",
        "## Z = μ +  ϵ ⊙ σ\n",
        "\n",
        "\n",
        "So, if an input data point is to be mapped into a latent variable `z` via sampling (after getting passed through a neural network), it has to follow the following equation:\n",
        "\n",
        "## z = z_mean + epsilon * std \n",
        "\n",
        "where \n",
        "\n",
        "### std = torch.exp(z_log_var / 2)\n",
        "\n",
        "----\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Reconciliation between 2 mathematical expression for z\n",
        "\n",
        "## 1st  -> z = z_mean + epsilon * std\n",
        "\n",
        "## 2nd (where the std is expressed as below )\n",
        "\n",
        "![Imgur](https://imgur.com/1G3iT3m.png)\n",
        "\n",
        "\n",
        "The σ in the first equation is the standard deviation which as you know is the square root of the variance. Then you can see that the multiplication of 0.5 outside the log equates to raising the variance inside of the log to the power of 0.5:\n",
        "\n",
        "\n",
        "![Imgur](https://imgur.com/CQxP66m.png)\n",
        "\n",
        "\n",
        "So they are the same."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "id": "45hmQrQZLoe4"
      },
      "outputs": [],
      "source": [
        "'''  REPARAMETERIZE \n",
        "\n",
        "The reparameterization trick is to learn two vectors σ and μ, sample ϵ from N(0,1) and \n",
        "then your latent vector Z would be as below (where ⊙ is the element-wise product.):\n",
        "\n",
        "z = μ +  ϵ ⊙  σ\n",
        "\n",
        "z = z_mean +  epsilon * sigma\n",
        "\n",
        "'''\n",
        "\n",
        "def reparameterization(z_mean, z_log_var):\n",
        "    std = torch.exp(z_log_var / 2)\n",
        "    sampled_z = Variable(Tensor(np.random.normal(0, 1, (z_mean.size(0), hp.latent_dim))))\n",
        "    z = z_mean + sampled_z * std\n",
        "    return z"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "So in above implementation the `sampled_z` is replacing 'epsilon'\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### `np.random.normal()`\n",
        "\n",
        "The `random.normal` method has the following syntax:\n",
        "\n",
        "`numpy.random.normal(m,s,n)`\n",
        "\n",
        "The random.normal function takes in three parameters:\n",
        "\n",
        "* m: the mean of the normal distribution.\n",
        "* s: the standard deviation of the distribution.\n",
        "* n: the total number of samples to be drawn."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 975
        },
        "id": "GiWH4SwoI7jA",
        "outputId": "a9c5a587-ce74-407b-b1ab-af378ba0d570"
      },
      "outputs": [],
      "source": [
        "#  TRAINING\n",
        "# Adversarial loss\n",
        "valid = 1\n",
        "fake = 0\n",
        "\n",
        "prev_time = time.time()\n",
        "for epoch in range(hp.epoch, hp.n_epochs):\n",
        "    for i, batch in enumerate(train_dataloader):\n",
        "\n",
        "        # Set model input\n",
        "        real_A = Variable(batch[\"A\"].type(Tensor))\n",
        "        real_B = Variable(batch[\"B\"].type(Tensor))\n",
        "\n",
        "        ################################\n",
        "        #  Train Generator and Encoder\n",
        "        #################################\n",
        "        '''  The first component of Bicycle-GAN is cVAE-GAN. It first encodes the ground truth image(B) into the latent space using encoder E.\n",
        "\n",
        "        ### Then input image(A) and encoded ground truth image(i.e. latent vector - Z) are passed into the Generator G which produces the output image(B^). That is, The generator attempts to map the input image A along with a sampled z back into the original image B.\n",
        "        the flow of cVAE-GAN => B -> Z -> B^\n",
        "        '''\n",
        "\n",
        "        optimizer_E.zero_grad()\n",
        "        optimizer_G.zero_grad()\n",
        "\n",
        "        #################\n",
        "        # cVAE-GAN\n",
        "        #################\n",
        "\n",
        "        # Produce output using encoding of B (cVAE-GAN)\n",
        "        mu, logvar = encoder(real_B)\n",
        "        # reparameterize so backprogation can be done on the\n",
        "        # stochastically generated z variable\n",
        "        encoded_z = reparameterization(mu, logvar)\n",
        "        fake_B = generator(real_A, encoded_z)\n",
        "\n",
        "        # Pixelwise loss of translated image by VAE\n",
        "        loss_pixel_L1_vae = mae_loss(fake_B, real_B)\n",
        "        \n",
        "        # Kullback-Leibler divergence of encoded B\n",
        "        # Refer - https://stackoverflow.com/questions/61597340/how-is-kl-divergence-in-pytorch-code-related-to-the-formula\n",
        "        loss_kl = 0.5 * torch.sum(torch.exp(logvar) + mu ** 2 - logvar - 1)\n",
        "        # Adversarial loss\n",
        "        # compute_loss() - Computes the MSE between model output and scalar ground_truth\n",
        "        loss_VAE_GAN = D_VAE.compute_loss(fake_B, valid)\n",
        "\n",
        "        ####################################\n",
        "        # cLR-GAN\n",
        "        # Conditional Latent Regressor GAN\n",
        "        ####################################\n",
        "        \n",
        "        ''' This is the second component of the Bicycle-GAN. Here a randomly drawn latent vector (sampled_z below) along with the input image(A) is provided to the generator. The generated output(B^) may not look like ground truth image(B), but it should look realistic.\n",
        "        Then the generated output is passed through the encoder, encoder tries to regain the latent vector from the output image. \n",
        "        \n",
        "        the flow of cLR-GAN is Z -> B^ -> Z^\n",
        "        '''\n",
        "\n",
        "        # real_A need to be a 4-D Tensor of Batch_size, Channel, Height, Width\n",
        "        # ic(real_A.size()) # torch.Size([8, 3, 128, 128])\n",
        "        # Produce output using sampled z (cLR-GAN)\n",
        "        # sampled_z need to be a 2-D Tensor of Batch_size (i.e. 8) and Latent_dim (i.e 8)\n",
        "        sampled_z = Variable(Tensor(np.random.normal(0, 1, (real_A.size(0), hp.latent_dim))))\n",
        "        # Draw random samples from a normal (Gaussian) distribution.\n",
        "        # ic(sampled_z.size()) # torch.Size([8, 8])\n",
        "        ''' ic(sampled_z) will output below\n",
        "        tensor([[ 0.2202,  0.3848, -1.0489, -0.5884, -0.0094,  0.1678, -1.5106, -0.2802],\n",
        "                       [ 1.9893, -0.8738, -2.4284,  1.0219,  0.4162,  0.3345, -1.7501,  0.0511],\n",
        "                       [-0.5878,  0.0200, -0.9107,  0.8697, -1.8777,  0.3819,  0.4788, -3.0111],\n",
        "                       [ 0.9784,  0.7068, -0.2902,  0.9084,  0.2643, -0.0090,  0.6117, -1.4820],\n",
        "                       [-0.1651, -0.5930,  0.1434,  1.6912, -0.6974, -1.7704, -3.2449, -0.4491],\n",
        "                       [-1.1474, -1.6918,  1.6326, -0.7691,  0.4847,  2.0244, -0.3476, -1.1350],\n",
        "                       [ 0.3965,  0.3789,  0.3177, -1.5840,  0.6150, -1.4996,  0.2707, -0.7499],\n",
        "                       [-0.1608, -0.8812, -0.0048, -0.6618, -0.5037, -0.8252,  0.7017,  0.4248]],\n",
        "                      device='cuda:0') '''\n",
        "        \n",
        "        _fake_B = generator(real_A, sampled_z)\n",
        "        # cLR Loss: Adversarial loss\n",
        "        loss_CLR_GAN = D_LR.compute_loss(_fake_B, valid)\n",
        "\n",
        "        #########################################\n",
        "        # Total Loss (Generator + Encoder)\n",
        "        #########################################\n",
        "\n",
        "        loss_total_gen_encoder = loss_VAE_GAN + loss_CLR_GAN + hp.lambda_pixel * loss_pixel_L1_vae + hp.lambda_kl * loss_kl\n",
        "\n",
        "        loss_total_gen_encoder.backward(retain_graph=True)\n",
        "        optimizer_E.step()\n",
        "\n",
        "        ######################\n",
        "        # Generator Only Loss\n",
        "        ######################\n",
        "        '''Under CLR-GAN - Then the generated output (_fake_B) is passed through the encoder, encoder tries to regain the latent vector from the output image. '''\n",
        "\n",
        "        # Latent L1 loss\n",
        "        _mu, _ = encoder(_fake_B)\n",
        "        loss_latent = hp.lambda_latent * mae_loss(_mu, sampled_z)\n",
        "\n",
        "        loss_latent.backward()\n",
        "        optimizer_G.step()\n",
        "\n",
        "        #######################################\n",
        "        #  Train Discriminator (cVAE-GAN)\n",
        "        #######################################\n",
        "\n",
        "        optimizer_D_VAE.zero_grad()\n",
        "\n",
        "        loss_D_VAE = D_VAE.compute_loss(real_B, valid) + D_VAE.compute_loss(fake_B.detach(), fake)\n",
        "\n",
        "        loss_D_VAE.backward()\n",
        "        optimizer_D_VAE.step()\n",
        "\n",
        "        ####################################\n",
        "        # Train Discriminator (cLR-GAN)\n",
        "        # Conditional Latent Regressor GAN\n",
        "        ####################################\n",
        "\n",
        "        optimizer_D_LR.zero_grad()\n",
        "\n",
        "        loss_D_LR = D_LR.compute_loss(real_B, valid) + D_LR.compute_loss(_fake_B.detach(), fake)\n",
        "\n",
        "        loss_D_LR.backward()\n",
        "        optimizer_D_LR.step()\n",
        "\n",
        "        #################\n",
        "        #  Log Progress\n",
        "        #################\n",
        "\n",
        "        # Determine approximate time left\n",
        "        batches_done = epoch * len(train_dataloader) + i\n",
        "        batches_left = hp.n_epochs * len(train_dataloader) - batches_done\n",
        "        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n",
        "        prev_time = time.time()\n",
        "\n",
        "        # Print log\n",
        "        sys.stdout.write(\n",
        "            \"\\r[Epoch %d/%d] [Batch %d/%d] [D VAE_loss: %f, LR_loss: %f] [G loss: %f, pixel: %f, kl: %f, latent: %f] ETA: %s\"\n",
        "            % (\n",
        "                epoch,\n",
        "                hp.n_epochs,\n",
        "                i,\n",
        "                len(train_dataloader),\n",
        "                loss_D_VAE.item(),\n",
        "                loss_D_LR.item(),\n",
        "                loss_total_gen_encoder.item(),\n",
        "                loss_pixel_L1_vae.item(),\n",
        "                loss_kl.item(),\n",
        "                loss_latent.item(),\n",
        "                time_left,\n",
        "            )\n",
        "        )\n",
        "        \n",
        "        # If at sample interval save image\n",
        "        if batches_done % hp.sample_interval == 0:\n",
        "          clear_output()\n",
        "          visualise_output(sample_images(batches_done), 30, 10)   "
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "GEN_5_BiCycleGAN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}