CartoonGAN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/TobiasSunderdiek/cartoon-gan/blob/master/CartoonGAN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "61gxheM2om8l"
      },
      "source": [
        "# CartoonGAN\n",
        "\n",
        "This notebook contains the implementation of the cartoon GAN model. It is implemented with PyTorch. See README [here](https://github.com/TobiasSunderdiek/cartoon-gan/blob/master/README.md) for more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U-GRcOxmm-8_"
      },
      "source": [
        "## Generate dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6PHbkMli7BFc"
      },
      "source": [
        "## Transfer data via google drive\n",
        "- all image data in this notebook is expected to be zipped to files on local computer as described in README of this project [here](https://github.com/TobiasSunderdiek/cartoon-gan/blob/master/README.md)\n",
        "- create folder `cartoonGAN` in `My Drive` in google drive\n",
        "- copy .zip-files `coco.zip`, `safebooru.zip` and `safebooru_smoothed.zip` to google drive `My Drive`/`cartoonGAN`\n",
        "- mount google drive in this notebook by executing cell below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_WzcH9ef4dc4"
      },
      "outputs": [
        {
          "ename": "ModuleNotFoundError",
          "evalue": "No module named 'google'",
          "output_type": "error",
          "traceback": [
            "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
            "\u001b[1;32m~\\AppData\\Local\\Temp/ipykernel_10836/2152992096.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mfrom\u001b[0m \u001b[0mgoogle\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcolab\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mdrive\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      2\u001b[0m \u001b[0mdrive\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmount\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'/content/data'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
            "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'google'"
          ]
        }
      ],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/data')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8Pee_KkYm_W1"
      },
      "source": [
        "### cartoons images\n",
        "\n",
        "- cartoon images are located in file `content/data/My Drive/cartoonGAN/safebooru.zip` of this notebook\n",
        "- extract images and place in folder `cartoons` by executing cell below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "b342SwfIR1ur"
      },
      "outputs": [],
      "source": [
        "!mkdir cartoons\n",
        "!mkdir cartoons/1\n",
        "!unzip -n -q /content/data/My\\ Drive/cartoonGAN/safebooru.zip -d cartoons/1/ #extract to subfolder due to DataLoader needs subdirectories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yPj9lzdNmvir"
      },
      "source": [
        "##### data-loader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ebY_JkCBVuyf"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.datasets import ImageFolder\n",
        "from torchvision import transforms\n",
        "from torch.utils.data import random_split\n",
        "import math\n",
        "\n",
        "image_size = 256\n",
        "batch_size = 16\n",
        "\n",
        "transformer = transforms.Compose([\n",
        "    transforms.CenterCrop(image_size),\n",
        "    transforms.ToTensor() # ToTensor() changes the range of the values from [0, 255] to [0.0, 1.0]\n",
        "])\n",
        "\n",
        "cartoon_dataset = ImageFolder('cartoons/', transformer)\n",
        "len_training_set = math.floor(len(cartoon_dataset) * 0.9)\n",
        "len_valid_set = len(cartoon_dataset) - len_training_set\n",
        "\n",
        "training_set, _ = random_split(cartoon_dataset, (len_training_set, len_valid_set))\n",
        "cartoon_image_dataloader_train = DataLoader(training_set, batch_size, shuffle=True, num_workers=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "G0kxqmDopf0P"
      },
      "source": [
        "#### show examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "AYapuAsuoOrO"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "def show_sample_image(dataloader):\n",
        "  iterator = iter(dataloader)\n",
        "  sample_batch, _ = iterator.next()\n",
        "  first_sample_image_of_batch = sample_batch[0]\n",
        "  print(first_sample_image_of_batch.size())\n",
        "  print(\"Current range: {} to {}\".format(first_sample_image_of_batch.min(), first_sample_image_of_batch.max()))\n",
        "  plt.imshow(np.transpose(first_sample_image_of_batch.numpy(), (1, 2, 0)))\n",
        "\n",
        "show_sample_image(cartoon_image_dataloader_train)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_PUKl2PHRYYP"
      },
      "source": [
        "### edge-smoothed cartoons images\n",
        "\n",
        "- edge-smoothed cartoon images are located in file `content/data/My Drive/cartoonGAN/safebooru_smoothed.zip` of this notebook\n",
        "- extract images and place in folder `cartoons_smoothed` by executing cell below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UQVv5Q-IRo1_"
      },
      "outputs": [],
      "source": [
        "!mkdir cartoons_smoothed\n",
        "!mkdir cartoons_smoothed/1\n",
        "!unzip -n -q /content/data/My\\ Drive/cartoonGAN/safebooru_smoothed.zip -d cartoons_smoothed/1/ #extract to subfolder due to DataLoader needs subdirectories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hqBoaA8tSZh8"
      },
      "source": [
        "##### data-loader\n",
        "\n",
        "same configuration as cartoon data loader above"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8aepIFEfSiZ3"
      },
      "outputs": [],
      "source": [
        "smoothed_cartoon_dataset = ImageFolder('cartoons_smoothed/', transformer)\n",
        "len_training_set = math.floor(len(smoothed_cartoon_dataset) * 0.9)\n",
        "len_valid_set = len(smoothed_cartoon_dataset) - len_training_set\n",
        "training_set, _ = random_split(smoothed_cartoon_dataset, (len_training_set, len_valid_set))\n",
        "smoothed_cartoon_image_dataloader_train = DataLoader(training_set, batch_size, shuffle=True, num_workers=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6IM7YUkgSE96"
      },
      "source": [
        "#### show examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "z_qeh0BqSE-B"
      },
      "outputs": [],
      "source": [
        "show_sample_image(smoothed_cartoon_image_dataloader_train)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5gL-nH1Q31vd"
      },
      "source": [
        "### photos"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_PlLTrVWKtqS"
      },
      "source": [
        "- photos are located in file `content/data/My Drive/cartoonGAN/coco.zip` of this notebook\n",
        "- extract images and place in folder `photos` by executing cell below"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "TQYRvLUWrPO9"
      },
      "outputs": [],
      "source": [
        "!mkdir photos\n",
        "!mkdir photos/1\n",
        "!unzip -n -q /content/data/My\\ Drive/cartoonGAN/coco.zip -d photos/1 #extract to subfolder due to DataLoader needs subdirectories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LvscdHFdNLbf"
      },
      "source": [
        "##### data-loader\n",
        "same configuration as cartoon data loader above"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5Db-YnYWNRIW"
      },
      "outputs": [],
      "source": [
        "photo_dataset = ImageFolder('photos/', transformer)\n",
        "len_training_set = math.floor(len(photo_dataset) * 0.9)\n",
        "len_valid_set = len(photo_dataset) - len_training_set\n",
        "training_set, validation_set = random_split(photo_dataset, (len_training_set, len_valid_set))\n",
        "photo_dataloader_train = DataLoader(training_set, batch_size, shuffle=True, num_workers=0)\n",
        "photo_dataloader_valid = DataLoader(validation_set, batch_size, shuffle=True, num_workers=0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qWreH4XwNplc"
      },
      "source": [
        "#### show examples"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sIO325gNNqIY"
      },
      "outputs": [],
      "source": [
        "show_sample_image(photo_dataloader_train)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rhuxp1xAq4CK"
      },
      "source": [
        "## Setup tensorboard\n",
        "\n",
        "Use tensorboard to have an eye on weights and losses."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "yYjbil6vrsnt"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/data/My\\ Drive/cartoonGAN/tensorboard/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0NqCxs_PrL1_"
      },
      "outputs": [],
      "source": [
        "%tensorflow_version 1.x\n",
        "\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "\n",
        "tensorboard_logdir = '/content/data/My Drive/cartoonGAN/tensorboard'\n",
        "writer = SummaryWriter(tensorboard_logdir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "q9mH063RfqHi"
      },
      "source": [
        "###Generator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "a2EmDnsucENx"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch import sigmoid\n",
        "\n",
        "class ResidualBlock(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(ResidualBlock, self).__init__()\n",
        "    self.conv_1 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)\n",
        "    self.conv_2 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)\n",
        "    self.norm_1 = nn.BatchNorm2d(256)\n",
        "    self.norm_2 = nn.BatchNorm2d(256)\n",
        "\n",
        "  def forward(self, x):\n",
        "    output = self.norm_2(self.conv_2(F.relu(self.norm_1(self.conv_1(x)))))\n",
        "    return output + x #ES\n",
        "\n",
        "class Generator(nn.Module):\n",
        "    def __init__(self):\n",
        "      super(Generator, self).__init__()\n",
        "      self.conv_1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=1, padding=3)\n",
        "      self.norm_1 = nn.BatchNorm2d(64)\n",
        "      \n",
        "      # down-convolution #\n",
        "      self.conv_2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1)\n",
        "      self.conv_3 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)\n",
        "      self.norm_2 = nn.BatchNorm2d(128)\n",
        "      \n",
        "      self.conv_4 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2, padding=1)\n",
        "      self.conv_5 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)\n",
        "      self.norm_3 = nn.BatchNorm2d(256)\n",
        "      \n",
        "      # residual blocks #\n",
        "      residualBlocks = []\n",
        "      for l in range(8):\n",
        "        residualBlocks.append(ResidualBlock())\n",
        "      self.res = nn.Sequential(*residualBlocks)\n",
        "      \n",
        "      # up-convolution #\n",
        "      self.conv_6 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
        "      self.conv_7 = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1)\n",
        "      self.norm_4 = nn.BatchNorm2d(128)\n",
        "\n",
        "      self.conv_8 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
        "      self.conv_9 = nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1)\n",
        "      self.norm_5 = nn.BatchNorm2d(64)\n",
        "      \n",
        "      self.conv_10 = nn.Conv2d(in_channels=64, out_channels=3, kernel_size=7, stride=1, padding=3)\n",
        "\n",
        "    def forward(self, x):\n",
        "      x = F.relu(self.norm_1(self.conv_1(x)))\n",
        "      \n",
        "      x = F.relu(self.norm_2(self.conv_3(self.conv_2(x))))\n",
        "      x = F.relu(self.norm_3(self.conv_5(self.conv_4(x))))\n",
        "      \n",
        "      x = self.res(x)\n",
        "      x = F.relu(self.norm_4(self.conv_7(self.conv_6(x))))\n",
        "      x = F.relu(self.norm_5(self.conv_9(self.conv_8(x))))\n",
        "\n",
        "      x = self.conv_10(x)\n",
        "\n",
        "      x = sigmoid(x)\n",
        "\n",
        "      return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ATqXFxcpz0eT"
      },
      "outputs": [],
      "source": [
        "G = Generator()\n",
        "print(G)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SjtvIkqrVy34"
      },
      "source": [
        "### Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "UxJ2dldGz5Nw"
      },
      "outputs": [],
      "source": [
        "class Discriminator(nn.Module):\n",
        "  def __init__(self):\n",
        "     super(Discriminator, self).__init__()\n",
        "     self.conv_1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)\n",
        "      \n",
        "     self.conv_2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1)\n",
        "     self.conv_3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)\n",
        "     self.norm_1 = nn.BatchNorm2d(128)\n",
        "      \n",
        "     self.conv_4 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=2, padding=1)\n",
        "     self.conv_5 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)\n",
        "     self.norm_2 = nn.BatchNorm2d(256)\n",
        "    \n",
        "     self.conv_6 = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1)\n",
        "     self.norm_3 = nn.BatchNorm2d(256)\n",
        "    \n",
        "     self.conv_7 = nn.Conv2d(in_channels=256, out_channels=1, kernel_size=3, stride=1, padding=1)\n",
        "\n",
        "  def forward(self, x):\n",
        "    x = F.leaky_relu(self.conv_1(x))\n",
        "    x = F.leaky_relu(self.norm_1(self.conv_3(F.leaky_relu(self.conv_2(x)))), negative_slope=0.2)\n",
        "    x = F.leaky_relu(self.norm_2(self.conv_5(F.leaky_relu(self.conv_4(x)))), negative_slope=0.2)\n",
        "    x = F.leaky_relu(self.norm_3(self.conv_6(x)), negative_slope=0.2)\n",
        "    x = self.conv_7(x)\n",
        "    x = sigmoid(x)\n",
        "    \n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "oI6tJdLmlUHp"
      },
      "outputs": [],
      "source": [
        "D = Discriminator()\n",
        "print(D)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nwVzSn9nB1jL"
      },
      "source": [
        "### use device CPU or GPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3IMrqAe1lVQ2"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "device = torch.device('cpu')\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  device = torch.device('cuda')\n",
        "  print(\"Train on GPU.\")\n",
        "else:\n",
        "  print(\"No cuda available\")\n",
        "\n",
        "G.to(device)\n",
        "D.to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "edVjF0YerHgM"
      },
      "source": [
        "### VGG-16\n",
        "Load already downloaded vgg-16 weights from drive, or download and save to drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZGwo1goXCW_K"
      },
      "outputs": [],
      "source": [
        "from torchvision import models\n",
        "\n",
        "path_to_pretrained_vgg16 = '/content/data/My Drive/cartoonGAN/vgg16-397923af.pth'\n",
        "\n",
        "try:\n",
        "  pretrained = torch.load(path_to_pretrained_vgg16)\n",
        "  vgg16 = models.vgg16(pretrained=False)\n",
        "  vgg16.load_state_dict(pretrained)\n",
        "  vgg16 = vgg16.to(device)\n",
        "except FileNotFoundError:\n",
        "  vgg16 = models.vgg16(pretrained=True)\n",
        "  torch.save(vgg16, path_to_pretrained_vgg16)\n",
        "  vgg16 = vgg16.to(device) # Credits to bdlneto (https://github.com/bdlneto) https://github.com/TobiasSunderdiek/cartoon-gan/issues/5 for finding this issue and testing the new version\n",
        "\n",
        "print(vgg16)\n",
        "\n",
        "# due VGG16 has 5 pooling-layer, I assume conv4_4 is the 4th pooling layer\n",
        "# (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
        "feature_extractor = vgg16.features[:24]\n",
        "for param in feature_extractor.parameters():\n",
        "  param.require_grad = False\n",
        "\n",
        "print(feature_extractor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "McDAYOutQ0Xh"
      },
      "source": [
        "### Two loss functions\n",
        "\n",
        "- discriminator loss\n",
        "- generator initialization phase loss and generator loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "mjkXoxfaRI5y"
      },
      "outputs": [],
      "source": [
        "from torchvision import models\n",
        "from torch.nn import BCELoss\n",
        "\n",
        "class DiscriminatorLoss(torch.nn.Module):\n",
        "  def __init__(self):\n",
        "      super(DiscriminatorLoss, self).__init__()\n",
        "      self.bce_loss = BCELoss()\n",
        "\n",
        "  def forward(self, discriminator_output_of_cartoon_input,\n",
        "              discriminator_output_of_cartoon_smoothed_input,\n",
        "              discriminator_output_of_generated_image_input,\n",
        "              epoch,\n",
        "              write_to_tensorboard=False):\n",
        "\n",
        "    return self._adversarial_loss(discriminator_output_of_cartoon_input,\n",
        "                     discriminator_output_of_cartoon_smoothed_input,\n",
        "                     discriminator_output_of_generated_image_input,\n",
        "                     epoch,\n",
        "                     write_to_tensorboard)\n",
        "\n",
        "  def _adversarial_loss(self, discriminator_output_of_cartoon_input,\n",
        "                     discriminator_output_of_cartoon_smoothed_input,\n",
        "                     discriminator_output_of_generated_image_input,\n",
        "                     epoch,\n",
        "                     write_to_tensorboard):\n",
        "\n",
        "    # define ones and zeros here instead within __init__ due to have same shape as input\n",
        "    # due to testing different batch_sizes, sometimes the \"last batch\" has < batch_size elements\n",
        "    actual_batch_size = discriminator_output_of_cartoon_input.size()[0]\n",
        "    zeros = torch.zeros([actual_batch_size, 1, 64, 64]).to(device)\n",
        "    ones = torch.ones([actual_batch_size, 1, 64, 64]).to(device)\n",
        "\n",
        "    d_loss_cartoon = self.bce_loss(discriminator_output_of_cartoon_input, ones)\n",
        "    d_loss_cartoon_smoothed = self.bce_loss(discriminator_output_of_cartoon_smoothed_input, zeros)\n",
        "    d_loss_generated_input = self.bce_loss(discriminator_output_of_generated_image_input, zeros)\n",
        "\n",
        "    d_loss = d_loss_cartoon + d_loss_cartoon_smoothed + d_loss_generated_input\n",
        "\n",
        "    if write_to_tensorboard:\n",
        "      writer.add_scalar('d_loss_cartoon', d_loss_cartoon,epoch)\n",
        "      writer.add_scalar('d_loss_cartoon_smoothed', d_loss_cartoon_smoothed, epoch)\n",
        "      writer.add_scalar('d_loss_generated_input', d_loss_generated_input, epoch)\n",
        "      writer.add_scalar('d_loss', d_loss, epoch)\n",
        "\n",
        "    return d_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "E-JzisW-65IN"
      },
      "outputs": [],
      "source": [
        "class GeneratorLoss(torch.nn.Module):\n",
        "  def __init__(self):\n",
        "      super(GeneratorLoss, self).__init__()\n",
        "      self.w = 0.000005\n",
        "      self.bce_loss = BCELoss()\n",
        "      self.feature_extractor = vgg16.features[:24]\n",
        "      for param in self.feature_extractor.parameters():\n",
        "        param.require_grad = False\n",
        "\n",
        "  def forward(self, discriminator_output_of_generated_image_input,\n",
        "              generator_input,\n",
        "              generator_output,\n",
        "              epoch,\n",
        "              is_init_phase=False,\n",
        "              write_to_tensorboard=False):\n",
        "    if is_init_phase:\n",
        "      g_content_loss = self._content_loss(generator_input, generator_output)\n",
        "      g_adversarial_loss = 0.0\n",
        "      g_loss = g_content_loss\n",
        "    else:\n",
        "      g_adversarial_loss = self._adversarial_loss_generator_part_only(discriminator_output_of_generated_image_input)\n",
        "      g_content_loss = self._content_loss(generator_input, generator_output)\n",
        "      g_loss = g_adversarial_loss + self.w * g_content_loss\n",
        "\n",
        "    if write_to_tensorboard:\n",
        "      writer.add_scalar('g_adversarial_loss', g_adversarial_loss, epoch)\n",
        "      writer.add_scalar('g_content_loss', g_content_loss, epoch)\n",
        "      writer.add_scalar('g_loss', g_loss, epoch)\n",
        "\n",
        "    return g_loss\n",
        "\n",
        "  def _adversarial_loss_generator_part_only(self, discriminator_output_of_generated_image_input):\n",
        "    actual_batch_size = discriminator_output_of_generated_image_input.size()[0]\n",
        "    ones = torch.ones([actual_batch_size, 1, 64, 64]).to(device)\n",
        "    return self.bce_loss(discriminator_output_of_generated_image_input, ones)\n",
        "\n",
        "  def _content_loss(self, generator_input, generator_output):\n",
        "    return (self.feature_extractor(generator_output) - self.feature_extractor(generator_input)).norm(p=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "XAlrFRkH2WaK"
      },
      "outputs": [],
      "source": [
        "discriminatorLoss = DiscriminatorLoss()\n",
        "generatorLoss = GeneratorLoss()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "CF8stH-X3e7j"
      },
      "source": [
        "## Optimizer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6MwrrY1z34WR"
      },
      "outputs": [],
      "source": [
        "import torch.optim as optim\n",
        "\n",
        "lr = 0.0002\n",
        "beta1 = 0.5\n",
        "beta2 = 0.999\n",
        "\n",
        "d_optimizer = optim.Adam(D.parameters(), lr, [beta1, beta2])\n",
        "g_optimizer = optim.Adam(G.parameters(), lr, [beta1, beta2])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pGP0eIqibdAN"
      },
      "source": [
        "## Saving\n",
        "To make training resumeable, I save some checkpoints to google drive and load them, if existing, before run the training.\n",
        "\n",
        "I also save weights and bias from generator and discriminator to tensorboard.\n",
        "\n",
        "For checking some intermediate images of the generator, I save them to google drive."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "K9b23aS94msx"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/data/My\\ Drive/cartoonGAN/checkpoints/\n",
        "!mkdir -p /content/data/My\\ Drive/cartoonGAN/intermediate_results/training/\n",
        "intermediate_results_training_path = \"/content/data/My Drive/cartoonGAN/intermediate_results/training/\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "s4zSMmblNxWR"
      },
      "outputs": [],
      "source": [
        "def save_training_result(input, output):\n",
        "  # input/output has batch-size number of images, get first one and detach from tensor\n",
        "  image_input = input[0].detach().cpu().numpy()\n",
        "  image_output = output[0].detach().cpu().numpy()\n",
        "  # transponse image from torch.Size([3, 256, 256]) to (256, 256, 3)\n",
        "  image_input = np.transpose(image_input, (1, 2, 0))\n",
        "  image_output = np.transpose(image_output, (1, 2, 0))\n",
        "\n",
        "  # generate filenames as timestamp, this orders the output by time\n",
        "  filename = str(int(time.time()))\n",
        "  path_input = intermediate_results_training_path + filename + \"_input.jpg\"\n",
        "  path_output = intermediate_results_training_path + filename + \".jpg\"\n",
        "  plt.imsave(path_input, image_input)\n",
        "  plt.imsave(path_output, image_output)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "SJYC1ec5ye3t"
      },
      "outputs": [],
      "source": [
        "def write_model_weights_and_bias_to_tensorboard(prefix, state_dict, epoch):\n",
        "  for param in state_dict:\n",
        "      writer.add_histogram(f\"{prefix}_{param}\", state_dict[param], epoch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "6weYKGsHB-Vo"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZoxZv-93DEXS"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "def train(_num_epochs, checkpoint_dir, best_valid_loss, epochs_already_done, losses, validation_losses):\n",
        "  init_epochs = 10\n",
        "  print_every = 100\n",
        "  start_time = time.time()\n",
        "\n",
        "  for epoch in range(_num_epochs - epochs_already_done):\n",
        "    epoch = epoch + epochs_already_done\n",
        "\n",
        "    for index, ((photo_images, _), (smoothed_cartoon_images, _), (cartoon_images, _)) in enumerate(zip(photo_dataloader_train, smoothed_cartoon_image_dataloader_train, cartoon_image_dataloader_train)):\n",
        "      batch_size = photo_images.size(0)\n",
        "      photo_images = photo_images.to(device)\n",
        "      smoothed_cartoon_images = smoothed_cartoon_images.to(device)\n",
        "      cartoon_images = cartoon_images.to(device)\n",
        "\n",
        "      # train the discriminator\n",
        "      d_optimizer.zero_grad()\n",
        "      \n",
        "      d_of_cartoon_input = D(cartoon_images)\n",
        "      d_of_cartoon_smoothed_input = D(smoothed_cartoon_images)\n",
        "      d_of_generated_image_input = D(G(photo_images))\n",
        "\n",
        "      write_only_one_loss_from_epoch_not_every_batch_loss = (index == 0)\n",
        "\n",
        "      d_loss = discriminatorLoss(d_of_cartoon_input,\n",
        "                                 d_of_cartoon_smoothed_input,\n",
        "                                 d_of_generated_image_input,\n",
        "                                 epoch,\n",
        "                                 write_to_tensorboard=write_only_one_loss_from_epoch_not_every_batch_loss)\n",
        "\n",
        "      d_loss.backward()\n",
        "      d_optimizer.step()\n",
        "\n",
        "      # train the generator\n",
        "      g_optimizer.zero_grad()\n",
        "\n",
        "      g_output = G(photo_images)\n",
        "\n",
        "      d_of_generated_image_input = D(g_output)\n",
        "\n",
        "      if epoch < init_epochs:\n",
        "        # init\n",
        "        init_phase = True\n",
        "      else:\n",
        "        # train\n",
        "        init_phase = False\n",
        "\n",
        "      g_loss = generatorLoss(d_of_generated_image_input,\n",
        "                              photo_images,\n",
        "                              g_output,\n",
        "                              epoch,\n",
        "                              is_init_phase=init_phase,\n",
        "                              write_to_tensorboard=write_only_one_loss_from_epoch_not_every_batch_loss)\n",
        "\n",
        "      g_loss.backward()\n",
        "      g_optimizer.step()\n",
        "\n",
        "      if (index % print_every) == 0:\n",
        "        losses.append((d_loss.item(), g_loss.item()))\n",
        "        now = time.time()\n",
        "        current_run_time = now - start_time\n",
        "        start_time = now\n",
        "        print(\"Epoch {}/{} | d_loss {:6.4f} | g_loss {:6.4f} | time {:2.0f}s | total no. of losses {}\".format(epoch+1, _num_epochs, d_loss.item(), g_loss.item(), current_run_time, len(losses)))\n",
        "    \n",
        "    # write to tensorboard\n",
        "      #write_model_weights_and_bias_to_tensorboard('D', D.state_dict(), epoch)\n",
        "      #write_model_weights_and_bias_to_tensorboard('G', G.state_dict(), epoch)\n",
        "    # save some intermediate results during training\n",
        "    save_training_result(photo_images, g_output)\n",
        "\n",
        "    # validate\n",
        "    with torch.no_grad():\n",
        "      D.eval()\n",
        "      G.eval()\n",
        "\n",
        "      for batch_index, (photo_images, _) in enumerate(photo_dataloader_valid):\n",
        "        photo_images = photo_images.to(device)\n",
        "\n",
        "        g_output = G(photo_images)\n",
        "        d_of_generated_image_input = D(g_output)\n",
        "        g_valid_loss = generatorLoss(d_of_generated_image_input,\n",
        "                                      photo_images,\n",
        "                                      g_output,\n",
        "                                      epoch,\n",
        "                                      is_init_phase=init_phase,\n",
        "                                      write_to_tensorboard=write_only_one_loss_from_epoch_not_every_batch_loss)\n",
        "\n",
        "        if batch_index % print_every == 0:\n",
        "          validation_losses.append(g_valid_loss.item())\n",
        "          now = time.time()\n",
        "          current_run_time = now - start_time\n",
        "          start_time = now\n",
        "          print(\"Epoch {}/{} | validation loss {:6.4f} | time {:2.0f}s | total no. of losses {}\".format(epoch+1, _num_epochs, g_valid_loss.item(), current_run_time, len(validation_losses)))\n",
        "\n",
        "    D.train()\n",
        "    G.train()\n",
        "    \n",
        "    if(g_valid_loss.item() < best_valid_loss):\n",
        "      print(\"Generator loss improved from {} to {}\".format(best_valid_loss, g_valid_loss.item()))\n",
        "      best_valid_loss = g_valid_loss.item()\n",
        "  \n",
        "    # save checkpoint\n",
        "    checkpoint = {'g_valid_loss': g_valid_loss.item(),\n",
        "                  'best_valid_loss': best_valid_loss,\n",
        "                  'losses': losses,\n",
        "                  'validation_losses': validation_losses,\n",
        "                  'last_epoch': epoch+1,\n",
        "                  'd_state_dict': D.state_dict(),\n",
        "                  'g_state_dict': G.state_dict(),\n",
        "                  'd_optimizer_state_dict': d_optimizer.state_dict(),\n",
        "                  'g_optimizer_state_dict': g_optimizer.state_dict()\n",
        "                }\n",
        "    print(\"Save checkpoint for validation loss of {}\".format(g_valid_loss.item()))\n",
        "    torch.save(checkpoint, checkpoint_dir + '/checkpoint_epoch_{:03d}.pth'.format(epoch+1))\n",
        "    if(best_valid_loss == g_valid_loss.item()):\n",
        "      print(\"Overwrite best checkpoint\")\n",
        "      torch.save(checkpoint, checkpoint_dir + '/best_checkpoint.pth')\n",
        "    \n",
        "  return losses, validation_losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "9P6_f7NTDXyM"
      },
      "outputs": [],
      "source": [
        "from os import listdir\n",
        "\n",
        "checkpoint_dir = '/content/data/My Drive/cartoonGAN/checkpoints'\n",
        "checkpoints = listdir(checkpoint_dir)\n",
        "num_epochs = 200 + 10 # training + init phase\n",
        "epochs_already_done = 0\n",
        "best_valid_loss = math.inf\n",
        "losses = []\n",
        "validation_losses = []\n",
        "\n",
        "if(len(checkpoints) > 0):\n",
        "  last_checkpoint = sorted(checkpoints)[-1]\n",
        "  checkpoint = torch.load(checkpoint_dir + '/' + last_checkpoint, map_location=torch.device(device))\n",
        "  best_valid_loss = checkpoint['best_valid_loss']\n",
        "  epochs_already_done = checkpoint['last_epoch']\n",
        "  losses = checkpoint['losses']\n",
        "  validation_losses = checkpoint['validation_losses']\n",
        "  \n",
        "  D.load_state_dict(checkpoint['d_state_dict'])\n",
        "  G.load_state_dict(checkpoint['g_state_dict'])\n",
        "  d_optimizer.load_state_dict(checkpoint['d_optimizer_state_dict'])\n",
        "  g_optimizer.load_state_dict(checkpoint['g_optimizer_state_dict'])\n",
        "  print('Load checkpoint {} with g_valid_loss {}, best_valid_loss {}, {} epochs and total no of losses {}'.format(last_checkpoint, checkpoint['g_valid_loss'], best_valid_loss, epochs_already_done, len(losses)))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ppzLtQGwBBAJ"
      },
      "outputs": [],
      "source": [
        "losses, validation_losses = train(num_epochs, checkpoint_dir, best_valid_loss, epochs_already_done, losses, validation_losses)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "3VxBiLKImhYq"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "%config InlineBackend.figure.format = 'retina'\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "d_losses = [x[0] for x in losses]\n",
        "g_losses = [x[1] for x in losses]\n",
        "plt.plot(d_losses, label='Discriminator training loss')\n",
        "plt.plot(g_losses, label='Generator training loss')\n",
        "plt.plot(validation_losses, label='Generator validation loss')\n",
        "plt.legend(frameon=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "VR4iOWxP7ZK4"
      },
      "source": [
        "### Show results in tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "H7QbonJE7WcK"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir='/content/data/My Drive/cartoonGAN/tensorboard'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JJ4wLFRfr0dS"
      },
      "source": [
        "##Test\n",
        "Bug: It seems that I make a mistake somewhere in the saving logic of the best checkpoint. The best checkpoint seems not to be the best comic-stlye looking result if compared with the model trained after 210 epochs. Therefore see both results tested below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Gg4bnUSksAiT"
      },
      "outputs": [],
      "source": [
        "checkpoint = torch.load(checkpoint_dir + '/best_checkpoint.pth', map_location=torch.device(device))\n",
        "G_inference = Generator()\n",
        "G_inference.load_state_dict(checkpoint['g_state_dict'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "k0nUlPZIALF-"
      },
      "outputs": [],
      "source": [
        "test_images = iter(photo_dataloader_valid).next()[0]\n",
        "result_images_best_checkpoint = G_inference(test_images)\n",
        "result_images_checkpoint_after_210_epochs = G(test_images)\n",
        "print(result_images_best_checkpoint[0])\n",
        "plt.imshow(np.transpose(result_images_best_checkpoint[0].detach().numpy(), (1, 2, 0)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "8kLL0_V_tcpB"
      },
      "outputs": [],
      "source": [
        "plt.imshow(np.transpose(result_images_checkpoint_after_210_epochs[0].detach().numpy(), (1, 2, 0)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aA13FqynYOHz"
      },
      "source": [
        "### Save weights for production release\n",
        "In order to have a small file size for a production release, save only the weights of the generator to get a smaller artifact.\n",
        "\n",
        "This should be the weights of the best checkpoint, but due to the bug described above, I take the latest version of the generator.\n",
        "\n",
        "Filesize comparison:\n",
        "\n",
        "Full checkpoint with Generator, Discriminator, Optimizer: 140MB\n",
        "\n",
        "Generator only: 43MB"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5d63A8uZ7vGc"
      },
      "outputs": [],
      "source": [
        "generator_release = {'last_epoch': epoch+1,\n",
        "                    'g_state_dict': G.state_dict(),\n",
        "                   }\n",
        "print(\"Save generator release\")\n",
        "torch.save(generator_release, checkpoint_dir + '/generator_release.pth')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "include_colab_link": true,
      "machine_shape": "hm",
      "name": "CartoonGAN.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}