srgan.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "srgan.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1706a71588fd459c912f8415d292f51e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_2ad8257f35ee42c9acb50ec37fff7b0b",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_cccbc200789243b7b5d37240305146a0",
              "IPY_MODEL_086368cc6615494ba48f21a6a8dd523e"
            ]
          }
        },
        "2ad8257f35ee42c9acb50ec37fff7b0b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "cccbc200789243b7b5d37240305146a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_79794a7717d14b20b52cb9398d148926",
            "_dom_classes": [],
            "description": "Training Epoch 0 : 100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 12410,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 12410,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_6195e3574c8f49a09ea0e816d62a2d7f"
          }
        },
        "086368cc6615494ba48f21a6a8dd523e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_f8d9878fa74a44f3b495d01e2508a8fb",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 12410/12410 [2:04:22&lt;00:00,  1.66it/s, disc_loss=0.00535, gen_loss=0.598]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_944f531e80f94e19a105063181ea2317"
          }
        },
        "79794a7717d14b20b52cb9398d148926": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "6195e3574c8f49a09ea0e816d62a2d7f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "f8d9878fa74a44f3b495d01e2508a8fb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "944f531e80f94e19a105063181ea2317": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "7804d4dff2ac4cccbab7621a23890de4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_aaffd65f21504553bd6cd51f20158eab",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_09bca5c6c53443f8a9d1340707008f63",
              "IPY_MODEL_79b4eb77c12142e3838fafcdebfe9da5"
            ]
          }
        },
        "aaffd65f21504553bd6cd51f20158eab": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "09bca5c6c53443f8a9d1340707008f63": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_d89e5d7fe52d485ab1289387e1b1ed98",
            "_dom_classes": [],
            "description": "Training Epoch 1 :  60%",
            "_model_name": "FloatProgressModel",
            "bar_style": "",
            "max": 12410,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 7420,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_730cdbfff47241ecbc1c7ce4e2a11c6c"
          }
        },
        "79b4eb77c12142e3838fafcdebfe9da5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_b8b84f38186e4361a7c7c0fdd8d241be",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 7420/12410 [1:15:03&lt;44:00,  1.89it/s, disc_loss=0.00166, gen_loss=0.551]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_45ce6fc8d2694d0792231c1f8716cf9b"
          }
        },
        "d89e5d7fe52d485ab1289387e1b1ed98": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "730cdbfff47241ecbc1c7ce4e2a11c6c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "b8b84f38186e4361a7c7c0fdd8d241be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "45ce6fc8d2694d0792231c1f8716cf9b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jxpHaFcukWFw"
      },
      "source": [
        "SRGAN은 high resolution이미지에서 low resolution으로 만들어서 super resolution으로 이미지를 복원 시키는 모델이라고 생각하시면 됩니다.\n",
        "\n",
        "참고 링크:\n",
        "\n",
        "1) https://github.com/kunalrdeshmukh/SRGAN/blob/master/SRGAN.ipynb\n",
        "\n",
        "2) https://www.kaggle.com/balraj98single-image-super-resolution-gan-srgan-pytorch\n",
        "\n",
        "3) https://github.com/leftthomas/SRGAN/blob/master/data_utils.py\n",
        "\n",
        "4) https://github.com/deepak112/Keras-SRGAN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vj1rh4_9o5K0"
      },
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import os, math, sys\n",
        "import glob, itertools\n",
        "import argparse, random\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "from torchvision.models import vgg19\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "import torchvision.datasets as dset\n",
        "from torchvision.utils import save_image, make_grid\n",
        "\n",
        "import plotly\n",
        "import plotly.express as px\n",
        "import plotly.graph_objects as go\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from PIL import Image\n",
        "from tqdm import tqdm_notebook as tqdm\n",
        "from sklearn.model_selection import train_test_split\n",
        "\n",
        "random.seed(42)\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1b13VH1pr6Tt"
      },
      "source": [
        "Setting"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-02ZJjhsquFl"
      },
      "source": [
        "# number of epochs of training\n",
        "n_epochs = 2\n",
        "\n",
        "# name of the dataset: 폴더에서 경로복사 해주세요!\n",
        "dataset_path = \"/content/training_data/img_align_celeba\"\n",
        "\n",
        "# size of the batches\n",
        "# batch_size = 16 # 지금 적용함 - 예시\n",
        "batch_size = 64 # train\n",
        "# batch_size = 1 # test \n",
        "\n",
        "# adam: learning rate\n",
        "lr = 0.0001\n",
        "# adam: decay of first order momentum of gradient\n",
        "b1 = 0.5\n",
        "# adam: decay of second order momentum of gradient\n",
        "b2 = 0.999\n",
        "# number of cpu threads to use during batch generation = number of workers\n",
        "n_cpu = 8\n",
        "\n",
        "# high res. image height\n",
        "hr_height = 256\n",
        "# high res. image width\n",
        "hr_width = 256\n",
        "# number of image channels = rgb\n",
        "channels = 3\n",
        "\n",
        "cuda = torch.cuda.is_available()\n",
        "hr_shape = (hr_height, hr_width)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjlmLjn-r96y"
      },
      "source": [
        "# **Dataset- CelebA**\n",
        "예시입니다.\n",
        "\n",
        "아마도 coco train 2017의 chair & sofa를 사용하는게 좋을 것 같습니다."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0f4nH9j90GwP"
      },
      "source": [
        "# Normalization parameters for pre-trained PyTorch models\n",
        "mean = np.array([0.485, 0.456, 0.406])\n",
        "std = np.array([0.229, 0.224, 0.225])\n",
        "\n",
        "class ImageDataset(Dataset):\n",
        "    def __init__(self, files, hr_shape):\n",
        "        hr_height, hr_width = hr_shape\n",
        "        # Transforms for low resolution images and high resolution images\n",
        "        # low : high = 4 배\n",
        "        self.lr_transform = transforms.Compose(\n",
        "            [\n",
        "                transforms.Resize((hr_height // 4, hr_height // 4), Image.BICUBIC),\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize(mean, std),\n",
        "            ]\n",
        "        )\n",
        "        self.hr_transform = transforms.Compose(\n",
        "            [\n",
        "                transforms.Resize((hr_height, hr_height), Image.BICUBIC),\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize(mean, std),\n",
        "            ]\n",
        "        )\n",
        "        self.files = files\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        img = Image.open(self.files[index % len(self.files)])\n",
        "        img_lr = self.lr_transform(img)\n",
        "        img_hr = self.hr_transform(img)\n",
        "\n",
        "        return {\"lr\": img_lr, \"hr\": img_hr}\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.files)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JTbAxPJL8NSA",
        "outputId": "ca1eb03b-ff56-4e80-f949-a6f3039347ac"
      },
      "source": [
        "!wget https://s3-us-west-1.amazonaws.com/udacity-dlnfd/datasets/celeba.zip\n",
        "!mkdir training_data/ && unzip celeba.zip -d training_data/\n",
        "\n",
        "# dataset = dset.ImageFolder(root=\"training_data\")\n",
        "# train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=n_cpu)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-11-19 09:42:13--  https://s3-us-west-1.amazonaws.com/udacity-dlnfd/datasets/celeba.zip\n",
            "Resolving s3-us-west-1.amazonaws.com (s3-us-west-1.amazonaws.com)... 52.219.112.16\n",
            "Connecting to s3-us-west-1.amazonaws.com (s3-us-west-1.amazonaws.com)|52.219.112.16|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1443490838 (1.3G) [application/zip]\n",
            "Saving to: ‘celeba.zip.4’\n",
            "\n",
            "celeba.zip.4        100%[===================>]   1.34G  19.9MB/s    in 71s     \n",
            "\n",
            "utime(celeba.zip.4): No such file or directory\n",
            "2020-11-19 09:43:25 (19.5 MB/s) - ‘celeba.zip.4’ saved [1443490838/1443490838]\n",
            "\n",
            "mkdir: cannot create directory ‘training_data/’: File exists\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qA3RggSRo7rq"
      },
      "source": [
        "train_paths, test_paths = train_test_split(sorted(glob.glob(dataset_path + \"/*.*\")), test_size=0.02, random_state=42)\n",
        "train_dataloader = DataLoader(ImageDataset(train_paths, hr_shape=hr_shape), batch_size=batch_size, shuffle=True, num_workers=n_cpu)\n",
        "# test_dataloader = DataLoader(ImageDataset(test_data, hr_shape=hr_shape), batch_size=int(batch_size*0.75), shuffle=True, num_workers=n_cpu)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "llGSPT32Hjc5"
      },
      "source": [
        "# **SRGAN** Model Define"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yMLZVTXGtm1j"
      },
      "source": [
        "def conv(ch_in, ch_out, k_size, stride=1, pad=1):\n",
        "    layers = []\n",
        "    layers.append(nn.Conv2d(ch_in, ch_out, k_size, stride, pad))\n",
        "    return nn.Sequential(*layers)\n",
        "\n",
        "def deconv(ch_in, ch_out, k_size, stride=2, pad=1, bn=False):\n",
        "    layers = []\n",
        "    layers.append(nn.ConvTranspose2d(ch_in, ch_out, k_size, stride, pad))\n",
        "    if bn:\n",
        "        layers.append(nn.BatchNorm2d(ch_out))\n",
        "    return nn.Sequential(*layers)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gUH8otfZLPdA"
      },
      "source": [
        "#VGG19를 사용한 Fixed Feature Extraction\n",
        "class FeatureExtractor(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(FeatureExtractor, self).__init__()\n",
        "\n",
        "    vgg19_model = vgg19(pretrained=True)\n",
        "    self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:18])\n",
        "\n",
        "  def forward(self, img):\n",
        "    return self.feature_extractor(img)\n",
        "\n",
        "class ResnetBlock(nn.Module):\n",
        "    def __init__(self, in_features):\n",
        "        super(ResnetBlock, self).__init__()\n",
        "                \n",
        "        self.conv_block = nn.Sequential(\n",
        "            conv(in_features, in_features, 3, 1, 1),\n",
        "            nn.BatchNorm2d(in_features, 0.8),\n",
        "            nn.PReLU(),\n",
        "            conv(in_features, in_features, 3, 1, 1),\n",
        "            nn.BatchNorm2d(in_features, 0.8))\n",
        "\n",
        "    def forward(self, x: torch.Tensor):\n",
        "        out = x + self.conv_block(x)\n",
        "        return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m9dn63hYuL9U"
      },
      "source": [
        "class Discriminator(nn.Module):\n",
        "  def __init__(self, input_shape):\n",
        "    super(Discriminator, self).__init__()\n",
        "\n",
        "    self.input_shape = input_shape\n",
        "    in_channels, in_height, in_width = self.input_shape\n",
        "    patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)\n",
        "    self.output_shape = (1, patch_h, patch_w)\n",
        "\n",
        "    def discriminator_block(in_filters, out_filters, first_block=False):\n",
        "        layers = []\n",
        "        layers.append(conv(in_filters, out_filters, 3, 1, 1))\n",
        "        if not first_block:\n",
        "            layers.append(nn.BatchNorm2d(out_filters))\n",
        "        layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
        "        layers.append(conv(out_filters, out_filters, 3, 2, 1))\n",
        "        layers.append(nn.BatchNorm2d(out_filters))\n",
        "        layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
        "        return layers\n",
        "\n",
        "    layers = []\n",
        "    in_filters = in_channels\n",
        "    for i, out_filters in enumerate([64, 128, 256, 512]):\n",
        "        layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))\n",
        "        in_filters = out_filters\n",
        "\n",
        "    layers.append(conv(out_filters, 1, 3, 1, 1))\n",
        "\n",
        "    self.model = nn.Sequential(*layers)\n",
        "  def forward(self, img):\n",
        "    return self.model(img)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wNrAJ-KtH58C"
      },
      "source": [
        "class GeneratorResNet(nn.Module):\n",
        "    def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16):\n",
        "        super(GeneratorResNet, self).__init__()\n",
        "\n",
        "        # First layer\n",
        "        self.conv1 = nn.Sequential(conv(in_channels, 64, 9, 1, 4), nn.PReLU())\n",
        "\n",
        "        # Residual blocks\n",
        "        res_blocks = []\n",
        "        for _ in range(n_residual_blocks):\n",
        "            res_blocks.append(ResnetBlock(64))\n",
        "        self.res_blocks = nn.Sequential(*res_blocks)\n",
        "\n",
        "        # Second conv layer post residual blocks\n",
        "        self.conv2 = nn.Sequential(conv(64, 64, 3, 1, 1), nn.BatchNorm2d(64, 0.8))\n",
        "\n",
        "        # Upsampling layers\n",
        "        upsampling = []\n",
        "        for out_features in range(2):\n",
        "            upsampling += [\n",
        "                # nn.Upsample(scale_factor=2),\n",
        "                nn.Conv2d(64, 256, 3, 1, 1),\n",
        "                nn.BatchNorm2d(256),\n",
        "                nn.PixelShuffle(upscale_factor=2),\n",
        "                nn.PReLU(),\n",
        "            ]\n",
        "        self.upsampling = nn.Sequential(*upsampling)\n",
        "\n",
        "        # Final output layer\n",
        "        self.conv3 = nn.Sequential(conv(64, out_channels, 9, 1, 4), nn.Tanh())\n",
        "\n",
        "    def forward(self, x):\n",
        "        out1 = self.conv1(x)\n",
        "        out = self.res_blocks(out1)\n",
        "        out2 = self.conv2(out)\n",
        "        out = torch.add(out1, out2)\n",
        "        out = self.upsampling(out)\n",
        "        out = self.conv3(out)\n",
        "        return out"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "N9qjr-1KXX8m",
        "outputId": "5b0c3eee-de32-4db2-c34a-f3108ec88b36"
      },
      "source": [
        "netG = GeneratorResNet().cuda()\n",
        "netD = Discriminator(input_shape=(channels, *hr_shape)).cuda()\n",
        "feature_extractor = FeatureExtractor().cuda()\n",
        "\n",
        "feature_extractor.eval()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "FeatureExtractor(\n",
              "  (feature_extractor): Sequential(\n",
              "    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (1): ReLU(inplace=True)\n",
              "    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (3): ReLU(inplace=True)\n",
              "    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (6): ReLU(inplace=True)\n",
              "    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (8): ReLU(inplace=True)\n",
              "    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
              "    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (11): ReLU(inplace=True)\n",
              "    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (13): ReLU(inplace=True)\n",
              "    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (15): ReLU(inplace=True)\n",
              "    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "    (17): ReLU(inplace=True)\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9P6SvahFIFjD"
      },
      "source": [
        "## **Loss Function** and **Optimizer**\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BApcmkC-ILNu"
      },
      "source": [
        "criterion_GAN = torch.nn.MSELoss().cuda()\n",
        "criterion_content = torch.nn.L1Loss().cuda()\n",
        "\n",
        "optimizer_G = torch.optim.Adam(netG.parameters(), lr=lr, betas=(b1, b2))\n",
        "optimizer_D = torch.optim.Adam(netD.parameters(), lr=lr, betas=(b1, b2))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iFP-4ZPUZdf3"
      },
      "source": [
        "# **Training**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 99,
          "referenced_widgets": [
            "1706a71588fd459c912f8415d292f51e",
            "2ad8257f35ee42c9acb50ec37fff7b0b",
            "cccbc200789243b7b5d37240305146a0",
            "086368cc6615494ba48f21a6a8dd523e",
            "79794a7717d14b20b52cb9398d148926",
            "6195e3574c8f49a09ea0e816d62a2d7f",
            "f8d9878fa74a44f3b495d01e2508a8fb",
            "944f531e80f94e19a105063181ea2317",
            "7804d4dff2ac4cccbab7621a23890de4",
            "aaffd65f21504553bd6cd51f20158eab",
            "09bca5c6c53443f8a9d1340707008f63",
            "79b4eb77c12142e3838fafcdebfe9da5",
            "d89e5d7fe52d485ab1289387e1b1ed98",
            "730cdbfff47241ecbc1c7ce4e2a11c6c",
            "b8b84f38186e4361a7c7c0fdd8d241be",
            "45ce6fc8d2694d0792231c1f8716cf9b"
          ]
        },
        "id": "2rKyMzJrZhH5",
        "outputId": "d2bc8721-ca7d-4097-f06a-b41c14507a87"
      },
      "source": [
        "os.makedirs('./results/images/', exist_ok=True)\n",
        "os.makedirs('./results/checkpoints/', exist_ok=True)\n",
        "\n",
        "Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n",
        "\n",
        "train_gen_losses, train_disc_losses, train_counter = [], [], []\n",
        "\n",
        "for epoch in range(n_epochs):\n",
        "  ### Training\n",
        "  gen_loss, disc_loss = 0, 0\n",
        "  netG.train()\n",
        "  netD.train()\n",
        "  tqdm_bar = tqdm(train_dataloader, desc=f'Training Epoch {epoch} ', total=int(len(train_dataloader)))\n",
        "\n",
        "  for batch_idx, imgs in enumerate(tqdm_bar):\n",
        "\n",
        "    # Configure model input\n",
        "    imgs_lr = Variable(imgs[\"lr\"].type(Tensor)).cuda() # 낮은 화질의 이미지\n",
        "    imgs_hr = Variable(imgs[\"hr\"].type(Tensor)).cuda() # 원래 이미지\n",
        "\n",
        "    valid = Variable(Tensor(np.ones((imgs_lr.size(0), *netD.output_shape))), requires_grad=False)\n",
        "    fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *netD.output_shape))), requires_grad=False)\n",
        "        \n",
        "    ########################### Train Generator ################################\n",
        "    optimizer_G.zero_grad()\n",
        "\n",
        "    # Generate a high resolution image from low resolution input\n",
        "    gen_hr = netG(imgs_lr)\n",
        "    \n",
        "    # Adversarial loss\n",
        "    fake_hr = netD(gen_hr) # super resolution\n",
        "    loss_GAN = criterion_GAN(fake_hr, valid)\n",
        "\n",
        "    # Content loss\n",
        "    gen_features = feature_extractor(gen_hr)\n",
        "    real_features = feature_extractor(imgs_hr).detach()\n",
        "    loss_content = criterion_content(gen_features, real_features)\n",
        "\n",
        "    # Total loss\n",
        "    # 기존 GAN loss와 다르게 generator로 부터 생성한 이미지를 HR 이미지로 구별할 확률을 정해줍니다.\n",
        "    # 아래의 식으로 최소화하면 결과가 더 좋다고 합니다..\n",
        "    loss_G = loss_content + 1e-3 * loss_GAN\n",
        "    loss_G.backward()\n",
        "    optimizer_G.step()\n",
        "\n",
        "    ########################### Train Discriminator ############################\n",
        "    optimizer_D.zero_grad()\n",
        "    # Loss of real and fake images\n",
        "    loss_real = criterion_GAN(netD(imgs_hr), valid)\n",
        "    loss_fake = criterion_GAN(netD(gen_hr.detach()), fake)\n",
        "\n",
        "    # Total loss\n",
        "    loss_D = (loss_real + loss_fake) / 2\n",
        "    loss_D.backward()\n",
        "    optimizer_D.step()\n",
        "\n",
        "    gen_loss += loss_G.item()\n",
        "    train_gen_losses.append(loss_G.item())\n",
        "    disc_loss += loss_D.item()\n",
        "    train_disc_losses.append(loss_D.item())\n",
        "\n",
        "    train_counter.append(batch_idx*batch_size + imgs_lr.size(0) + epoch*len(train_dataloader.dataset))\n",
        "    tqdm_bar.set_postfix(gen_loss=gen_loss/(batch_idx+1), disc_loss=disc_loss/(batch_idx+1))\n",
        "\n",
        "    # Test코드는 아직 안했습니다!\n",
        "    \n",
        "    # Save image grid with upsampled inputs and SRGAN outputs\n",
        "    if random.uniform(0,1)<0.1:\n",
        "        imgs_lr = nn.functional.interpolate(imgs_lr, scale_factor=4)\n",
        "        imgs_hr = make_grid(imgs_hr, nrow=1, normalize=True)\n",
        "        gen_hr = make_grid(gen_hr, nrow=1, normalize=True)\n",
        "        imgs_lr = make_grid(imgs_lr, nrow=1, normalize=True)\n",
        "        img_grid = torch.cat((imgs_hr, imgs_lr, gen_hr), -1)\n",
        "        # 한 이미지에 원래 high, low, 만든 high가 들어갑니다!\n",
        "        # batch size만큼 저장해서 한 이미징 16개씩 들어갑니다.\n",
        "        save_image(img_grid, os.path.join('./results/images', 'fake-{:03d}.png'.format(batch_idx)), normalize=False)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1706a71588fd459c912f8415d292f51e",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Training Epoch 0 ', max=12410.0, style=ProgressStyle(desc…"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "7804d4dff2ac4cccbab7621a23890de4",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, description='Training Epoch 1 ', max=12410.0, style=ProgressStyle(desc…"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    }
  ]
}