retargetting_basic_encoder_decoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "retargetting_basic_encoder_decoder.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "! pip install kornia"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hXyHeruE6k6q",
        "outputId": "05a3857d-f8ce-48ee-c88a-67f48b4ec645"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: kornia in /usr/local/lib/python3.7/dist-packages (0.6.5)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from kornia) (21.3)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.7/dist-packages (from kornia) (1.11.0+cu113)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.8.1->kornia) (4.2.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->kornia) (3.0.9)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sRCw1eISxI31"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import torch\n",
        "import  torchvision\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision.transforms as tt\n",
        "from torchvision.datasets import VOCDetection\n",
        "from torchvision.transforms import ToTensor\n",
        "from torchvision.utils import make_grid\n",
        "from torch.utils.data.dataloader import DataLoader\n",
        "from torch.utils.data import random_split\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from kornia.geometry.transform import get_perspective_transform, warp_perspective\n",
        "from skimage import transform as trf"
      ],
      "metadata": {
        "id": "JQfoK4NH-gO0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dataset = VOCDetection(root='data/',year ='2012',\n",
        "                       download=True, transform=ToTensor())\n",
        "test_dataset = VOCDetection(root='data/', year = '2012',\n",
        "                            image_set='val', transform=ToTensor())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pNUjFGULxMUs",
        "outputId": "cb375ed4-fb53-4a1d-d1d0-105de67f3f5f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using downloaded and verified file: data/VOCtrainval_11-May-2012.tar\n",
            "Extracting data/VOCtrainval_11-May-2012.tar to data/\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_categories(labels_dir):\n",
        "    if not os.path.isdir(labels_dir):\n",
        "        raise FileNotFoundError\n",
        "    else:\n",
        "        categories = []\n",
        "        for file in os.listdir(labels_dir):\n",
        "            if file.endswith(\"_train.txt\"):\n",
        "                categories.append(file.split(\"_\")[0])\n",
        "        return categories\n",
        "\n",
        "object_categories = get_categories('./data/VOCdevkit/VOC2012/ImageSets/Main')\n",
        "\n",
        "def encode_labels(target):\n",
        "    ls = target['annotation']['object']\n",
        "    j = []\n",
        "    if type(ls) == dict:\n",
        "        if int(ls['difficult']) == 0:\n",
        "            j.append(object_categories.index(ls['name']))\n",
        "    else:\n",
        "        for i in range(len(ls)):\n",
        "            if int(ls[i]['difficult']) == 0:\n",
        "                j.append(object_categories.index(ls[i]['name']))\n",
        "    k = np.zeros(len(object_categories))\n",
        "    k[j] = 1\n",
        "    return torch.from_numpy(k)\n",
        "\n",
        "def denormalize(images, means, stds):\n",
        "    means = torch.tensor(means).reshape(1, 3, 1, 1)\n",
        "    stds = torch.tensor(stds).reshape(1, 3, 1, 1)\n",
        "    return images * stds + means\n",
        "\n",
        "def show_batch(dl):\n",
        "    for images, labels in dl:\n",
        "        fig, ax = plt.subplots(figsize=(12, 12))\n",
        "        ax.set_xticks([]); ax.set_yticks([])\n",
        "        denorm_images = denormalize(images, *stats)\n",
        "        ax.imshow(make_grid(denorm_images[:64], nrow=8).permute(1, 2, 0).clamp(0,1))\n",
        "        break"
      ],
      "metadata": {
        "id": "9Fh7JO1GyP4o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Data transforms (normalization & data augmentation)\n",
        "stats = ((0.45704722, 0.43824774, 0.4061733),(0.23908591, 0.23509644, 0.2397309))\n",
        "train_tfms=tt.Compose([tt.Resize((256, 256)),\n",
        "                       tt.RandomChoice([tt.ColorJitter(brightness=(0.80, 1.20)),\n",
        "                                        tt.RandomGrayscale(p = 0.25)]),\n",
        "                       tt.RandomHorizontalFlip(p = 0.25),\n",
        "                       tt.ToTensor(), \n",
        "                       tt.Normalize(*stats,inplace=True)])\n",
        "\n",
        "valid_tfms = tt.Compose([tt.Resize(256), \n",
        "                         tt.CenterCrop(200),\n",
        "                         tt.ToTensor(), \n",
        "                         tt.Normalize(*stats)])"
      ],
      "metadata": {
        "id": "lbFd2uW-xN5g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_ds = VOCDetection(root='data/', year ='2012',\n",
        "                        download=True,\n",
        "                        transform=train_tfms, target_transform=encode_labels)\n",
        "test_dataset = VOCDetection(root='data/', year = '2012',\n",
        "                            image_set='val',\n",
        "                            transform=valid_tfms,target_transform=encode_labels)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cXTScrN8xuiv",
        "outputId": "cac42e48-404e-444f-c2f9-900398d25172"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using downloaded and verified file: data/VOCtrainval_11-May-2012.tar\n",
            "Extracting data/VOCtrainval_11-May-2012.tar to data/\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(43)\n",
        "val_size = 4500\n",
        "batch_size = 16"
      ],
      "metadata": {
        "id": "2EY4VE8Sx-S_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_size = len(test_dataset) - val_size\n",
        "test_ds, val_ds = random_split(test_dataset, [test_size, val_size])\n",
        "len(test_ds), len(val_ds)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ewORs1Oex9S6",
        "outputId": "9085a8bf-fb1a-4834-efec-3dc88d383ea6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1323, 4500)"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=4, pin_memory=True)\n",
        "val_dl = DataLoader(val_ds, batch_size, num_workers=4, pin_memory=True)\n",
        "test_dl = DataLoader(test_ds, batch_size, num_workers=4, pin_memory=True)"
      ],
      "metadata": {
        "id": "ZT6tHSS02zdn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def to_device(data, device):\n",
        "    \"\"\"Move tensor(s) to chosen device\"\"\"\n",
        "    if isinstance(data, (list,tuple)):\n",
        "        return [to_device(x, device) for x in data]\n",
        "    return data.to(device, non_blocking=True)\n",
        "\n",
        "class DeviceDataLoader():\n",
        "    \"\"\"Wrap a dataloader to move data to a device\"\"\"\n",
        "    def __init__(self, dl, device):\n",
        "        self.dl = dl\n",
        "        self.device = device\n",
        "        \n",
        "    def __iter__(self):\n",
        "        \"\"\"Yield a batch of data after moving it to device\"\"\"\n",
        "        for b in self.dl: \n",
        "            yield to_device(b, self.device)\n",
        "\n",
        "    def __len__(self):\n",
        "        \"\"\"Number of batches\"\"\"\n",
        "        return len(self.dl)"
      ],
      "metadata": {
        "id": "fx0QNlRA2kOE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
        "torch.cuda.is_available()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rAXl11QqyDGl",
        "outputId": "597f3e09-57e7-4914-ceeb-cdab3c7ea234"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_dl = DeviceDataLoader(train_dl, torch.device('cpu'))\n",
        "val_dl   = DeviceDataLoader(val_dl, torch.device('cpu'))\n",
        "test_dl  = DeviceDataLoader(test_dl, torch.device('cpu'))"
      ],
      "metadata": {
        "id": "cliJaMhH2lO2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "HEIGHT = WIDTH = 256\n",
        "AR = 0.8"
      ],
      "metadata": {
        "id": "mRAL1oRF4SY4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def cumulative_norm(x):\n",
        "    y = torch.sum(x, -1).view(-1, 3, 256, 1)\n",
        "    x = torch.cumsum(x, -1)\n",
        "    return x / y"
      ],
      "metadata": {
        "id": "oT9mvafs6B01"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def inverse_warp(input, flow):\n",
        "    shape = input.shape     # [1, 3, 256, 256]\n",
        "    i_H = shape[2] # 256\n",
        "    i_W = shape[3] # 256\n",
        "    shape = flow.shape      # [1, 3, 256, 408]\n",
        "    N = shape[1] # 3\n",
        "    H = shape[2] # 256\n",
        "    W = shape[3] # 408\n",
        "\n",
        "    N_i = torch.range(0, N - 1)  # [0, ..., 1]\n",
        "    W_i = torch.range(0, W - 1)  # [0, ..., 408]\n",
        "    H_i = torch.range(0, H - 1)  # [0, ..., 256]\n",
        "\n",
        "    n, h, w = torch.meshgrid(N_i, H_i, W_i, indexing='ij')\n",
        "    n = n.repeat(1,1,1,1)  # torch.unsqueeze(n, dim=-1)  # [1, 3, 256, 408]\n",
        "    h = h.repeat(1,1,1,1)  # [1, 3, 256, 408]\n",
        "    w = w.repeat(1,1,1,1)  # [1, 3, 256, 408]\n",
        "\n",
        "    n = n.double()\n",
        "    h = h.double()\n",
        "    w = w.double()\n",
        "\n",
        "    # print('flow:', flow.shape)  # [1, 3, 256, 408]\n",
        "    v_col, v_row = torch.split(flow, [204, 204], dim=-1)\n",
        "    # print('v_col', v_col.shape)  # [1, 3, 256, 204]\n",
        "    # print('v_row', v_row.shape)  # [1, 3, 256, 204]\n",
        "    \"\"\" calculate index \"\"\"\n",
        "    v_r0 = torch.floor(v_row)\n",
        "    v_r1 = v_r0 + 1\n",
        "    v_c0 = torch.floor(v_col)\n",
        "    v_c1 = v_c0 + 1\n",
        "\n",
        "    H_ = float(i_H - 1)\n",
        "    W_ = float(i_W - 1)\n",
        "    \n",
        "    i_r0 = torch.clamp(v_r0, 0., H_)\n",
        "    i_r1 = torch.clamp(v_r1, 0., H_)\n",
        "    i_c0 = torch.clamp(v_c0, 0., W_)\n",
        "    i_c1 = torch.clamp(v_c1, 0., W_)\n",
        "\n",
        "    i_r0c0 = torch.concat([n.cuda(), i_r0.cuda(), i_c0.cuda()], -1).double() # [N, H, W, 3]\n",
        "    i_r0c1 = torch.concat([n.cuda(), i_r0.cuda(), i_c1.cuda()], -1).double()\n",
        "    i_r1c0 = torch.concat([n.cuda(), i_r1.cuda(), i_c0.cuda()], -1).double()\n",
        "    i_r1c1 = torch.concat([n.cuda(), i_r1.cuda(), i_c1.cuda()], -1).double()\n",
        "    print('i_r0c0', i_r0c0.shape)\n",
        "    \"\"\" take value from index \"\"\"\n",
        "    f00 = torch.index_select(input, 0, i_r0c0) # [N, H, W, C]\n",
        "    f01 = torch.index_select(input, -1, i_r0c1)\n",
        "    f10 = torch.index_select(input, -1, i_r1c0)\n",
        "    f11 = torch.index_select(input, -1, i_r1c1)\n",
        "\n",
        "    \"\"\" calculate coeff \"\"\"\n",
        "    w00 = (v_r1 - v_row) * (v_c1 - v_col)\n",
        "    w01 = (v_r1 - v_row) * (v_col - v_c0)\n",
        "    w10 = (v_row - v_r0) * (v_c1 - v_col)\n",
        "    w11 = (v_row - v_r0) * (v_col - v_c0)\n",
        "\n",
        "    out = w00 * f00 + w01 * f01 + w10 * f10 + w11 * f11\n",
        "    return out"
      ],
      "metadata": {
        "id": "d2USjfzN7HbX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MySkimageTransform(object):\n",
        "    def __init__(self):\n",
        "        '''\n",
        "        initialize your transformation here, if necessary\n",
        "        '''\n",
        "    def __call__(self, pic):\n",
        "        batch = pic.cpu().detach().numpy()\n",
        "        afine_tf = trf.AffineTransform(shear=0.4,\n",
        "                                       translation=[(1-AR) * WIDTH, 0])\n",
        "        for img in batch:\n",
        "            # apply your transformation\n",
        "            img = trf.warp(img, inverse_map=afine_tf)\n",
        "        return batch"
      ],
      "metadata": {
        "id": "T5PhYVQGOhcK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "SsdsZJN33ZVu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class RetargetModel(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super(RetargetModel, self).__init__()\n",
        "        self.vgg16_model = torchvision.models.vgg16(pretrained=True)\n",
        "        self.vgg16_model.classifier[6] = nn.Linear(in_features=4096,\n",
        "                                                   out_features=20, bias=True)\n",
        "        self.decoder = nn.Sequential(\n",
        "            nn.Upsample(size=(16, 16)),\n",
        "            nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            # nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            # nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "\n",
        "            nn.Upsample(size=(32, 32)),\n",
        "            nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            # nn.ConvTranspose2d(512, 512, kernel_size=(3, 3)), nn.ELU(),\n",
        "\n",
        "            nn.Upsample(scale_factor=(64, 64)),\n",
        "            nn.ConvTranspose2d(512, 256, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            nn.ConvTranspose2d(256, 256, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((1, 1)),\n",
        "            # nn.ConvTranspose2d(256, 256, kernel_size=(3, 3)), nn.ELU(),\n",
        "\n",
        "            nn.Upsample(scale_factor=(128, 128)),\n",
        "            nn.ConvTranspose2d(256, 128, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # nn.ConvTranspose2d(128, 128, kernel_size=(3, 3)), nn.ELU(),\n",
        "\n",
        "            nn.Upsample(size=(HEIGHT, WIDTH)),\n",
        "            nn.ConvTranspose2d(128, 64, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((HEIGHT - 1, WIDTH - 1)),\n",
        "            nn.ConvTranspose2d(64, 3, kernel_size=(3, 3)), nn.ELU(),\n",
        "            # tt.CenterCrop((HEIGHT - 1, WIDTH - 1)),\n",
        "        )\n",
        "        \n",
        "        self.resize = tt.Resize((HEIGHT, int(AR*WIDTH)))\n",
        "        self.revert_resize = tt.Resize((HEIGHT, WIDTH))\n",
        "\n",
        "        self.conv1d = nn.Sequential(\n",
        "            nn.Conv2d(3, 3, kernel_size=(1, HEIGHT),\n",
        "                      padding='valid'),\n",
        "            nn.ReLU()\n",
        "        )\n",
        "        # Spatial transformer localization-network\n",
        "        self.localization = nn.Sequential(\n",
        "            nn.Conv2d(3, 8, kernel_size=7),\n",
        "            nn.MaxPool2d(2, stride=2),\n",
        "            nn.ReLU(True),\n",
        "            nn.Conv2d(8, 16, kernel_size=5),\n",
        "            nn.MaxPool2d(2, stride=2),\n",
        "            nn.ReLU(True),\n",
        "            nn.Conv2d(16, 32, kernel_size=5),\n",
        "            nn.MaxPool2d(2, stride=2),\n",
        "            nn.ReLU(True)\n",
        "        )\n",
        "\n",
        "        # Regressor for the 3 * 2 affine matrix\n",
        "        self.fc_loc = nn.Sequential(\n",
        "            nn.Linear(24 * 16 * 7 * 7, 32),\n",
        "            nn.ReLU(True),\n",
        "            nn.Linear(32, 3 * 2)\n",
        "        )\n",
        "        # Initialize the weights/bias with identity transformation\n",
        "        self.fc_loc[2].weight.data.zero_()\n",
        "        self.fc_loc[2].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0],\n",
        "                                                    dtype=torch.float))\n",
        "\n",
        "    def forward(self, input_image):\n",
        "        encoded_image = self.vgg16_model.features(input_image)\n",
        "        print('encded: ', encoded_image.size())\n",
        "        output_image = self.decoder(encoded_image)\n",
        "        print('decoded: ', output_image.size())\n",
        "        resized_image = self.resize(output_image)    # [32, 3, 256, 204]\n",
        "        \n",
        "        # DUPLICATE LAYER IMPLEMENTATION\n",
        "        duplicate_image = self.conv1d(output_image)    # [32, 3, 256, 1]\n",
        "        duplicate_image = duplicate_image.repeat(1, 1, 1, int(AR*WIDTH)) # [32, 3, 256, 204]\n",
        "        output_map = torch.add(resized_image, duplicate_image) # [1, 3, 256, 204]\n",
        "        \n",
        "        cum_map = torch.cumsum(output_map, -1) # [1, 3, 256, 204]\n",
        "        cum_image = torch.cumsum(input_image, -1) # [1, 3, 256, 256]\n",
        "\n",
        "        # NORM LAYER IMPLEMENTATION\n",
        "        output_map = np.abs(WIDTH - int(AR*WIDTH)) * cumulative_norm(output_map)\n",
        "        \n",
        "        # WARP LAYER IMPLEMENTATION\n",
        "        xs = self.localization(output_map)\n",
        "        xs = xs.view(-1, 24 * 16 * 7 * 7)\n",
        "        theta = self.fc_loc(xs)\n",
        "        theta = theta.view(-1, 2, 3)\n",
        "        grid = F.affine_grid(theta, output_map.size(), align_corners=True)\n",
        "        resized_image = F.grid_sample(resized_image, grid)\n",
        "\n",
        "        reverted_sized_image = self.revert_resize(resized_image)\n",
        "        output_label = self.vgg16_model(reverted_sized_image)\n",
        "\n",
        "        # getting structure loss\n",
        "        a1 = self.unet_model.encoder2(self.unet_model.encoder1(input_image))\n",
        "        a2 = self.vgg16_model.features[0](reverted_sized_image)\n",
        "        return resized_image, output_label, a1, a2"
      ],
      "metadata": {
        "id": "slOgSiIb6Bya"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = RetargetModel()\n",
        "model = model.to(device)\n",
        "out = model(torch.zeros((1, 3, 256, 256)).cuda())\n",
        "print(out[2].size(), out[3].size())\n",
        "plt.imshow(out[0].cpu().detach().numpy()[0][0])"
      ],
      "metadata": {
        "id": "9gyat38fGap2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "eYJXhA8IpHZ5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "id": "X9mEAUMhCWdl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "learning_rate = 1e-5\n",
        "# Create adam optimizer\n",
        "optimizer=torch.optim.Adam(params=model.parameters(), lr=learning_rate)"
      ],
      "metadata": {
        "id": "WoGlSffk089M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ContentLoss = nn.CrossEntropyLoss()\n",
        "# ContentLoss = nn.BCEWithLogitsLoss()\n",
        "StructureLoss = nn.CosineLoss()"
      ],
      "metadata": {
        "id": "_02gR5yR2Uil"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\n",
        "\n",
        "for epoch in tqdm(range(500)):\n",
        "    running_loss = 0.0\n",
        "    # loss = None\n",
        "    for i, data in enumerate(train_dl, 0):\n",
        "        # get the inputs; data is a list of [inputs, labels]\n",
        "        inputs, labels = data\n",
        "        inputs, labels = inputs.cuda(), labels.cuda()\n",
        "\n",
        "        # zero the parameter gradients\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        # forward + backward + optimize\n",
        "        images, predictions, a1, a2 = model(inputs)\n",
        "        content_loss = ContentLoss(predictions, labels)\n",
        "        structure_loss = StructureLoss(a1, a2)\n",
        "        avg = torch.tensor(0.5, requires_grad=True)\n",
        "        loss = (1 - avg) * content_loss + avg * structure_loss\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # print statistics\n",
        "        running_loss = running_loss + float(content_loss.item()) + float(structure_loss.item())\n",
        "        if i % 100 == 99:\n",
        "            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')\n",
        "            running_loss = 0.0\n",
        "            \n",
        "            plt.imshow(denormalize(inputs[0].cpu(), *stats).permute(0, 2, 3, 1).clamp(0,1).detach().numpy()[0])\n",
        "            plt.show()\n",
        "            plt.imshow(denormalize(images[0].cpu(), *stats).permute(0, 2, 3, 1).clamp(0,1).detach().numpy()[0])\n",
        "            plt.show()\n",
        "        # if epoch % 50 == 0:\n",
        "        #     PATH = \"model.pt\"\n",
        "        #     torch.save({\n",
        "        #         'epoch': epoch,\n",
        "        #         'model_state_dict': model.state_dict(),\n",
        "        #         'optimizer_state_dict': optimizer.state_dict(),\n",
        "        #         'loss': loss,\n",
        "        #         }, PATH)\n",
        "print('Finished Training')"
      ],
      "metadata": {
        "id": "G3750eX71VtV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "content_loss, structure_loss"
      ],
      "metadata": {
        "id": "PmEG3NZOvjCe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "denormalize(images[0].cpu(), *stats).size()"
      ],
      "metadata": {
        "id": "6n5uTZq16B87"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.permute(denormalize(images[0].cpu(), *stats), (0, 2,3,1)).size()"
      ],
      "metadata": {
        "id": "Dnjn8rkr5R_3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.imshow(denormalize(images[0].cpu(), *stats).permute(0, 2, 3, 1).clamp(0,1).detach().numpy()[0])\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "DcdACeor5GUS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "id": "UFXGSO6cBpib"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "r2Mq6U-4KVhp"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}