FACPA_attack_training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "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",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QSWwErWsaDCQ",
        "outputId": "5ff70f92-a2f0-4000-e4f7-c9b625d8c41f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download target NR model"
      ],
      "metadata": {
        "id": "dByk9fhZXf2h"
      }
    },
    {
      "cell_type": "code",
      "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://github.com/raw/baidut/PaQ-2-PiQ_GAE/master/paq2piq_standalone.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v4ERZzgUmB2e",
        "outputId": "d1e67862-f15e-4296-9ff3-f3fd240aa0e1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "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://github.com/raw/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"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install opencv-python-headless==4.5.2.52"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-KE1h9ukmESH",
        "outputId": "3d3e25a1-8ff4-487c-806a-673fc9a6177a"
      },
      "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",
            "\u001b[31mERROR: Could not find a version that satisfies the requirement opencv-python-headless==4.5.2.52 (from versions: 3.4.10.37, 3.4.11.39, 3.4.11.41, 3.4.11.43, 3.4.11.45, 3.4.13.47, 3.4.15.55, 3.4.16.59, 3.4.17.61, 3.4.17.63, 3.4.18.65, 4.3.0.38, 4.4.0.40, 4.4.0.42, 4.4.0.44, 4.4.0.46, 4.5.1.48, 4.5.3.56, 4.5.4.58, 4.5.4.60, 4.5.5.62, 4.5.5.64, 4.6.0.66, 4.7.0.68, 4.7.0.72)\u001b[0m\u001b[31m\n",
            "\u001b[0m\u001b[31mERROR: No matching distribution found for opencv-python-headless==4.5.2.52\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BxOM3_LRmGIz",
        "outputId": "1d1ee7bd-550e-44dc-a5bf-29c1decbbe17"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/content/paq2piq_standalone.py:46: SyntaxWarning: \"is\" with a literal. Did you mean \"==\"?\n",
            "  if backbone is 'resnet18':\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "cuda\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "BXSlXT3dmH8o"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "31OVJ8IBmKNe",
        "outputId": "8b5ce309-7198-434c-da14-d9b7518ad208"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "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"
          ]
        },
        {
          "output_type": "execute_result",
          "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",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Define attack model and some util functions for training"
      ],
      "metadata": {
        "id": "W3vkWESPXpHm"
      }
    },
    {
      "cell_type": "code",
      "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)]"
      ],
      "metadata": {
        "id": "thsMwDD3lA26"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "V0QXYA9VlFID"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "enOPsWAPlH-E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def normalize_and_scale(delta_im, mode='train'):\n",
        "    delta_im = (delta_im) * 10.0/255.0\n",
        "    return delta_im"
      ],
      "metadata": {
        "id": "nUjOOz5PlIf9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "netG = UnetGenerator(3, 3, 64, norm_type='instance', act_type='relu').to(device)\n",
        "netG.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fUdMMHqDlLbP",
        "outputId": "cfc0874f-7582-4a4d-e06e-0542278d2cda"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "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",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "id": "ivo9i4_6ld6W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Attacked model training loop"
      ],
      "metadata": {
        "id": "WESa16onYHB6"
      }
    },
    {
      "cell_type": "code",
      "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')"
      ],
      "metadata": {
        "id": "l4C0SMdVllI5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "5IONEw6mZr5o"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}