train.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CycleGAN_train.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PUZCMsH5LTQ2",
        "colab_type": "text"
      },
      "source": [
        "# Notebook for training CycleGAN to transform one architecture style to another"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aVq9m4Szz9Kk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%bash\n",
        "pip3 install torch torchvision\n",
        "pip3 install pillow==4.1.1\n",
        "git clone https://github.com/sergeisoly/CycleGAN"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LMLbb5sF0Ke_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "from PIL import Image\n",
        "import itertools\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "plt.rcParams['axes.grid'] = 'True'\n",
        "\n",
        "import glob\n",
        "import random\n",
        "import os\n",
        "\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.models as models\n",
        "\n",
        "import copy\n",
        "import warnings\n",
        "from tqdm import tqdm, tqdm_notebook\n",
        "from IPython.display import clear_output\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(action='ignore')\n",
        "\n",
        "from CycleGAN import nets, functions, data"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8vfokw5d0L2H",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "print(device)\n",
        "!nvidia-smi"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyr1JtlXPmOa",
        "colab_type": "text"
      },
      "source": [
        "Download dataset with different architecture styles \\\n",
        "Dataset was taken from https://sites.google.com/site/zhexuutssjtu/projects/arch \\\n",
        "Instruction for downloading large files (>100 MB) from Google Drive \\\n",
        "https://medium.com/@acpanjan/download-google-drive-files-using-wget-3c2c025a8b99"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KhaWByfbLDHD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%bash\n",
        "wget --load-cookies /tmp/cookies.txt \\\n",
        "\"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt \\\n",
        "--keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=0Bwo0SFiZwl3JVGRlWGZUaW5va00' \\\n",
        "-O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=0Bwo0SFiZwl3JVGRlWGZUaW5va00\" -O arcDataset.zip  && rm -rf /tmp/cookies.txt"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5nXXUyNxan9w",
        "colab_type": "text"
      },
      "source": [
        "Unpack dataset to current colab directory\n",
        "\n",
        "I tried transforming from Gothic Style to International and back,\n",
        "you can try other."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ic8hVB9j0N7B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%bash\n",
        "unzip arcDataset.zip -d data\n",
        "cp -r data/arcDataset/Gothic\\ architecture data/trainB\n",
        "cp -r data/arcDataset/International\\ style data/trainA\n",
        "mkdir data/testA\n",
        "mkdir data/testB"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7t06BJ3-1HpO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "N_EPOCHS = 200\n",
        "DECAY_START_EPOCH = 100\n",
        "BATCH_SIZE = 2\n",
        "DATA_PATH = \"data/\"\n",
        "PATH_TO_SAVE = '/content/drive/My Drive/DLSchool/CycleGAN/final/'\n",
        "# PATH_TO_SAVE = 'data/output/'\n",
        "NUM_WORKERS = 4\n",
        "IMAGE_SIZE = 256\n",
        "LEARNING_RATE = 0.0002"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RvOohKvvTgYj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "os.makedirs(os.path.dirname(PATH_TO_SAVE), exist_ok=True)\n",
        "with open(PATH_TO_SAVE + 'history.txt', 'a') as f:\n",
        "        f.write('loss_gen, loss_disc, loss_identity, loss_gan, loss_cycle')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tToIw_b-1R-W",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# This functions creates test set\n",
        "# and remove grayscale images\n",
        "data.prepare_dataset(DATA_PATH)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1LkgEsTaDn31",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_transforms = [transforms.Resize((600, 800)),\n",
        "                   transforms.ToTensor(),\n",
        "                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ]\n",
        "test_loader = DataLoader(data.ImageDataset(DATA_PATH, transforms_=test_transforms, mode='test', unaligned=False), \n",
        "                        batch_size=1, shuffle=False, num_workers=NUM_WORKERS)\n",
        "# Take one image from test set to display results after each epoch\n",
        "test_iter = iter(test_loader)\n",
        "sample_Y = next(test_iter)['B']\n",
        "sample_X = next(test_iter)['A']"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2AGjc8Fgi1Lb",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def lambda_lr(epoch):\n",
        "    # Linear decay\n",
        "    return 1.0 - max(0, epoch - DECAY_START_EPOCH)/(N_EPOCHS - DECAY_START_EPOCH)\n",
        "\n",
        "def set_requires_grad(nets, requires_grad=False):\n",
        "        \"\"\"Set requies_grad=False for all the networks to avoid unnecessary computations\n",
        "        Parameters:\n",
        "            nets (network list)   -- a list of networks\n",
        "            requires_grad (bool)  -- whether the networks require gradients or not\n",
        "        \"\"\"\n",
        "        if not isinstance(nets, list):\n",
        "            nets = [nets]\n",
        "        for net in nets:\n",
        "            if net is not None:\n",
        "                for param in net.parameters():\n",
        "                    param.requires_grad = requires_grad\n",
        "\n",
        "def plot_history(history):\n",
        "    loss_gen, loss_disc, loss_identity, loss_gan, loss_cycle = zip(*history)\n",
        "    plt.figure(figsize=(20, 5))\n",
        "    plt.subplot(151)\n",
        "    plt.title(\"loss G\")\n",
        "    plt.plot(loss_gen)\n",
        "    plt.subplot(152)\n",
        "    plt.title(\"loss D\")\n",
        "    plt.plot(loss_disc)\n",
        "    plt.subplot(153)\n",
        "    plt.title(\"loss Identity\")\n",
        "    plt.plot(loss_identity)\n",
        "    plt.subplot(154)\n",
        "    plt.title(\"loss GAN\")\n",
        "    plt.plot(loss_gan)\n",
        "    plt.subplot(155)\n",
        "    plt.title(\"loss Cycle\")\n",
        "    plt.plot(loss_cycle)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q4kDPEvMI4Ff",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Initialize models and init weights\n",
        "# Forward X->Y\n",
        "# Backward Y->X\n",
        "gen_forward = nets.Generator().to(device)\n",
        "gen_backward = nets.Generator().to(device)\n",
        "discX = nets.Discriminator().to(device)\n",
        "discY = nets.Discriminator().to(device)\n",
        "\n",
        "gen_forward.apply(functions.weights_init_normal)\n",
        "gen_backward.apply(functions.weights_init_normal)\n",
        "discY.apply(functions.weights_init_normal)\n",
        "discX.apply(functions.weights_init_normal)\n",
        "\n",
        "models = {'netG_A2B': gen_forward,\n",
        "          'netG_B2A': gen_backward,\n",
        "          'netD_A': discY,\n",
        "          'netD_B': discX}\n",
        "\n",
        "ones = torch.FloatTensor(BATCH_SIZE).fill_(1.0).to(device)\n",
        "zeros = torch.FloatTensor(BATCH_SIZE).fill_(0.0).to(device)\n",
        "\n",
        "\n",
        "transforms_ = [ transforms.Resize(int(IMAGE_SIZE*1.12)), \n",
        "                transforms.RandomCrop(IMAGE_SIZE), \n",
        "                transforms.RandomHorizontalFlip(),\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ]\n",
        "\n",
        "dataloader = DataLoader(data.ImageDataset(DATA_PATH, transforms_=transforms_, unaligned=True), \n",
        "                        batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS)\n",
        "\n",
        "criterion_GAN = nn.MSELoss()\n",
        "criterion_cycle = nn.L1Loss()\n",
        "criterion_identity = nn.L1Loss()\n",
        "\n",
        "# Optimizers & LR schedulers\n",
        "gen_optim = optim.Adam(itertools.chain(gen_forward.parameters(), gen_backward.parameters()),\n",
        "                                lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
        "discX_optim = optim.Adam(discX.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
        "discY_optim = optim.Adam(discY.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))\n",
        "\n",
        "lr_scheduler_gen = optim.lr_scheduler.LambdaLR(gen_optim, lr_lambda=lambda_lr)\n",
        "lr_scheduler_discX = optim.lr_scheduler.LambdaLR(discX_optim, lr_lambda=lambda_lr)\n",
        "lr_scheduler_discY = optim.lr_scheduler.LambdaLR(discY_optim, lr_lambda=lambda_lr)\n",
        "\n",
        "\n",
        "###### Training ######\n",
        "history = []\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "    loss_gan_av = 0\n",
        "    loss_cycle_av = 0\n",
        "    loss_identity_av = 0\n",
        "    loss_disc_av = 0\n",
        "    loss_gen_av = 0\n",
        "    with tqdm_notebook(desc=f\"Epoch {epoch}/{N_EPOCHS}\", total=len(dataloader), position=0) as pbar:\n",
        "        for i, batch in enumerate(dataloader):\n",
        "            \n",
        "            real_X = batch['A'].to(device)\n",
        "            real_Y = batch['B'].to(device)\n",
        "\n",
        "            # Generators\n",
        "            set_requires_grad([discX, discY], False)     \n",
        "            gen_optim.zero_grad()\n",
        "\n",
        "            # Identity loss\n",
        "            same_Y = gen_forward(real_Y)\n",
        "            loss_identity_Y = criterion_identity(same_Y, real_Y)\n",
        "\n",
        "            same_X = gen_backward(real_X)\n",
        "            loss_identity_X = criterion_identity(same_X, real_X)\n",
        "\n",
        "            # Forward GAN loss\n",
        "            fake_Y = gen_forward(real_X)\n",
        "            pred_fake = discY(fake_Y)\n",
        "            loss_GAN_X2Y = criterion_GAN(pred_fake, ones) # [(D(G(x)) - 1)^2]\n",
        "\n",
        "            # Backward GAN loss\n",
        "            fake_X = gen_backward(real_Y)\n",
        "            pred_fake = discX(fake_X)\n",
        "            loss_GAN_Y2X = criterion_GAN(pred_fake, ones) # [(D'(F(y)) - 1)^2]\n",
        "\n",
        "            # Cycle loss\n",
        "            recovered_X = gen_backward(fake_Y)\n",
        "            recovered_Y = gen_forward(fake_X)\n",
        "\n",
        "            loss_cycle = criterion_cycle(recovered_X, real_X) + criterion_cycle(recovered_Y, real_Y)\n",
        "\n",
        "            gen_loss = 5.0*loss_identity_X + 5.0*loss_identity_Y + loss_GAN_X2Y + loss_GAN_Y2X + 10.0*loss_cycle\n",
        "            gen_loss.backward()\n",
        "            \n",
        "            gen_optim.step()\n",
        "\n",
        "            # Backward discriminator loss\n",
        "            set_requires_grad([discX, discY], True)   \n",
        "            discX_optim.zero_grad()\n",
        "\n",
        "            loss_D_real = criterion_GAN(discX(real_X), ones)    # [(D'(x) - 1)^2]      \n",
        "            loss_D_fake = criterion_GAN(discX(fake_X.detach()), zeros) # [D'(F(y))^2]\n",
        "\n",
        "            loss_discX = (loss_D_real + loss_D_fake)*0.5\n",
        "            loss_discX.backward()\n",
        "\n",
        "            discX_optim.step()\n",
        "\n",
        "            # Forward discriminator loss\n",
        "            discY_optim.zero_grad()\n",
        "\n",
        "            loss_D_real = criterion_GAN(discY(real_Y), ones)   # [(D(y) - 1)^2]      \n",
        "            loss_D_fake = criterion_GAN(discY(fake_Y.detach()), zeros) # [D(G(x))^2]\n",
        "\n",
        "            loss_discY = (loss_D_real + loss_D_fake)*0.5\n",
        "            loss_discY.backward()\n",
        "\n",
        "            discY_optim.step()\n",
        "\n",
        "            pbar.update(1)\n",
        "\n",
        "            loss_gan_av += loss_GAN_Y2X.item() + loss_GAN_X2Y.item()\n",
        "            loss_cycle_av += loss_cycle.item()\n",
        "            loss_identity_av += loss_identity_X.item() + loss_identity_Y.item()\n",
        "            loss_disc_av += loss_discY.item() + loss_discX.item()\n",
        "            loss_gen_av += gen_loss.item()\n",
        "\n",
        "    loss_gan_av /= len(dataloader)\n",
        "    loss_cycle_av /= len(dataloader)\n",
        "    loss_identity_av /= len(dataloader)\n",
        "    loss_disc_av /= len(dataloader)\n",
        "    loss_gen_av /= len(dataloader)\n",
        "    \n",
        "    history.append([loss_gen_av, loss_disc_av, loss_identity_av,\n",
        "                    loss_gan_av, loss_cycle_av])\n",
        "\n",
        "                \n",
        "    clear_output(wait=True)\n",
        "    plot_history(history)\n",
        "    # Using a consistent image (sample) so that the progress of the model\n",
        "    # is clearly visible.\n",
        "    print(\"Generation A -> B\")\n",
        "    functions.generate_images(gen_forward, sample_X, device)\n",
        "    print(\"Generation B -> A\")\n",
        "    functions.generate_images(gen_backward, sample_Y, device)\n",
        "\n",
        "    # Update learning rates\n",
        "    lr_scheduler_gen.step()\n",
        "    lr_scheduler_discX.step()\n",
        "    lr_scheduler_discY.step()\n",
        "\n",
        "    # Save models checkpoints and history\n",
        "    with open(PATH_TO_SAVE + 'history.txt', 'a') as f:\n",
        "        f.write(','.join([str(i) for i in history[-1]]))\n",
        "    functions.save_models(PATH_TO_SAVE, models)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N5IiDIxi__Ha",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}