CartoonGAN_Experiment.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CartoonGAN Experiments.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 935
        },
        "id": "JGExvxRID2DV",
        "outputId": "4c7bf8bd-4bf2-4a50-9f2c-976fed9aad25"
      },
      "source": [
        "!pip install -r '/content/drive/MyDrive/11785 Project/requirements.txt'"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting imageio==2.9.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/6e/57/5d899fae74c1752f52869b613a8210a2480e1a69688e65df6cb26117d45d/imageio-2.9.0-py3-none-any.whl (3.3MB)\n",
            "\u001b[K     |████████████████████████████████| 3.3MB 12.9MB/s \n",
            "\u001b[?25hCollecting numpy==1.20.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/73/ef/8967d406f3f85018ceb5efab50431e901683188f1741ceb053efcab26c87/numpy-1.20.2-cp37-cp37m-manylinux2010_x86_64.whl (15.3MB)\n",
            "\u001b[K     |████████████████████████████████| 15.3MB 344kB/s \n",
            "\u001b[?25hCollecting opencv-python==4.5.1.48\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/0f/13/192104516c4a3d92dc6b5e106ffcfbf0fe35f3c4faa49650205ff652af72/opencv_python-4.5.1.48-cp37-cp37m-manylinux2014_x86_64.whl (50.4MB)\n",
            "\u001b[K     |████████████████████████████████| 50.4MB 113kB/s \n",
            "\u001b[?25hCollecting Pillow==8.2.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/33/34/542152297dcc6c47a9dcb0685eac6d652d878ed3cea83bf2b23cb988e857/Pillow-8.2.0-cp37-cp37m-manylinux1_x86_64.whl (3.0MB)\n",
            "\u001b[K     |████████████████████████████████| 3.0MB 40.0MB/s \n",
            "\u001b[?25hCollecting scipy==1.6.2\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/75/91/ee427c42957f8c4cbe477bf4f8b7f608e003a17941e509d1777e58648cb3/scipy-1.6.2-cp37-cp37m-manylinux1_x86_64.whl (27.4MB)\n",
            "\u001b[K     |████████████████████████████████| 27.4MB 94kB/s \n",
            "\u001b[?25hRequirement already satisfied: torch==1.8.1 in /usr/local/lib/python3.7/dist-packages (from -r /content/drive/MyDrive/11785 Project/requirements.txt (line 6)) (1.8.1+cu101)\n",
            "Requirement already satisfied: torchvision==0.9.1 in /usr/local/lib/python3.7/dist-packages (from -r /content/drive/MyDrive/11785 Project/requirements.txt (line 7)) (0.9.1+cu101)\n",
            "Collecting tqdm==4.59.0\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f8/3e/2730d0effc282960dbff3cf91599ad0d8f3faedc8e75720fdf224b31ab24/tqdm-4.59.0-py2.py3-none-any.whl (74kB)\n",
            "\u001b[K     |████████████████████████████████| 81kB 11.5MB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions==3.7.4.3 in /usr/local/lib/python3.7/dist-packages (from -r /content/drive/MyDrive/11785 Project/requirements.txt (line 9)) (3.7.4.3)\n",
            "\u001b[31mERROR: tensorflow 2.4.1 has requirement numpy~=1.19.2, but you'll have numpy 1.20.2 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: datascience 0.10.6 has requirement folium==0.2.1, but you'll have folium 0.8.3 which is incompatible.\u001b[0m\n",
            "\u001b[31mERROR: albumentations 0.1.12 has requirement imgaug<0.2.7,>=0.2.5, but you'll have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Installing collected packages: Pillow, numpy, imageio, opencv-python, scipy, tqdm\n",
            "  Found existing installation: Pillow 7.1.2\n",
            "    Uninstalling Pillow-7.1.2:\n",
            "      Successfully uninstalled Pillow-7.1.2\n",
            "  Found existing installation: numpy 1.19.5\n",
            "    Uninstalling numpy-1.19.5:\n",
            "      Successfully uninstalled numpy-1.19.5\n",
            "  Found existing installation: imageio 2.4.1\n",
            "    Uninstalling imageio-2.4.1:\n",
            "      Successfully uninstalled imageio-2.4.1\n",
            "  Found existing installation: opencv-python 4.1.2.30\n",
            "    Uninstalling opencv-python-4.1.2.30:\n",
            "      Successfully uninstalled opencv-python-4.1.2.30\n",
            "  Found existing installation: scipy 1.4.1\n",
            "    Uninstalling scipy-1.4.1:\n",
            "      Successfully uninstalled scipy-1.4.1\n",
            "  Found existing installation: tqdm 4.41.1\n",
            "    Uninstalling tqdm-4.41.1:\n",
            "      Successfully uninstalled tqdm-4.41.1\n",
            "Successfully installed Pillow-8.2.0 imageio-2.9.0 numpy-1.20.2 opencv-python-4.5.1.48 scipy-1.6.2 tqdm-4.59.0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "PIL",
                  "numpy"
                ]
              }
            }
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-91BmVIGgOos",
        "outputId": "f2240787-31db-40d1-899b-c34ac6022142"
      },
      "source": [
        "!git clone https://[GITHUB_TOKEN]@github.com/hans-lizihan/11785-project.git"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into '11785-project'...\n",
            "remote: Enumerating objects: 62682, done.\u001b[K\n",
            "remote: Counting objects: 100% (10267/10267), done.\u001b[K\n",
            "remote: Compressing objects: 100% (10238/10238), done.\u001b[K\n",
            "remote: Total 62682 (delta 22), reused 10267 (delta 22), pack-reused 52415\u001b[K\n",
            "Receiving objects: 100% (62682/62682), 2.24 GiB | 11.95 MiB/s, done.\n",
            "Resolving deltas: 100% (126/126), done.\n",
            "Checking out files: 100% (79926/79926), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "faLGrCyb2cHL",
        "outputId": "3a0ca668-9558-4219-cfa7-a3b8c7081f94"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oMeRfBPhFOLD"
      },
      "source": [
        "import imageio\n",
        "import itertools\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torchvision import datasets, transforms\n",
        "from torch.autograd import Variable\n",
        "import cv2, os\n",
        "from tqdm import tqdm\n",
        "import os, time, pickle, argparse\n",
        "import torch.optim as optim\n",
        "from google.colab.patches import cv2_imshow\n",
        "from pathlib import Path"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-X3isBewUVaE"
      },
      "source": [
        "def data_load(path, subfolder, transform, batch_size, shuffle=False, drop_last=True):\n",
        "    dset = datasets.ImageFolder(path, transform)\n",
        "    ind = dset.class_to_idx[subfolder]\n",
        "\n",
        "    n = 0\n",
        "    for i in range(dset.__len__()):\n",
        "        if ind != dset.imgs[n][1]:\n",
        "            del dset.imgs[n]\n",
        "            n -= 1\n",
        "\n",
        "        n += 1\n",
        "\n",
        "    return torch.utils.data.DataLoader(dset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)\n",
        "\n",
        "def print_network(net):\n",
        "    num_params = 0\n",
        "    for param in net.parameters():\n",
        "        num_params += param.numel()\n",
        "    print(net)\n",
        "    print('Total number of parameters: %d' % num_params)\n",
        "\n",
        "def initialize_weights(net):\n",
        "    for m in net.modules():\n",
        "        if isinstance(m, nn.Conv2d):\n",
        "            m.weight.data.normal_(0, 0.02)\n",
        "            m.bias.data.zero_()\n",
        "        elif isinstance(m, nn.ConvTranspose2d):\n",
        "            m.weight.data.normal_(0, 0.02)\n",
        "            m.bias.data.zero_()\n",
        "        elif isinstance(m, nn.Linear):\n",
        "            m.weight.data.normal_(0, 0.02)\n",
        "            m.bias.data.zero_()\n",
        "        elif isinstance(m, nn.BatchNorm2d):\n",
        "            m.weight.data.fill_(1)\n",
        "            m.bias.data.zero_()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3Q2PqsnwUOMb"
      },
      "source": [
        "class resnet_block(nn.Module):\n",
        "    def __init__(self, channel, kernel, stride, padding):\n",
        "        super(resnet_block, self).__init__()\n",
        "        self.channel = channel\n",
        "        self.kernel = kernel\n",
        "        self.strdie = stride\n",
        "        self.padding = padding\n",
        "        self.conv1 = nn.Conv2d(channel, channel, kernel, stride, padding)\n",
        "        self.conv1_norm = nn.InstanceNorm2d(channel)\n",
        "        self.conv2 = nn.Conv2d(channel, channel, kernel, stride, padding)\n",
        "        self.conv2_norm = nn.InstanceNorm2d(channel)\n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input):\n",
        "        x = F.relu(self.conv1_norm(self.conv1(input)), True)\n",
        "        x = self.conv2_norm(self.conv2(x))\n",
        "\n",
        "        return input + x #Elementwise Sum\n",
        " \n",
        "\n",
        "class generator(nn.Module):\n",
        "    # initializers\n",
        "    def __init__(self, in_nc, out_nc, nf=32, nb=6):\n",
        "        super(generator, self).__init__()\n",
        "        self.input_nc = in_nc\n",
        "        self.output_nc = out_nc\n",
        "        self.nf = nf\n",
        "        self.nb = nb\n",
        "        self.down_convs = nn.Sequential(\n",
        "            nn.Conv2d(in_nc, nf, 7, 1, 3), #k7n64s1\n",
        "            nn.InstanceNorm2d(nf),\n",
        "            nn.ReLU(True),\n",
        "            nn.Conv2d(nf, nf * 2, 3, 2, 1), #k3n128s2\n",
        "            nn.Conv2d(nf * 2, nf * 2, 3, 1, 1), #k3n128s1\n",
        "            nn.InstanceNorm2d(nf * 2),\n",
        "            nn.ReLU(True),\n",
        "            nn.Conv2d(nf * 2, nf * 4, 3, 2, 1), #k3n256s1\n",
        "            nn.Conv2d(nf * 4, nf * 4, 3, 1, 1), #k3n256s1\n",
        "            nn.InstanceNorm2d(nf * 4),\n",
        "            nn.ReLU(True),\n",
        "        )\n",
        "\n",
        "        self.resnet_blocks = []\n",
        "        for i in range(nb):\n",
        "            self.resnet_blocks.append(resnet_block(nf * 4, 3, 1, 1))\n",
        "\n",
        "        self.resnet_blocks = nn.Sequential(*self.resnet_blocks)\n",
        "\n",
        "        self.up_convs = nn.Sequential(\n",
        "            nn.ConvTranspose2d(nf * 4, nf * 2, 3, 2, 1, 1), #k3n128s1/2\n",
        "            nn.Conv2d(nf * 2, nf * 2, 3, 1, 1), #k3n128s1\n",
        "            nn.InstanceNorm2d(nf * 2),\n",
        "            nn.ReLU(True),\n",
        "            nn.ConvTranspose2d(nf * 2, nf, 3, 2, 1, 1), #k3n64s1/2\n",
        "            nn.Conv2d(nf, nf, 3, 1, 1), #k3n64s1\n",
        "            nn.InstanceNorm2d(nf),\n",
        "            nn.ReLU(True),\n",
        "            nn.Conv2d(nf, out_nc, 7, 1, 3), #k7n3s1\n",
        "            nn.Tanh(),\n",
        "        )\n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    # forward method\n",
        "    def forward(self, input):\n",
        "        x = self.down_convs(input)\n",
        "        x = self.resnet_blocks(x)\n",
        "        output = self.up_convs(x)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class discriminator(nn.Module):\n",
        "    # initializers\n",
        "    def __init__(self, in_nc, out_nc, nf=32):\n",
        "        super(discriminator, self).__init__()\n",
        "        self.input_nc = in_nc\n",
        "        self.output_nc = out_nc\n",
        "        self.nf = nf\n",
        "        self.convs = nn.Sequential(\n",
        "            nn.Conv2d(in_nc, nf, 3, 1, 1),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf, nf * 2, 3, 2, 1),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf * 2, nf * 4, 3, 1, 1),\n",
        "            nn.InstanceNorm2d(nf * 4),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf * 4, nf * 4, 3, 2, 1),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf * 4, nf * 8, 3, 1, 1),\n",
        "            nn.InstanceNorm2d(nf * 8),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf * 8, nf * 8, 3, 1, 1),\n",
        "            nn.InstanceNorm2d(nf * 8),\n",
        "            nn.LeakyReLU(0.2, True),\n",
        "            nn.Conv2d(nf * 8, out_nc, 3, 1, 1),\n",
        "            nn.Sigmoid(),\n",
        "        )\n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    # forward method\n",
        "    def forward(self, input):\n",
        "        # input = torch.cat((input1, input2), 1)\n",
        "        output = self.convs(input)\n",
        "\n",
        "        return output\n",
        "\n",
        "\n",
        "class VGG19(nn.Module):\n",
        "    def __init__(self, init_weights=None, feature_mode=False, batch_norm=False, num_classes=1000):\n",
        "        super(VGG19, self).__init__()\n",
        "        self.cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M']\n",
        "        self.init_weights = init_weights\n",
        "        self.feature_mode = feature_mode\n",
        "        self.batch_norm = batch_norm\n",
        "        self.num_clases = num_classes\n",
        "        self.features = self.make_layers(self.cfg, batch_norm)\n",
        "        self.classifier = nn.Sequential(\n",
        "            nn.Linear(512 * 7 * 7, 4096),\n",
        "            nn.ReLU(True),\n",
        "            nn.Dropout(),\n",
        "            nn.Linear(4096, 4096),\n",
        "            nn.ReLU(True),\n",
        "            nn.Dropout(),\n",
        "            nn.Linear(4096, num_classes),\n",
        "        )\n",
        "        if not init_weights == None:\n",
        "            self.load_state_dict(torch.load(init_weights))\n",
        "\n",
        "    def make_layers(self, cfg, batch_norm=False):\n",
        "        layers = []\n",
        "        in_channels = 3\n",
        "        for v in cfg:\n",
        "            if v == 'M':\n",
        "                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n",
        "            else:\n",
        "                conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)\n",
        "                if batch_norm:\n",
        "                    layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]\n",
        "                else:\n",
        "                    layers += [conv2d, nn.ReLU(inplace=True)]\n",
        "                in_channels = v\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.feature_mode:\n",
        "            module_list = list(self.features.modules())\n",
        "            for l in module_list[1:27]:                 # conv4_4\n",
        "                x = l(x)\n",
        "        if not self.feature_mode:\n",
        "            x = x.view(x.size(0), -1)\n",
        "            x = self.classifier(x)\n",
        "\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JgSiOhn77Udl"
      },
      "source": [
        "def edge_promoting_v1(root, save):\n",
        "    file_list = os.listdir(root)\n",
        "    if not os.path.isdir(save):\n",
        "        os.makedirs(save)\n",
        "    kernel_size = 5\n",
        "    kernel = np.ones((kernel_size, kernel_size), np.uint8)\n",
        "    gauss = cv2.getGaussianKernel(kernel_size, 0)\n",
        "    gauss = gauss * gauss.transpose(1, 0)\n",
        "    n = 1\n",
        "    for f in tqdm(file_list):\n",
        "        rgb_img = cv2.imread(os.path.join(root, f))\n",
        "        gray_img = cv2.imread(os.path.join(root, f), 0)\n",
        "        rgb_img = cv2.resize(rgb_img, (256, 256))\n",
        "        pad_img = np.pad(rgb_img, ((2,2), (2,2), (0,0)), mode='reflect')\n",
        "        gray_img = cv2.resize(gray_img, (256, 256))\n",
        "        edges = cv2.Canny(gray_img, 100, 200)\n",
        "        dilation = cv2.dilate(edges, kernel)\n",
        "\n",
        "        gauss_img = np.copy(rgb_img)\n",
        "        idx = np.where(dilation != 0)\n",
        "        for i in range(np.sum(dilation != 0)):\n",
        "            gauss_img[idx[0][i], idx[1][i], 0] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 0], gauss))\n",
        "            gauss_img[idx[0][i], idx[1][i], 1] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 1], gauss))\n",
        "            gauss_img[idx[0][i], idx[1][i], 2] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 2], gauss))\n",
        "\n",
        "        result = np.concatenate((rgb_img, gauss_img), 1)\n",
        "\n",
        "        cv2.imwrite(os.path.join(save, str(n) + '.png'), result)\n",
        "        n += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7yhhYzcURHfh"
      },
      "source": [
        "def to_picture(argv):\n",
        "\n",
        "  parser = argparse.ArgumentParser()\n",
        "  parser.add_argument('--name', required=False, default='project_name',  help='')\n",
        "  parser.add_argument('--src_data', required=False, default='src_data_path',  help='sec data path')\n",
        "  parser.add_argument('--tgt_data', required=False, default='tgt_data_path',  help='tgt data path')\n",
        "  parser.add_argument('--vgg_model', required=False, default='pre_trained_VGG19_model_path/vgg19.pth', help='pre-trained VGG19 model path')\n",
        "  parser.add_argument('--in_ngc', type=int, default=3, help='input channel for generator')\n",
        "  parser.add_argument('--out_ngc', type=int, default=3, help='output channel for generator')\n",
        "  parser.add_argument('--in_ndc', type=int, default=3, help='input channel for discriminator')\n",
        "  parser.add_argument('--out_ndc', type=int, default=1, help='output channel for discriminator')\n",
        "  parser.add_argument('--batch_size', type=int, default=8, help='batch size')\n",
        "  parser.add_argument('--ngf', type=int, default=64)\n",
        "  parser.add_argument('--ndf', type=int, default=32)\n",
        "  parser.add_argument('--nb', type=int, default=8, help='the number of resnet block layer for generator')\n",
        "  parser.add_argument('--input_size', type=int, default=256, help='input size')\n",
        "  parser.add_argument('--train_epoch', type=int, default=100)\n",
        "  parser.add_argument('--pre_train_epoch', type=int, default=20)\n",
        "  parser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002')\n",
        "  parser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002')\n",
        "  parser.add_argument('--con_lambda', type=float, default=15, help='lambda for content loss')\n",
        "  parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for Adam optimizer')\n",
        "  parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')\n",
        "  parser.add_argument('--latest_generator_model', required=False, default='', help='the latest trained model path')\n",
        "  parser.add_argument('--latest_discriminator_model', required=False, default='', help='the latest trained model path')\n",
        "  args = parser.parse_args(argv)\n",
        "\n",
        "  print('------------ Options -------------')\n",
        "  for k, v in sorted(vars(args).items()):\n",
        "      print('%s: %s' % (str(k), str(v)))\n",
        "  print('-------------- End ----------------')\n",
        "\n",
        "  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "  if torch.backends.cudnn.enabled:\n",
        "      torch.backends.cudnn.benchmark = True\n",
        "\n",
        "  # results save path\n",
        "  if not os.path.isdir(os.path.join(args.name + '_to_picture_results', 'Reconstruction')):\n",
        "      os.makedirs(os.path.join(args.name + '_to_picture_results', 'Reconstruction'))\n",
        "  if not os.path.isdir(os.path.join(args.name + '_to_picture_results', 'Transfer')):\n",
        "      os.makedirs(os.path.join(args.name + '_to_picture_results', 'Transfer'))\n",
        "\n",
        "  # # edge-promoting\n",
        "  # if not os.path.isdir(os.path.join('data', args.tgt_data, 'pair')):\n",
        "  #     print('edge-promoting start!!')\n",
        "  #     edge_promoting(os.path.join('data', args.tgt_data), os.path.join('data', args.tgt_data, 'pair'))\n",
        "  # else:\n",
        "  #     print('edge-promoting already done')\n",
        "\n",
        "  src_transform = transforms.Compose([\n",
        "          transforms.Resize((args.input_size, args.input_size)),\n",
        "          transforms.ToTensor(),\n",
        "          # transforms.GaussianBlur(kernel_size= 5, sigma=0.1),\n",
        "          transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
        "  ])\n",
        "  tgt_transform = transforms.Compose([\n",
        "          transforms.ToTensor(),\n",
        "          transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
        "  ])\n",
        "  train_loader_src = data_load(os.path.join('data', args.src_data), 'train', src_transform, args.batch_size, shuffle=True, drop_last=True)\n",
        "  train_loader_tgt = data_load(os.path.join('data', args.tgt_data), 'train', tgt_transform, args.batch_size, shuffle=True, drop_last=True)\n",
        "  test_loader_src = data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True)\n",
        "\n",
        "  # network\n",
        "  G = generator(args.in_ngc, args.out_ngc, args.ngf, args.nb)\n",
        "  if args.latest_generator_model != '':\n",
        "      if torch.cuda.is_available():\n",
        "          G.load_state_dict(torch.load(args.latest_generator_model))\n",
        "      else:\n",
        "          # cpu mode\n",
        "          G.load_state_dict(torch.load(args.latest_generator_model, map_location=lambda storage, loc: storage))\n",
        "\n",
        "  D = discriminator(args.in_ndc, args.out_ndc, args.ndf)\n",
        "  if args.latest_discriminator_model != '':\n",
        "      if torch.cuda.is_available():\n",
        "          D.load_state_dict(torch.load(args.latest_discriminator_model))\n",
        "      else:\n",
        "          D.load_state_dict(torch.load(args.latest_discriminator_model, map_location=lambda storage, loc: storage))\n",
        "  VGG = VGG19(init_weights=args.vgg_model, feature_mode=True)\n",
        "  G.to(device)\n",
        "  D.to(device)\n",
        "  VGG.to(device)\n",
        "  G.train()\n",
        "  D.train()\n",
        "  VGG.eval()\n",
        "  print('---------- Networks initialized -------------')\n",
        "  print_network(G)\n",
        "  print_network(D)\n",
        "  print_network(VGG)\n",
        "  print('-----------------------------------------------')\n",
        "\n",
        "  # loss\n",
        "  BCE_loss = nn.BCELoss().to(device)\n",
        "  L1_loss = nn.L1Loss().to(device)\n",
        "\n",
        "  # Adam optimizer\n",
        "  G_optimizer = optim.Adam(G.parameters(), lr=args.lrG, betas=(args.beta1, args.beta2))\n",
        "  D_optimizer = optim.Adam(D.parameters(), lr=args.lrD, betas=(args.beta1, args.beta2))\n",
        "  G_scheduler = optim.lr_scheduler.MultiStepLR(optimizer=G_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1)\n",
        "  D_scheduler = optim.lr_scheduler.MultiStepLR(optimizer=D_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1)\n",
        "\n",
        "  pre_train_hist = {}\n",
        "  pre_train_hist['Recon_loss'] = []\n",
        "  pre_train_hist['per_epoch_time'] = []\n",
        "  pre_train_hist['total_time'] = []\n",
        "\n",
        "  \"\"\" Pre-train reconstruction \"\"\"\n",
        "  if args.latest_generator_model == '':\n",
        "      print('Pre-training start!')\n",
        "      start_time = time.time()\n",
        "      for epoch in range(args.pre_train_epoch):\n",
        "          epoch_start_time = time.time()\n",
        "          Recon_losses = []\n",
        "          for x, _ in train_loader_src:\n",
        "              x = x.to(device)\n",
        "\n",
        "              # train generator G\n",
        "              G_optimizer.zero_grad()\n",
        "\n",
        "              x_feature = VGG((x + 1) / 2)\n",
        "              G_ = G(x)\n",
        "              G_feature = VGG((G_ + 1) / 2)\n",
        "\n",
        "              Recon_loss = 10 * L1_loss(G_feature, x_feature.detach())\n",
        "              Recon_losses.append(Recon_loss.item())\n",
        "              pre_train_hist['Recon_loss'].append(Recon_loss.item())\n",
        "\n",
        "              Recon_loss.backward()\n",
        "              G_optimizer.step()\n",
        "\n",
        "          per_epoch_time = time.time() - epoch_start_time\n",
        "          pre_train_hist['per_epoch_time'].append(per_epoch_time)\n",
        "          print('[%d/%d] - time: %.2f, Recon loss: %.3f' % ((epoch + 1), args.pre_train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Recon_losses))))\n",
        "\n",
        "      total_time = time.time() - start_time\n",
        "      pre_train_hist['total_time'].append(total_time)\n",
        "      with open(os.path.join(args.name + '_to_picture_results',  'pre_train_hist.pkl'), 'wb') as f:\n",
        "          pickle.dump(pre_train_hist, f)\n",
        "\n",
        "      with torch.no_grad():\n",
        "          G.eval()\n",
        "          for n, (x, _) in enumerate(train_loader_src):\n",
        "              x = x.to(device)\n",
        "              G_recon = G(x)\n",
        "              result = torch.cat((x[0], G_recon[0]), 2)\n",
        "              path = os.path.join(args.name + '_to_picture_results', 'Reconstruction', args.name + '_train_recon_' + str(n + 1) + '.png')\n",
        "              plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "              if n == 4:\n",
        "                  break\n",
        "\n",
        "          for n, (x, _) in enumerate(test_loader_src):\n",
        "              x = x.to(device)\n",
        "              G_recon = G(x)\n",
        "              result = torch.cat((x[0], G_recon[0]), 2)\n",
        "              path = os.path.join(args.name + '_to_picture_results', 'Reconstruction', args.name + '_test_recon_' + str(n + 1) + '.png')\n",
        "              plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "              if n == 4:\n",
        "                  break\n",
        "  else:\n",
        "      print('Load the latest generator model, no need to pre-train')\n",
        "\n",
        "\n",
        "  train_hist = {}\n",
        "  train_hist['Disc_loss'] = []\n",
        "  train_hist['Gen_loss'] = []\n",
        "  train_hist['Con_loss'] = []\n",
        "  train_hist['per_epoch_time'] = []\n",
        "  train_hist['total_time'] = []\n",
        "  print('training start!')\n",
        "  start_time = time.time()\n",
        "  real = torch.ones(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device)\n",
        "  fake = torch.zeros(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device)\n",
        "  for epoch in range(args.train_epoch):\n",
        "      epoch_start_time = time.time()\n",
        "      G.train()\n",
        "      G_scheduler.step()\n",
        "      D_scheduler.step()\n",
        "      Disc_losses = []\n",
        "      Gen_losses = []\n",
        "      Con_losses = []\n",
        "      for (x, _), (y, _) in zip(train_loader_src, train_loader_tgt):\n",
        "          # e = y[:, :, :, args.input_size:]\n",
        "          y = y[:, :, :, :args.input_size]\n",
        "          x, y = x.to(device), y.to(device)\n",
        "\n",
        "          # train D\n",
        "          D_optimizer.zero_grad()\n",
        "\n",
        "          D_real = D(y)\n",
        "          D_real_loss = BCE_loss(D_real, real)\n",
        "\n",
        "          G_ = G(x)\n",
        "          D_fake = D(G_)\n",
        "          D_fake_loss = BCE_loss(D_fake, fake)\n",
        "\n",
        "          # D_edge = D(e)\n",
        "          # D_edge_loss = BCE_loss(D_edge, fake)\n",
        "\n",
        "          Disc_loss = D_real_loss + D_fake_loss\n",
        "          Disc_losses.append(Disc_loss.item())\n",
        "          train_hist['Disc_loss'].append(Disc_loss.item())\n",
        "\n",
        "          Disc_loss.backward()\n",
        "          D_optimizer.step()\n",
        "\n",
        "          # train G\n",
        "          G_optimizer.zero_grad()\n",
        "\n",
        "          G_ = G(x)\n",
        "          D_fake = D(G_)\n",
        "          D_fake_loss = BCE_loss(D_fake, real)\n",
        "\n",
        "          x_feature = VGG((x + 1) / 2)\n",
        "          G_feature = VGG((G_ + 1) / 2)\n",
        "          Con_loss = args.con_lambda * L1_loss(G_feature, x_feature.detach())\n",
        "\n",
        "          Gen_loss = D_fake_loss + Con_loss\n",
        "          Gen_losses.append(D_fake_loss.item())\n",
        "          train_hist['Gen_loss'].append(D_fake_loss.item())\n",
        "          Con_losses.append(Con_loss.item())\n",
        "          train_hist['Con_loss'].append(Con_loss.item())\n",
        "\n",
        "          Gen_loss.backward()\n",
        "          G_optimizer.step()\n",
        "\n",
        "\n",
        "      per_epoch_time = time.time() - epoch_start_time\n",
        "      train_hist['per_epoch_time'].append(per_epoch_time)\n",
        "      print(\n",
        "      '[%d/%d] - time: %.2f, Disc loss: %.3f, Gen loss: %.3f, Con loss: %.3f' % ((epoch + 1), args.train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Disc_losses)),\n",
        "          torch.mean(torch.FloatTensor(Gen_losses)), torch.mean(torch.FloatTensor(Con_losses))))\n",
        "\n",
        "      if epoch % 2 == 1 or epoch == args.train_epoch - 1:\n",
        "          with torch.no_grad():\n",
        "              G.eval()\n",
        "              for n, (x, _) in enumerate(train_loader_src):\n",
        "                  x = x.to(device)\n",
        "                  G_recon = G(x)\n",
        "                  result = torch.cat((x[0], G_recon[0]), 2)\n",
        "                  path = os.path.join(args.name + '_to_picture_results', 'Transfer', str(epoch+1) + '_epoch_' + args.name + '_train_' + str(n + 1) + '.png')\n",
        "                  plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "                  if n == 4:\n",
        "                      break\n",
        "\n",
        "              for n, (x, _) in enumerate(test_loader_src):\n",
        "                  x = x.to(device)\n",
        "                  G_recon = G(x)\n",
        "                  result = torch.cat((x[0], G_recon[0]), 2)\n",
        "                  path = os.path.join(args.name + '_to_picture_results', 'Transfer', str(epoch+1) + '_epoch_' + args.name + '_test_' + str(n + 1) + '.png')\n",
        "                  plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "                  if n == 4:\n",
        "                      break\n",
        "\n",
        "              torch.save(G.state_dict(), os.path.join(args.name + '_to_picture_results', 'generator_latest.pkl'))\n",
        "              torch.save(D.state_dict(), os.path.join(args.name + '_to_picture_results', 'discriminator_latest.pkl'))\n",
        "\n",
        "  total_time = time.time() - start_time\n",
        "  train_hist['total_time'].append(total_time)\n",
        "\n",
        "  print(\"Avg one epoch time: %.2f, total %d epochs time: %.2f\" % (torch.mean(torch.FloatTensor(train_hist['per_epoch_time'])), args.train_epoch, total_time))\n",
        "  print(\"Training finish!... save training results\")\n",
        "\n",
        "  torch.save(G.state_dict(), os.path.join(args.name + '_to_picture_results',  'generator_param.pkl'))\n",
        "  torch.save(D.state_dict(), os.path.join(args.name + '_to_picture_results',  'discriminator_param.pkl'))\n",
        "  with open(os.path.join(args.name + '_to_picture_results',  'train_hist.pkl'), 'wb') as f:\n",
        "      pickle.dump(train_hist, f)\n",
        "  return G"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fUec_d3ATA1j",
        "outputId": "11080763-3ffb-42d3-f8e5-ad071bf82a38"
      },
      "source": [
        "picture_G = to_picture([\n",
        "    \"--name\", \"11785\",\n",
        "    \"--src_data\", \"/content/drive/MyDrive/11785 Project/src_shinkai Makoto\",  \n",
        "    \"--tgt_data\", \"/content/drive/MyDrive/11785 Project/src_all\",\n",
        "    \"--vgg_model\",\"/content/drive/MyDrive/11785 Project/vgg19-dcbb9e9d.pth\",\n",
        "    \"--train_epoch\", \"100\"])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "------------ Options -------------\n",
            "batch_size: 8\n",
            "beta1: 0.5\n",
            "beta2: 0.999\n",
            "con_lambda: 15\n",
            "in_ndc: 3\n",
            "in_ngc: 3\n",
            "input_size: 256\n",
            "latest_discriminator_model: \n",
            "latest_generator_model: \n",
            "lrD: 0.0002\n",
            "lrG: 0.0002\n",
            "name: 11785\n",
            "nb: 8\n",
            "ndf: 32\n",
            "ngf: 64\n",
            "out_ndc: 1\n",
            "out_ngc: 3\n",
            "pre_train_epoch: 20\n",
            "src_data: /content/drive/MyDrive/11785 Project/src_shinkai Makoto\n",
            "tgt_data: /content/drive/MyDrive/11785 Project/src_all\n",
            "train_epoch: 100\n",
            "vgg_model: /content/drive/MyDrive/11785 Project/vgg19-dcbb9e9d.pth\n",
            "-------------- End ----------------\n",
            "---------- Networks initialized -------------\n",
            "generator(\n",
            "  (down_convs): Sequential(\n",
            "    (0): Conv2d(3, 64, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3))\n",
            "    (1): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (2): ReLU(inplace=True)\n",
            "    (3): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
            "    (4): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (6): ReLU(inplace=True)\n",
            "    (7): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
            "    (8): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (9): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (10): ReLU(inplace=True)\n",
            "  )\n",
            "  (resnet_blocks): Sequential(\n",
            "    (0): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (1): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (2): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (3): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (4): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (5): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (6): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "    (7): resnet_block(\n",
            "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv1_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "      (conv2_norm): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    )\n",
            "  )\n",
            "  (up_convs): Sequential(\n",
            "    (0): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
            "    (1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (2): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (3): ReLU(inplace=True)\n",
            "    (4): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
            "    (5): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (6): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (7): ReLU(inplace=True)\n",
            "    (8): Conv2d(64, 3, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3))\n",
            "    (9): Tanh()\n",
            "  )\n",
            ")\n",
            "Total number of parameters: 11120195\n",
            "discriminator(\n",
            "  (convs): Sequential(\n",
            "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (1): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (2): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
            "    (3): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (4): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (6): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
            "    (8): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (9): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (10): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (11): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (13): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
            "    (14): LeakyReLU(negative_slope=0.2, inplace=True)\n",
            "    (15): Conv2d(256, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (16): Sigmoid()\n",
            "  )\n",
            ")\n",
            "Total number of parameters: 1128385\n",
            "VGG19(\n",
            "  (features): 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",
            "    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            "    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (20): ReLU(inplace=True)\n",
            "    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (22): ReLU(inplace=True)\n",
            "    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (24): ReLU(inplace=True)\n",
            "    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (26): ReLU(inplace=True)\n",
            "    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            "    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (29): ReLU(inplace=True)\n",
            "    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (31): ReLU(inplace=True)\n",
            "    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (33): ReLU(inplace=True)\n",
            "    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
            "    (35): ReLU(inplace=True)\n",
            "    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
            "  )\n",
            "  (classifier): Sequential(\n",
            "    (0): Linear(in_features=25088, out_features=4096, bias=True)\n",
            "    (1): ReLU(inplace=True)\n",
            "    (2): Dropout(p=0.5, inplace=False)\n",
            "    (3): Linear(in_features=4096, out_features=4096, bias=True)\n",
            "    (4): ReLU(inplace=True)\n",
            "    (5): Dropout(p=0.5, inplace=False)\n",
            "    (6): Linear(in_features=4096, out_features=1000, bias=True)\n",
            "  )\n",
            ")\n",
            "Total number of parameters: 143667240\n",
            "-----------------------------------------------\n",
            "Pre-training start!\n",
            "[1/20] - time: 258.95, Recon loss: 14.323\n",
            "[2/20] - time: 257.49, Recon loss: 7.591\n",
            "[3/20] - time: 257.51, Recon loss: 5.625\n",
            "[4/20] - time: 257.51, Recon loss: 4.700\n",
            "[5/20] - time: 257.49, Recon loss: 4.182\n",
            "[6/20] - time: 257.52, Recon loss: 3.794\n",
            "[7/20] - time: 257.49, Recon loss: 3.589\n",
            "[8/20] - time: 257.48, Recon loss: 3.333\n",
            "[9/20] - time: 257.52, Recon loss: 3.187\n",
            "[10/20] - time: 257.50, Recon loss: 3.079\n",
            "[11/20] - time: 257.46, Recon loss: 2.965\n",
            "[12/20] - time: 257.48, Recon loss: 2.858\n",
            "[13/20] - time: 257.49, Recon loss: 2.758\n",
            "[14/20] - time: 257.49, Recon loss: 2.694\n",
            "[15/20] - time: 257.48, Recon loss: 2.632\n",
            "[16/20] - time: 257.48, Recon loss: 2.540\n",
            "[17/20] - time: 257.48, Recon loss: 2.516\n",
            "[18/20] - time: 257.49, Recon loss: 2.437\n",
            "[19/20] - time: 257.50, Recon loss: 2.400\n",
            "[20/20] - time: 257.49, Recon loss: 2.360\n",
            "training start!\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/optim/lr_scheduler.py:134: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
            "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "[1/100] - time: 504.20, Disc loss: 0.335, Gen loss: 3.038, Con loss: 4.594\n",
            "[2/100] - time: 504.20, Disc loss: 0.016, Gen loss: 5.457, Con loss: 3.982\n",
            "[3/100] - time: 503.96, Disc loss: 0.431, Gen loss: 4.230, Con loss: 3.790\n",
            "[4/100] - time: 504.17, Disc loss: 0.704, Gen loss: 2.063, Con loss: 3.933\n",
            "[5/100] - time: 503.96, Disc loss: 0.512, Gen loss: 2.748, Con loss: 4.347\n",
            "[6/100] - time: 504.14, Disc loss: 0.017, Gen loss: 5.475, Con loss: 3.801\n",
            "[7/100] - time: 504.02, Disc loss: 0.006, Gen loss: 6.376, Con loss: 3.448\n",
            "[8/100] - time: 504.33, Disc loss: 0.004, Gen loss: 6.655, Con loss: 3.605\n",
            "[9/100] - time: 504.19, Disc loss: 0.003, Gen loss: 7.098, Con loss: 3.605\n",
            "[10/100] - time: 504.52, Disc loss: 0.002, Gen loss: 7.552, Con loss: 3.544\n",
            "[11/100] - time: 503.93, Disc loss: 0.001, Gen loss: 7.966, Con loss: 3.575\n",
            "[12/100] - time: 504.10, Disc loss: 0.074, Gen loss: 6.705, Con loss: 3.518\n",
            "[13/100] - time: 503.90, Disc loss: 0.073, Gen loss: 6.635, Con loss: 4.252\n",
            "[14/100] - time: 504.12, Disc loss: 0.006, Gen loss: 6.575, Con loss: 4.470\n",
            "[15/100] - time: 503.95, Disc loss: 0.003, Gen loss: 7.428, Con loss: 4.053\n",
            "[16/100] - time: 504.17, Disc loss: 0.001, Gen loss: 8.238, Con loss: 4.189\n",
            "[17/100] - time: 503.92, Disc loss: 0.001, Gen loss: 8.760, Con loss: 3.567\n",
            "[18/100] - time: 504.11, Disc loss: 0.001, Gen loss: 8.931, Con loss: 3.852\n",
            "[19/100] - time: 503.89, Disc loss: 0.006, Gen loss: 8.154, Con loss: 3.882\n",
            "[20/100] - time: 504.46, Disc loss: 0.001, Gen loss: 9.216, Con loss: 3.760\n",
            "[21/100] - time: 503.92, Disc loss: 0.000, Gen loss: 9.654, Con loss: 4.121\n",
            "[22/100] - time: 504.05, Disc loss: 0.414, Gen loss: 7.390, Con loss: 3.512\n",
            "[23/100] - time: 503.70, Disc loss: 0.692, Gen loss: 2.095, Con loss: 3.687\n",
            "[24/100] - time: 503.89, Disc loss: 0.657, Gen loss: 2.107, Con loss: 3.513\n",
            "[25/100] - time: 503.64, Disc loss: 0.631, Gen loss: 2.149, Con loss: 3.598\n",
            "[26/100] - time: 503.82, Disc loss: 0.638, Gen loss: 2.153, Con loss: 3.487\n",
            "[27/100] - time: 503.67, Disc loss: 0.622, Gen loss: 2.193, Con loss: 3.724\n",
            "[28/100] - time: 503.94, Disc loss: 0.615, Gen loss: 2.219, Con loss: 3.674\n",
            "[29/100] - time: 503.75, Disc loss: 0.614, Gen loss: 2.222, Con loss: 3.534\n",
            "[30/100] - time: 503.92, Disc loss: 0.637, Gen loss: 2.192, Con loss: 3.604\n",
            "[31/100] - time: 503.74, Disc loss: 0.607, Gen loss: 2.255, Con loss: 3.462\n",
            "[32/100] - time: 503.97, Disc loss: 0.621, Gen loss: 2.237, Con loss: 3.693\n",
            "[33/100] - time: 503.74, Disc loss: 0.617, Gen loss: 2.269, Con loss: 3.470\n",
            "[34/100] - time: 503.94, Disc loss: 0.605, Gen loss: 2.304, Con loss: 3.622\n",
            "[35/100] - time: 503.72, Disc loss: 0.618, Gen loss: 2.257, Con loss: 3.524\n",
            "[36/100] - time: 503.93, Disc loss: 0.601, Gen loss: 2.327, Con loss: 3.659\n",
            "[37/100] - time: 503.86, Disc loss: 0.599, Gen loss: 2.332, Con loss: 3.605\n",
            "[38/100] - time: 503.98, Disc loss: 0.582, Gen loss: 2.419, Con loss: 3.693\n",
            "[39/100] - time: 503.81, Disc loss: 0.587, Gen loss: 2.411, Con loss: 3.495\n",
            "[40/100] - time: 504.01, Disc loss: 0.585, Gen loss: 2.438, Con loss: 3.609\n",
            "[41/100] - time: 503.87, Disc loss: 0.569, Gen loss: 2.506, Con loss: 3.717\n",
            "[42/100] - time: 504.10, Disc loss: 0.562, Gen loss: 2.539, Con loss: 3.611\n",
            "[43/100] - time: 503.89, Disc loss: 0.569, Gen loss: 2.542, Con loss: 3.720\n",
            "[44/100] - time: 504.16, Disc loss: 0.572, Gen loss: 2.579, Con loss: 3.700\n",
            "[45/100] - time: 503.83, Disc loss: 0.563, Gen loss: 2.582, Con loss: 3.653\n",
            "[46/100] - time: 504.07, Disc loss: 0.550, Gen loss: 2.596, Con loss: 3.825\n",
            "[47/100] - time: 503.89, Disc loss: 0.551, Gen loss: 2.674, Con loss: 3.794\n",
            "[48/100] - time: 504.08, Disc loss: 0.554, Gen loss: 2.650, Con loss: 3.723\n",
            "[49/100] - time: 503.88, Disc loss: 0.551, Gen loss: 2.672, Con loss: 3.846\n",
            "[50/100] - time: 504.09, Disc loss: 0.417, Gen loss: 2.606, Con loss: 3.199\n",
            "[51/100] - time: 503.84, Disc loss: 0.440, Gen loss: 2.615, Con loss: 3.348\n",
            "[52/100] - time: 504.11, Disc loss: 0.448, Gen loss: 2.625, Con loss: 3.348\n",
            "[53/100] - time: 503.86, Disc loss: 0.465, Gen loss: 2.612, Con loss: 3.438\n",
            "[54/100] - time: 504.10, Disc loss: 0.459, Gen loss: 2.633, Con loss: 3.482\n",
            "[55/100] - time: 503.86, Disc loss: 0.478, Gen loss: 2.600, Con loss: 3.519\n",
            "[56/100] - time: 504.11, Disc loss: 0.514, Gen loss: 2.535, Con loss: 3.650\n",
            "[57/100] - time: 503.94, Disc loss: 0.498, Gen loss: 2.532, Con loss: 3.569\n",
            "[58/100] - time: 504.16, Disc loss: 0.496, Gen loss: 2.540, Con loss: 3.592\n",
            "[59/100] - time: 503.93, Disc loss: 0.497, Gen loss: 2.554, Con loss: 3.621\n",
            "[60/100] - time: 504.17, Disc loss: 0.512, Gen loss: 2.529, Con loss: 3.642\n",
            "[61/100] - time: 503.94, Disc loss: 0.508, Gen loss: 2.533, Con loss: 3.640\n",
            "[62/100] - time: 504.27, Disc loss: 0.524, Gen loss: 2.496, Con loss: 3.670\n",
            "[63/100] - time: 504.14, Disc loss: 0.510, Gen loss: 2.518, Con loss: 3.679\n",
            "[64/100] - time: 504.33, Disc loss: 0.529, Gen loss: 2.454, Con loss: 3.679\n",
            "[65/100] - time: 504.12, Disc loss: 0.526, Gen loss: 2.471, Con loss: 3.675\n",
            "[66/100] - time: 504.33, Disc loss: 0.525, Gen loss: 2.467, Con loss: 3.677\n",
            "[67/100] - time: 504.15, Disc loss: 0.529, Gen loss: 2.469, Con loss: 3.695\n",
            "[68/100] - time: 504.25, Disc loss: 0.530, Gen loss: 2.465, Con loss: 3.702\n",
            "[69/100] - time: 503.83, Disc loss: 0.531, Gen loss: 2.462, Con loss: 3.701\n",
            "[70/100] - time: 503.80, Disc loss: 0.533, Gen loss: 2.461, Con loss: 3.702\n",
            "[71/100] - time: 503.58, Disc loss: 0.537, Gen loss: 2.464, Con loss: 3.698\n",
            "[72/100] - time: 503.76, Disc loss: 0.521, Gen loss: 2.471, Con loss: 3.704\n",
            "[73/100] - time: 503.56, Disc loss: 0.536, Gen loss: 2.448, Con loss: 3.689\n",
            "[74/100] - time: 503.76, Disc loss: 0.531, Gen loss: 2.474, Con loss: 3.721\n",
            "[75/100] - time: 503.50, Disc loss: 0.507, Gen loss: 2.412, Con loss: 3.661\n",
            "[76/100] - time: 503.76, Disc loss: 0.514, Gen loss: 2.441, Con loss: 3.668\n",
            "[77/100] - time: 503.65, Disc loss: 0.514, Gen loss: 2.421, Con loss: 3.667\n",
            "[78/100] - time: 503.76, Disc loss: 0.506, Gen loss: 2.458, Con loss: 3.689\n",
            "[79/100] - time: 503.59, Disc loss: 0.519, Gen loss: 2.448, Con loss: 3.698\n",
            "[80/100] - time: 503.81, Disc loss: 0.517, Gen loss: 2.427, Con loss: 3.691\n",
            "[81/100] - time: 503.57, Disc loss: 0.513, Gen loss: 2.453, Con loss: 3.689\n",
            "[82/100] - time: 503.78, Disc loss: 0.512, Gen loss: 2.446, Con loss: 3.696\n",
            "[83/100] - time: 503.60, Disc loss: 0.509, Gen loss: 2.454, Con loss: 3.697\n",
            "[84/100] - time: 503.86, Disc loss: 0.514, Gen loss: 2.460, Con loss: 3.703\n",
            "[85/100] - time: 503.67, Disc loss: 0.528, Gen loss: 2.400, Con loss: 3.706\n",
            "[86/100] - time: 503.87, Disc loss: 0.514, Gen loss: 2.443, Con loss: 3.711\n",
            "[87/100] - time: 503.66, Disc loss: 0.521, Gen loss: 2.433, Con loss: 3.711\n",
            "[88/100] - time: 503.92, Disc loss: 0.525, Gen loss: 2.418, Con loss: 3.715\n",
            "[89/100] - time: 503.82, Disc loss: 0.526, Gen loss: 2.418, Con loss: 3.713\n",
            "[90/100] - time: 504.17, Disc loss: 0.527, Gen loss: 2.416, Con loss: 3.717\n",
            "[91/100] - time: 503.93, Disc loss: 0.520, Gen loss: 2.421, Con loss: 3.715\n",
            "[92/100] - time: 504.17, Disc loss: 0.525, Gen loss: 2.424, Con loss: 3.721\n",
            "[93/100] - time: 503.91, Disc loss: 0.526, Gen loss: 2.401, Con loss: 3.720\n",
            "[94/100] - time: 504.15, Disc loss: 0.517, Gen loss: 2.438, Con loss: 3.722\n",
            "[95/100] - time: 503.92, Disc loss: 0.527, Gen loss: 2.429, Con loss: 3.722\n",
            "[96/100] - time: 504.12, Disc loss: 0.527, Gen loss: 2.434, Con loss: 3.731\n",
            "[97/100] - time: 503.91, Disc loss: 0.518, Gen loss: 2.425, Con loss: 3.720\n",
            "[98/100] - time: 504.12, Disc loss: 0.532, Gen loss: 2.419, Con loss: 3.734\n",
            "[99/100] - time: 503.91, Disc loss: 0.525, Gen loss: 2.413, Con loss: 3.721\n",
            "[100/100] - time: 504.11, Disc loss: 0.521, Gen loss: 2.427, Con loss: 3.727\n",
            "Avg one epoch time: 503.95, total 100 epochs time: 50487.43\n",
            "Training finish!... save training results\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8POMgepCu1sX",
        "outputId": "e262850b-207c-42b5-cb73-1510a7ef5c63"
      },
      "source": [
        "picture_G = generator(3, 3, 64, 8)\n",
        "picture_G.load_state_dict(torch.load(\"/content/drive/MyDrive/11785 Project/11785_to_picture_results/generator_latest.pkl\"))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KEF_O41Xv6oH"
      },
      "source": [
        "def promoting(root, save, Generator):\n",
        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "    Generator.to(device)\n",
        "\n",
        "    if not os.path.isdir(save):\n",
        "        os.makedirs(save)\n",
        "    \n",
        "    src_transform = transforms.Compose([\n",
        "          transforms.Resize((256, 256)), # \n",
        "          transforms.ToTensor(),\n",
        "          # transforms.GaussianBlur(kernel_size= 5, sigma=0.1),\n",
        "          transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
        "    ])\n",
        "\n",
        "    train_loader_src = data_load(root, 'train', src_transform, 1, shuffle=False, drop_last=False) # \n",
        "\n",
        "    for n, (x, _) in enumerate(train_loader_src):\n",
        "        x = x.to(device)\n",
        "        G_recon = Generator(x)\n",
        "        G_recon = G_recon[0].detach().cpu()\n",
        "        path = os.path.join(save, str(n) + '.png')\n",
        "\n",
        "        G_recon = (G_recon.numpy().transpose(1, 2, 0) + 1) / 2\n",
        "        promo_img = np.concatenate(((x[0].detach().cpu().permute(1,2,0)+1)/2, G_recon), 1)\n",
        "        \n",
        "        plt.imsave(path, promo_img)\n",
        "\n",
        "        torch.cuda.empty_cache()\n",
        "        del n\n",
        "        del x\n",
        "        del G_recon\n",
        "        del promo_img\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zMusUf_s0tLS"
      },
      "source": [
        "promoting(os.path.join('data', '/content/drive/MyDrive/11785 Project/src_shinkai Makoto'), os.path.join('data', '/content/drive/MyDrive/11785 Project/src_shinkai Makoto', 'promo_pair'), picture_G)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ICqn4fRmdaQw"
      },
      "source": [
        "import numpy as np\n",
        "import os\n",
        "import shutil\n",
        "\n",
        "# list all files in dir\n",
        "files = os.listdir('/content/drive/MyDrive/11785 Project/src_all/train')\n",
        "\n",
        "random_files = np.random.choice(files, size = 5000, replace=False)\n",
        "for f in random_files:\n",
        "  shutil.copy(os.path.join('/content/drive/MyDrive/11785 Project/src_all/train',f), '/content/drive/MyDrive/11785 Project/random_src/train')\n",
        "  print(f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CVgfycXjjePH"
      },
      "source": [
        "def edge_promoting_v2(root, save):\n",
        "    file_list = os.listdir(root)\n",
        "    if not os.path.isdir(save):\n",
        "        os.makedirs(save)\n",
        "    kernel_size = 5\n",
        "    kernel = np.ones((kernel_size, kernel_size), np.uint8)\n",
        "    gauss = cv2.getGaussianKernel(kernel_size, 0)\n",
        "    gauss = gauss * gauss.transpose(1, 0)\n",
        "    n = 1\n",
        "    for f in tqdm(file_list):\n",
        "        img = cv2.imread(os.path.join(root, f))\n",
        "        real_img = img[:, :256]\n",
        "        rgb_img = img[:, 256:]\n",
        "        gray_img = cv2.imread(os.path.join(root, f), 0)\n",
        "        gray_img = gray_img[:, 256:]\n",
        "        rgb_img = cv2.resize(rgb_img, (256, 256))\n",
        "        pad_img = np.pad(rgb_img, ((2,2), (2,2), (0,0)), mode='reflect')\n",
        "        gray_img = cv2.resize(gray_img, (256, 256))\n",
        "        edges = cv2.Canny(gray_img, 100, 200)\n",
        "        dilation = cv2.dilate(edges, kernel)\n",
        "\n",
        "        gauss_img = np.copy(rgb_img)\n",
        "        idx = np.where(dilation != 0)\n",
        "        for i in range(np.sum(dilation != 0)):\n",
        "            gauss_img[idx[0][i], idx[1][i], 0] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 0], gauss))\n",
        "            gauss_img[idx[0][i], idx[1][i], 1] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 1], gauss))\n",
        "            gauss_img[idx[0][i], idx[1][i], 2] = np.sum(np.multiply(pad_img[idx[0][i]:idx[0][i] + kernel_size, idx[1][i]:idx[1][i] + kernel_size, 2], gauss))\n",
        "\n",
        "        result = np.concatenate((real_img, gauss_img), 1)\n",
        "\n",
        "        cv2.imwrite(os.path.join(save, str(n) + '.png'), result)\n",
        "        n += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yS5Ev7kBjgnF"
      },
      "source": [
        "edge_promoting_v2('/content/drive/MyDrive/11785 Project/src_shinkai Makoto/promo_pair', '/content/drive/MyDrive/11785 Project/src_shinkai Makoto/pair')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2wGF6eEyDtAD"
      },
      "source": [
        "def main(argv):\n",
        "\n",
        "  parser = argparse.ArgumentParser()\n",
        "  parser.add_argument('--name', required=False, default='project_name',  help='')\n",
        "  parser.add_argument('--src_data', required=False, default='src_data_path',  help='sec data path')\n",
        "  parser.add_argument('--tgt_data', required=False, default='tgt_data_path',  help='tgt data path')\n",
        "  parser.add_argument('--vgg_model', required=False, default='pre_trained_VGG19_model_path/vgg19.pth', help='pre-trained VGG19 model path')\n",
        "  parser.add_argument('--in_ngc', type=int, default=3, help='input channel for generator')\n",
        "  parser.add_argument('--out_ngc', type=int, default=3, help='output channel for generator')\n",
        "  parser.add_argument('--in_ndc', type=int, default=3, help='input channel for discriminator')\n",
        "  parser.add_argument('--out_ndc', type=int, default=1, help='output channel for discriminator')\n",
        "  parser.add_argument('--batch_size', type=int, default=8, help='batch size')\n",
        "  parser.add_argument('--ngf', type=int, default=64)\n",
        "  parser.add_argument('--ndf', type=int, default=32)\n",
        "  parser.add_argument('--nb', type=int, default=8, help='the number of resnet block layer for generator')\n",
        "  parser.add_argument('--input_size', type=int, default=256, help='input size')\n",
        "  parser.add_argument('--train_epoch', type=int, default=100)\n",
        "  parser.add_argument('--pre_train_epoch', type=int, default=20)\n",
        "  parser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002')\n",
        "  parser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002')\n",
        "  parser.add_argument('--con_lambda', type=float, default=10, help='lambda for content loss')\n",
        "  parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for Adam optimizer')\n",
        "  parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')\n",
        "  parser.add_argument('--latest_generator_model', required=False, default='', help='the latest trained model path')\n",
        "  parser.add_argument('--latest_discriminator_model', required=False, default='', help='the latest trained model path')\n",
        "  args = parser.parse_args(argv)\n",
        "\n",
        "  print('------------ Options -------------')\n",
        "  for k, v in sorted(vars(args).items()):\n",
        "      print('%s: %s' % (str(k), str(v)))\n",
        "  print('-------------- End ----------------')\n",
        "\n",
        "  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "  if torch.backends.cudnn.enabled:\n",
        "      torch.backends.cudnn.benchmark = True\n",
        "\n",
        "  # results save path\n",
        "  if not os.path.isdir(os.path.join(args.name + '_results', 'Reconstruction')):\n",
        "      os.makedirs(os.path.join(args.name + '_results', 'Reconstruction'))\n",
        "  if not os.path.isdir(os.path.join(args.name + '_results', 'Transfer')):\n",
        "      os.makedirs(os.path.join(args.name + '_results', 'Transfer'))\n",
        "\n",
        "  # edge-promoting\n",
        "  if not os.path.isdir(os.path.join('data', args.tgt_data, 'pair')):\n",
        "      print('edge-promoting start!!')\n",
        "      edge_promoting_v1(os.path.join('data', args.tgt_data), os.path.join('data', args.tgt_data, 'pair'))\n",
        "  else:\n",
        "      print('edge-promoting already done')\n",
        "\n",
        "  # data_loader\n",
        "  src_transform = transforms.Compose([\n",
        "          transforms.Resize((args.input_size, args.input_size)),\n",
        "          transforms.ToTensor(),\n",
        "          transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
        "  ])\n",
        "  tgt_transform = transforms.Compose([\n",
        "          transforms.ToTensor(),\n",
        "          transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))\n",
        "  ])\n",
        "  train_loader_src = data_load(os.path.join('data', args.src_data), 'train', src_transform, args.batch_size, shuffle=True, drop_last=True)\n",
        "  train_loader_tgt = data_load(os.path.join('data', args.tgt_data), 'pair', tgt_transform, args.batch_size, shuffle=True, drop_last=True)\n",
        "  # train_loader_tgt = data_load(os.path.join('data', args.tgt_data), 'promo_pair', tgt_transform, args.batch_size, shuffle=True, drop_last=True)\n",
        "  test_loader_src = data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True)\n",
        "\n",
        "  # network\n",
        "  G = generator(args.in_ngc, args.out_ngc, args.ngf, args.nb)\n",
        "  if args.latest_generator_model != '':\n",
        "      if torch.cuda.is_available():\n",
        "          G.load_state_dict(torch.load(args.latest_generator_model))\n",
        "      else:\n",
        "          # cpu mode\n",
        "          G.load_state_dict(torch.load(args.latest_generator_model, map_location=lambda storage, loc: storage))\n",
        "\n",
        "  D = discriminator(args.in_ndc, args.out_ndc, args.ndf)\n",
        "  if args.latest_discriminator_model != '':\n",
        "      if torch.cuda.is_available():\n",
        "          D.load_state_dict(torch.load(args.latest_discriminator_model))\n",
        "      else:\n",
        "          D.load_state_dict(torch.load(args.latest_discriminator_model, map_location=lambda storage, loc: storage))\n",
        "  VGG = VGG19(init_weights=args.vgg_model, feature_mode=True)\n",
        "  G.to(device)\n",
        "  D.to(device)\n",
        "  VGG.to(device)\n",
        "  G.train()\n",
        "  D.train()\n",
        "  VGG.eval()\n",
        "  print('---------- Networks initialized -------------')\n",
        "  print_network(G)\n",
        "  print_network(D)\n",
        "  print_network(VGG)\n",
        "  print('-----------------------------------------------')\n",
        "\n",
        "  # loss\n",
        "  BCE_loss = nn.BCELoss().to(device)\n",
        "  L1_loss = nn.L1Loss().to(device)\n",
        "\n",
        "  # Adam optimizer\n",
        "  G_optimizer = optim.Adam(G.parameters(), lr=args.lrG, betas=(args.beta1, args.beta2))\n",
        "  D_optimizer = optim.Adam(D.parameters(), lr=args.lrD, betas=(args.beta1, args.beta2))\n",
        "  G_scheduler = optim.lr_scheduler.MultiStepLR(optimizer=G_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1)\n",
        "  D_scheduler = optim.lr_scheduler.MultiStepLR(optimizer=D_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1)\n",
        "\n",
        "  pre_train_hist = {}\n",
        "  pre_train_hist['Recon_loss'] = []\n",
        "  pre_train_hist['per_epoch_time'] = []\n",
        "  pre_train_hist['total_time'] = []\n",
        "\n",
        "  \"\"\" Pre-train reconstruction \"\"\"\n",
        "  if args.latest_generator_model == '':\n",
        "      print('Pre-training start!')\n",
        "      start_time = time.time()\n",
        "      for epoch in range(args.pre_train_epoch):\n",
        "          epoch_start_time = time.time()\n",
        "          Recon_losses = []\n",
        "          for x, _ in train_loader_src:\n",
        "              x = x.to(device)\n",
        "\n",
        "              # train generator G\n",
        "              G_optimizer.zero_grad()\n",
        "\n",
        "              x_feature = VGG((x + 1) / 2)\n",
        "              G_ = G(x)\n",
        "              G_feature = VGG((G_ + 1) / 2)\n",
        "\n",
        "              Recon_loss = 10 * L1_loss(G_feature, x_feature.detach())\n",
        "              Recon_losses.append(Recon_loss.item())\n",
        "              pre_train_hist['Recon_loss'].append(Recon_loss.item())\n",
        "\n",
        "              Recon_loss.backward()\n",
        "              G_optimizer.step()\n",
        "\n",
        "          per_epoch_time = time.time() - epoch_start_time\n",
        "          pre_train_hist['per_epoch_time'].append(per_epoch_time)\n",
        "          print('[%d/%d] - time: %.2f, Recon loss: %.3f' % ((epoch + 1), args.pre_train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Recon_losses))))\n",
        "\n",
        "      total_time = time.time() - start_time\n",
        "      pre_train_hist['total_time'].append(total_time)\n",
        "      with open(os.path.join(args.name + '_results',  'pre_train_hist.pkl'), 'wb') as f:\n",
        "          pickle.dump(pre_train_hist, f)\n",
        "\n",
        "      with torch.no_grad():\n",
        "          G.eval()\n",
        "          for n, (x, _) in enumerate(train_loader_src):\n",
        "              x = x.to(device)\n",
        "              G_recon = G(x)\n",
        "              result = torch.cat((x[0], G_recon[0]), 2)\n",
        "              path = os.path.join(args.name + '_results', 'Reconstruction', '_train_recon_' + str(n + 1) + '.png')\n",
        "              plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "              if n == 4:\n",
        "                  break\n",
        "\n",
        "          for n, (x, _) in enumerate(test_loader_src):\n",
        "              x = x.to(device)\n",
        "              G_recon = G(x)\n",
        "              result = torch.cat((x[0], G_recon[0]), 2)\n",
        "              path = os.path.join(args.name + '_results', 'Reconstruction',  '_test_recon_' + str(n + 1) + '.png')\n",
        "              plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "              if n == 4:\n",
        "                  break\n",
        "  else:\n",
        "      print('Load the latest generator model, no need to pre-train')\n",
        "\n",
        "\n",
        "  train_hist = {}\n",
        "  train_hist['Disc_loss'] = []\n",
        "  train_hist['Gen_loss'] = []\n",
        "  train_hist['Con_loss'] = []\n",
        "  train_hist['per_epoch_time'] = []\n",
        "  train_hist['total_time'] = []\n",
        "  print('training start!')\n",
        "  start_time = time.time()\n",
        "  real = torch.ones(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device)\n",
        "  fake = torch.zeros(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device)\n",
        "  for epoch in range(args.train_epoch):\n",
        "      epoch_start_time = time.time()\n",
        "      G.train()\n",
        "      G_scheduler.step()\n",
        "      D_scheduler.step()\n",
        "      Disc_losses = []\n",
        "      Gen_losses = []\n",
        "      Con_losses = []\n",
        "      for (x, _), (y, _) in zip(train_loader_src, train_loader_tgt):\n",
        "          e = y[:, :, :, args.input_size:]\n",
        "          y = y[:, :, :, :args.input_size]\n",
        "          x, y, e = x.to(device), y.to(device), e.to(device)\n",
        "\n",
        "          # train D\n",
        "          D_optimizer.zero_grad()\n",
        "\n",
        "          D_real = D(y)\n",
        "          D_real_loss = BCE_loss(D_real, real)\n",
        "\n",
        "          G_ = G(x)\n",
        "          D_fake = D(G_)\n",
        "          D_fake_loss = BCE_loss(D_fake, fake)\n",
        "\n",
        "          D_edge = D(e)\n",
        "          D_edge_loss = BCE_loss(D_edge, fake)\n",
        "\n",
        "          Disc_loss = D_real_loss + D_fake_loss + D_edge_loss\n",
        "          Disc_losses.append(Disc_loss.item())\n",
        "          train_hist['Disc_loss'].append(Disc_loss.item())\n",
        "\n",
        "          Disc_loss.backward()\n",
        "          D_optimizer.step()\n",
        "\n",
        "          # train G\n",
        "          G_optimizer.zero_grad()\n",
        "\n",
        "          G_ = G(x)\n",
        "          D_fake = D(G_)\n",
        "          D_fake_loss = BCE_loss(D_fake, real)\n",
        "\n",
        "          x_feature = VGG((x + 1) / 2)\n",
        "          G_feature = VGG((G_ + 1) / 2)\n",
        "          Con_loss = args.con_lambda * L1_loss(G_feature, x_feature.detach())\n",
        "\n",
        "          Gen_loss = D_fake_loss + Con_loss\n",
        "          Gen_losses.append(D_fake_loss.item())\n",
        "          train_hist['Gen_loss'].append(D_fake_loss.item())\n",
        "          Con_losses.append(Con_loss.item())\n",
        "          train_hist['Con_loss'].append(Con_loss.item())\n",
        "\n",
        "          Gen_loss.backward()\n",
        "          G_optimizer.step()\n",
        "\n",
        "\n",
        "      per_epoch_time = time.time() - epoch_start_time\n",
        "      train_hist['per_epoch_time'].append(per_epoch_time)\n",
        "      print(\n",
        "      '[%d/%d] - time: %.2f, Disc loss: %.3f, Gen loss: %.3f, Con loss: %.3f' % ((epoch + 1), args.train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Disc_losses)),\n",
        "          torch.mean(torch.FloatTensor(Gen_losses)), torch.mean(torch.FloatTensor(Con_losses))))\n",
        "\n",
        "      if epoch % 2 == 1 or epoch == args.train_epoch - 1:\n",
        "          with torch.no_grad():\n",
        "              G.eval()\n",
        "              for n, (x, _) in enumerate(train_loader_src):\n",
        "                  x = x.to(device)\n",
        "                  G_recon = G(x)\n",
        "                  result = torch.cat((x[0], G_recon[0]), 2)\n",
        "                  path = os.path.join(args.name + '_results', 'Transfer', str(epoch+1) + '_epoch_'  + '_train_' + str(n + 1) + '.png')\n",
        "                  plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "                  if n == 4:\n",
        "                      break\n",
        "\n",
        "              for n, (x, _) in enumerate(test_loader_src):\n",
        "                  x = x.to(device)\n",
        "                  G_recon = G(x)\n",
        "                  result = torch.cat((x[0], G_recon[0]), 2)\n",
        "                  path = os.path.join(args.name + '_results', 'Transfer', str(epoch+1) + '_epoch_' + '_test_' + str(n + 1) + '.png')\n",
        "                  plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2)\n",
        "                  if n == 4:\n",
        "                      break\n",
        "\n",
        "              torch.save(G.state_dict(), os.path.join(args.name + '_results', 'generator_latest.pkl'))\n",
        "              torch.save(D.state_dict(), os.path.join(args.name + '_results', 'discriminator_latest.pkl'))\n",
        "\n",
        "  total_time = time.time() - start_time\n",
        "  train_hist['total_time'].append(total_time)\n",
        "\n",
        "  print(\"Avg one epoch time: %.2f, total %d epochs time: %.2f\" % (torch.mean(torch.FloatTensor(train_hist['per_epoch_time'])), args.train_epoch, total_time))\n",
        "  print(\"Training finish!... save training results\")\n",
        "\n",
        "  torch.save(G.state_dict(), os.path.join(args.name + '_results',  'generator_param.pkl'))\n",
        "  torch.save(D.state_dict(), os.path.join(args.name + '_results',  'discriminator_param.pkl'))\n",
        "  with open(os.path.join(args.name + '_results',  'train_hist.pkl'), 'wb') as f:\n",
        "      pickle.dump(train_hist, f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-GIS38NzaLU5"
      },
      "source": [
        "main([\n",
        "    \"--name\", \"/content/drive/MyDrive/11785\",\n",
        "    \"--src_data\", \"/content/drive/MyDrive/11785 Project/random_src\",  \n",
        "    \"--tgt_data\", \"/content/drive/MyDrive/11785 Project/src_shinkai Makoto\",\n",
        "    \"--vgg_model\",\"/content/drive/MyDrive/11785 Project/vgg19-dcbb9e9d.pth\",\n",
        "    \"--train_epoch\", \"110\",\n",
        "    \"--batch_size\", \"8\"])"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}