CNN-based_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dCBoVaqzk2R0"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QSWwErWsaDCQ",
    "outputId": "5ff70f92-a2f0-4000-e4f7-c9b625d8c41f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mounted at /content/drive\n"
     ]
    }
   ],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dByk9fhZXf2h"
   },
   "source": [
    "## Download target NR model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "v4ERZzgUmB2e",
    "outputId": "d1e67862-f15e-4296-9ff3-f3fd240aa0e1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: timestamping does nothing in combination with -O. See the manual\n",
      "for details.\n",
      "\n",
      "--2023-04-29 09:44:20--  https://github.com/baidut/PaQ-2-PiQ/releases/download/v1.0/RoIPoolModel-fit.10.bs.120.pth\n",
      "Resolving github.com (github.com)... 140.82.114.3\n",
      "Connecting to github.com (github.com)|140.82.114.3|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/237024974/a1c42500-4755-11ea-9c0e-7bf2246fe9e5?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230429%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230429T094420Z&X-Amz-Expires=300&X-Amz-Signature=f0fa15b8c744f2754db1ceffce079646ae90d898602b4254f89a44f93e4a1898&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=237024974&response-content-disposition=attachment%3B%20filename%3DRoIPoolModel-fit.10.bs.120.pth&response-content-type=application%2Foctet-stream [following]\n",
      "--2023-04-29 09:44:20--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/237024974/a1c42500-4755-11ea-9c0e-7bf2246fe9e5?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230429%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230429T094420Z&X-Amz-Expires=300&X-Amz-Signature=f0fa15b8c744f2754db1ceffce079646ae90d898602b4254f89a44f93e4a1898&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=237024974&response-content-disposition=attachment%3B%20filename%3DRoIPoolModel-fit.10.bs.120.pth&response-content-type=application%2Foctet-stream\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 140552190 (134M) [application/octet-stream]\n",
      "Saving to: ‘RoIPoolModel.pth’\n",
      "\n",
      "RoIPoolModel.pth    100%[===================>] 134.04M   199MB/s    in 0.7s    \n",
      "\n",
      "2023-04-29 09:44:21 (199 MB/s) - ‘RoIPoolModel.pth’ saved [140552190/140552190]\n",
      "\n",
      "--2023-04-29 09:44:22--  https://github.com/baidut/PaQ-2-PiQ/releases/download/v1.0/Picture1.jpg\n",
      "Resolving github.com (github.com)... 140.82.114.3\n",
      "Connecting to github.com (github.com)|140.82.114.3|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/237024974/b488b880-6231-11ea-8328-cd6123862fda?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230429%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230429T094422Z&X-Amz-Expires=300&X-Amz-Signature=7cf445fae2b752b58b7b9540c0b7c7ca07edfba6ab39d589c334fbb0509d83ac&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=237024974&response-content-disposition=attachment%3B%20filename%3DPicture1.jpg&response-content-type=application%2Foctet-stream [following]\n",
      "--2023-04-29 09:44:22--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/237024974/b488b880-6231-11ea-8328-cd6123862fda?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230429%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230429T094422Z&X-Amz-Expires=300&X-Amz-Signature=7cf445fae2b752b58b7b9540c0b7c7ca07edfba6ab39d589c334fbb0509d83ac&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=237024974&response-content-disposition=attachment%3B%20filename%3DPicture1.jpg&response-content-type=application%2Foctet-stream\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 110901 (108K) [application/octet-stream]\n",
      "Saving to: ‘Picture1.jpg’\n",
      "\n",
      "Picture1.jpg        100%[===================>] 108.30K  --.-KB/s    in 0.02s   \n",
      "\n",
      "2023-04-29 09:44:22 (5.12 MB/s) - ‘Picture1.jpg’ saved [110901/110901]\n",
      "\n",
      "--2023-04-29 09:44:22--  https://raw.githubusercontent.com/baidut/PaQ-2-PiQ_GAE/master/paq2piq_standalone.py\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 6486 (6.3K) [text/plain]\n",
      "Saving to: ‘paq2piq_standalone.py’\n",
      "\n",
      "paq2piq_standalone. 100%[===================>]   6.33K  --.-KB/s    in 0s      \n",
      "\n",
      "Last-modified header missing -- time-stamps turned off.\n",
      "2023-04-29 09:44:22 (60.5 MB/s) - ‘paq2piq_standalone.py’ saved [6486/6486]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# download the model\n",
    "!wget -O RoIPoolModel.pth -N https://github.com/baidut/PaQ-2-PiQ/releases/download/v1.0/RoIPoolModel-fit.10.bs.120.pth\n",
    "\n",
    "# download a test image\n",
    "!wget -N https://github.com/baidut/PaQ-2-PiQ/releases/download/v1.0/Picture1.jpg\n",
    "\n",
    "# download the standalone version of code\n",
    "!wget -N https://raw.githubusercontent.com/baidut/PaQ-2-PiQ_GAE/master/paq2piq_standalone.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BxOM3_LRmGIz",
    "outputId": "1d1ee7bd-550e-44dc-a5bf-29c1decbbe17"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/content/paq2piq_standalone.py:46: SyntaxWarning: \"is\" with a literal. Did you mean \"==\"?\n",
      "  if backbone is 'resnet18':\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from paq2piq_standalone import *\n",
    "from torch.autograd import Variable\n",
    "import imageio\n",
    "import os\n",
    "import subprocess \n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BXSlXT3dmH8o"
   },
   "outputs": [],
   "source": [
    "IMAGE_NET_MEAN = [0.485, 0.456, 0.406]\n",
    "IMAGE_NET_STD = [0.229, 0.224, 0.225]\n",
    "\n",
    "\n",
    "class Transform:\n",
    "    def __init__(self):\n",
    "        # normalize = transforms.Normalize(mean=IMAGE_NET_MEAN, std=IMAGE_NET_STD)\n",
    "\n",
    "        self._train_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.ToTensor(),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self._val_transform = transforms.Compose([transforms.ToTensor()])\n",
    "\n",
    "    @property\n",
    "    def train_transform(self):\n",
    "        return self._train_transform\n",
    "\n",
    "    @property\n",
    "    def val_transform(self):\n",
    "        return self._val_transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "31OVJ8IBmKNe",
    "outputId": "8b5ce309-7198-434c-da14-d9b7518ad208"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=None`.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RoIPoolModel(\n",
       "  (body): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU(inplace=True)\n",
       "    (3): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "    (4): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (5): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (6): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (7): Sequential(\n",
       "      (0): BasicBlock(\n",
       "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (downsample): Sequential(\n",
       "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (1): BasicBlock(\n",
       "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head): Sequential(\n",
       "    (0): AdaptiveConcatPool2d(\n",
       "      (ap): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "      (mp): AdaptiveMaxPool2d(output_size=(1, 1))\n",
       "    )\n",
       "    (1): Flatten(start_dim=1, end_dim=-1)\n",
       "    (2): BatchNorm1d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (3): Dropout(p=0.25, inplace=False)\n",
       "    (4): Linear(in_features=1024, out_features=512, bias=True)\n",
       "    (5): ReLU(inplace=True)\n",
       "    (6): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (7): Dropout(p=0.5, inplace=False)\n",
       "    (8): Linear(in_features=512, out_features=1, bias=True)\n",
       "  )\n",
       "  (roi_pool): RoIPool(output_size=(2, 2), spatial_scale=0.03125)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_state = torch.load('RoIPoolModel.pth', map_location=lambda storage, loc: storage)\n",
    "model = RoIPoolModel()\n",
    "model.load_state_dict(model_state[\"model\"])\n",
    "model = model.to(device)\n",
    "transform = Transform().val_transform\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W3vkWESPXpHm"
   },
   "source": [
    "## Define attack model and some util functions for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "thsMwDD3lA26"
   },
   "outputs": [],
   "source": [
    "def center_crop(image):\n",
    "  center = image.shape[0] / 2, image.shape[1] / 2\n",
    "  if center[1] < 128 or center[0] < 128:\n",
    "    return cv2.resize(image, (256, 256))\n",
    "  x = center[1] - 128\n",
    "  y = center[0] - 128\n",
    "\n",
    "  return image[int(y):int(y+256), int(x):int(x+256)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V0QXYA9VlFID"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt\n",
    "from torch.utils.data import Dataset\n",
    "from typing import Optional\n",
    "import os\n",
    "import subprocess \n",
    "import random\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.autograd import Variable\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torchvision import transforms\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "enOPsWAPlH-E"
   },
   "outputs": [],
   "source": [
    "class UnetGenerator(nn.Module):\n",
    "    def __init__(self, input_nc, output_nc, ngf, norm_type='batch', act_type='selu'):\n",
    "        super(UnetGenerator, self).__init__()\n",
    "        self.name = 'unet'\n",
    "        self.conv1 = nn.Conv2d(input_nc, ngf, 4, 2, 1)\n",
    "        self.conv2 = nn.Conv2d(ngf, ngf * 2, 4, 2, 1)\n",
    "        self.conv3 = nn.Conv2d(ngf * 2, ngf * 4, 4, 2, 1)\n",
    "        self.conv4 = nn.Conv2d(ngf * 4, ngf * 8, 4, 2, 1)\n",
    "        self.conv5 = nn.Conv2d(ngf * 8, ngf * 8, 4, 2, 1)\n",
    "        self.conv6 = nn.Conv2d(ngf * 8, ngf * 8, 4, 2, 1)\n",
    "        self.conv7 = nn.Conv2d(ngf * 8, ngf * 8, 4, 2, 1)\n",
    "        self.conv8 = nn.Conv2d(ngf * 8, ngf * 8, 4, 2, 1)\n",
    "        self.dconv1 = nn.ConvTranspose2d(ngf * 8, ngf * 8, 4, 2, 1)\n",
    "        self.dconv2 = nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, 4, 2, 1)\n",
    "        self.dconv3 = nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, 4, 2, 1)\n",
    "        self.dconv4 = nn.ConvTranspose2d(ngf * 8 * 2, ngf * 8, 4, 2, 1)\n",
    "        self.dconv5 = nn.ConvTranspose2d(ngf * 8 * 2, ngf * 4, 4, 2, 1)\n",
    "        self.dconv6 = nn.ConvTranspose2d(ngf * 4 * 2, ngf * 2, 4, 2, 1)\n",
    "        self.dconv7 = nn.ConvTranspose2d(ngf * 2 * 2, ngf, 4, 2, 1)\n",
    "        self.dconv8 = nn.ConvTranspose2d(ngf * 2, output_nc, 4, 2, 1)\n",
    "\n",
    "        if norm_type == 'batch':\n",
    "            self.norm = nn.BatchNorm2d(ngf)\n",
    "            self.norm2 = nn.BatchNorm2d(ngf * 2)\n",
    "            self.norm4 = nn.BatchNorm2d(ngf * 4)\n",
    "            self.norm8 = nn.BatchNorm2d(ngf * 8)\n",
    "        elif norm_type == 'instance':\n",
    "            self.norm = nn.InstanceNorm2d(ngf)\n",
    "            self.norm2 = nn.InstanceNorm2d(ngf * 2)\n",
    "            self.norm4 = nn.InstanceNorm2d(ngf * 4)\n",
    "            self.norm8 = nn.InstanceNorm2d(ngf * 8)\n",
    "        self.leaky_relu = nn.LeakyReLU(0.2, True)\n",
    "\n",
    "        if act_type == 'selu':\n",
    "            self.act = nn.SELU(True)\n",
    "        else:\n",
    "            self.act = nn.ReLU(True)\n",
    "\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "        self.tanh = nn.Tanh()\n",
    "\n",
    "    def forward(self, input):\n",
    "        # Encoder\n",
    "        # Convolution layers:\n",
    "        # input is (nc) x 512 x 1024\n",
    "        e1 = self.conv1(input)\n",
    "        # state size is (ngf) x 256 x 512\n",
    "        e2 = self.norm2(self.conv2(self.leaky_relu(e1)))\n",
    "        # state size is (ngf x 2) x 128 x 256\n",
    "        e3 = self.norm4(self.conv3(self.leaky_relu(e2)))\n",
    "        # state size is (ngf x 4) x 64 x 128\n",
    "        e4 = self.norm8(self.conv4(self.leaky_relu(e3)))\n",
    "        # state size is (ngf x 8) x 32 x 64\n",
    "        e5 = self.norm8(self.conv5(self.leaky_relu(e4)))\n",
    "        # state size is (ngf x 8) x 16 x 32\n",
    "        e6 = self.norm8(self.conv6(self.leaky_relu(e5)))\n",
    "        # state size is (ngf x 8) x 8 x 16\n",
    "        e7 = self.norm8(self.conv7(self.leaky_relu(e6)))\n",
    "        # state size is (ngf x 8) x 4 x 8\n",
    "        # No batch norm on output of Encoder\n",
    "        e8 = self.conv8(self.leaky_relu(e7))\n",
    "\n",
    "        # Decoder\n",
    "        # Deconvolution layers:\n",
    "        # state size is (ngf x 8) x 2 x 4\n",
    "        d1_ = self.dropout(self.norm8(self.dconv1(self.act(e8))))\n",
    "        # state size is (ngf x 8) x 4 x 8\n",
    "        d1 = torch.cat((d1_, e7), 1)\n",
    "        d2_ = self.dropout(self.norm8(self.dconv2(self.act(d1))))\n",
    "        # state size is (ngf x 8) x 8 x 16\n",
    "        d2 = torch.cat((d2_, e6), 1)\n",
    "        d3_ = self.dropout(self.norm8(self.dconv3(self.act(d2))))\n",
    "        # state size is (ngf x 8) x 16 x 32\n",
    "        d3 = torch.cat((d3_, e5), 1)\n",
    "        d4_ = self.norm8(self.dconv4(self.act(d3)))\n",
    "        # state size is (ngf x 8) x 32 x 64\n",
    "        d4 = torch.cat((d4_, e4), 1)\n",
    "        d5_ = self.norm4(self.dconv5(self.act(d4)))\n",
    "        # state size is (ngf x 4) x 64 x 128\n",
    "        d5 = torch.cat((d5_, e3), 1)\n",
    "        d6_ = self.norm2(self.dconv6(self.act(d5)))\n",
    "        # state size is (ngf x 2) x 128 x 256\n",
    "        d6 = torch.cat((d6_, e2), 1)\n",
    "        d7_ = self.norm(self.dconv7(self.act(d6)))\n",
    "        # state size is (ngf) x 256 x 512\n",
    "        d7 = torch.cat((d7_, e1), 1)\n",
    "        d8 = self.dconv8(self.act(d7))\n",
    "        # state size is (nc) x 512 x 1024\n",
    "        output = self.tanh(d8)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nUjOOz5PlIf9"
   },
   "outputs": [],
   "source": [
    "def normalize_and_scale(delta_im, mode='train'):\n",
    "    delta_im = (delta_im) * 10.0/255.0\n",
    "    return delta_im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fUdMMHqDlLbP",
    "outputId": "cfc0874f-7582-4a4d-e06e-0542278d2cda"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "UnetGenerator(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv3): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv4): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv5): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv6): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv7): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv8): Conv2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv1): ConvTranspose2d(512, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv2): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv3): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv4): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv5): ConvTranspose2d(1024, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv6): ConvTranspose2d(512, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv7): ConvTranspose2d(256, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (dconv8): ConvTranspose2d(128, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (norm): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "  (norm2): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "  (norm4): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "  (norm8): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "  (leaky_relu): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "  (act): ReLU(inplace=True)\n",
       "  (dropout): Dropout(p=0.5, inplace=False)\n",
       "  (tanh): Tanh()\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "netG = UnetGenerator(3, 3, 64, norm_type='instance', act_type='relu').to(device)\n",
    "netG.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ivo9i4_6ld6W"
   },
   "outputs": [],
   "source": [
    "class MyCustomDataset(Dataset):\n",
    "    def __init__(self, \n",
    "                 path_gt,\n",
    "                 mode\n",
    "                ):\n",
    "        \n",
    "        self._items = [] \n",
    "        self._index = 0\n",
    "        self.mode = mode\n",
    "        if mode == 'train':\n",
    "          dir_img = sorted(os.listdir(path_gt))\n",
    "        else:\n",
    "          dir_img = sorted(os.listdir(path_gt))\n",
    "        img_pathes = dir_img\n",
    "\n",
    "        for img_path in img_pathes:\n",
    "          self._items.append((\n",
    "            os.path.join(path_gt, img_path)\n",
    "          ))\n",
    "        random.shuffle(self._items)\n",
    "\n",
    "    def __len__(self):\n",
    "      return len(self._items)\n",
    "\n",
    "    def next_data(self):\n",
    "      gt_path = self._items[self._index]\n",
    "      self._index += 1 \n",
    "      if self._index == len(self._items):\n",
    "        self._index = 0\n",
    "        random.shuffle(self._items)\n",
    "\n",
    "      image = Image.open(gt_path).convert('RGB')\n",
    "      image = np.array(image).astype(np.float32) \n",
    "      image = center_crop(image)\n",
    "\n",
    "      image = image / 255.\n",
    "      image = transforms.ToTensor()(image)\n",
    "      y = image.to(device)\n",
    "      return y\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "      gt_path = self._items[index]\n",
    "      image = Image.open(gt_path).convert('RGB')\n",
    "      image = np.array(image).astype(np.float32) \n",
    "\n",
    "      image = center_crop(image)\n",
    "\n",
    "      image = image / 255.\n",
    "      image = transforms.ToTensor()(image)\n",
    "      y = image.to(device)\n",
    "      return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WESa16onYHB6"
   },
   "source": [
    "# Attacked model training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l4C0SMdVllI5"
   },
   "outputs": [],
   "source": [
    "lr = 0.001\n",
    "path_train = '/content/drive/MyDrive/metrics/train'\n",
    "path_test = '/content/drive/MyDrive/metrics/val'\n",
    "eps = 10/255\n",
    "\n",
    "optimizer = torch.optim.Adam(netG.parameters(), lr=lr)\n",
    "\n",
    "ds_train = MyCustomDataset(path_gt=path_train, mode='train')\n",
    "ds_val = MyCustomDataset(path_gt=path_test, mode='test')\n",
    "dl_train = DataLoader(ds_train, batch_size=4, shuffle=True)\n",
    "dl_val = DataLoader(ds_val, batch_size=4, shuffle=False)\n",
    "n_epoch = 200\n",
    "id = 0\n",
    "mean_train_loss = []\n",
    "mean_test_loss = []\n",
    "for epoch in range(n_epoch):\n",
    "  loss_list = []\n",
    "  netG.train()\n",
    "  for u in tqdm(range(len(dl_train))):\n",
    "    y = next(iter(dl_train))\n",
    "    delta_im = netG(y)\n",
    "    delta_im = normalize_and_scale(delta_im, 'train')\n",
    "    res = torch.add(y, delta_im.to(device))\n",
    "    res[res > 1] = 1\n",
    "    res[res < 0] = 0\n",
    "    a = 1 - model(res).mean() / 100\n",
    "    loss_list.append(a.item())\n",
    "    optimizer.zero_grad()\n",
    "    a.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "  print(epoch)\n",
    "  print('train loss = ' , np.array(loss_list).mean())\n",
    "  test_loss_list = []\n",
    "  netG.eval()\n",
    "  with torch.no_grad():\n",
    "    for u in tqdm(range(len(dl_val))):\n",
    "      y = next(iter(dl_val))\n",
    "      delta_im = netG(y)\n",
    "      delta_im = normalize_and_scale(delta_im, 'train')\n",
    "      out = torch.add(y, delta_im.to(device))\n",
    "      res[res > 1] = 1\n",
    "      res[res < 0] = 0\n",
    "      a = 1 - model(res).mean() / 100\n",
    "      test_loss_list.append(a.item())\n",
    "  print('valid loss = ' , np.array(test_loss_list).mean())\n",
    "  torch.save(netG.state_dict(), f'/content/drive/MyDrive/metrics/cnn_methods/unet/p2p/model_{epoch:02}.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5IONEw6mZr5o"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}