train_SRCNN_coord_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import log10\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from dataset.data_loader_YCbCr_resize import *\n",
    "from utils.pytorch_ssim import *\n",
    "\n",
    "from SRCNN_coord_model import Net\n",
    "\n",
    "torch.manual_seed(1)\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SRCNN CoordConv parameters\n",
    "\n",
    "batch_size = 32\n",
    "epochs = 40\n",
    "lr = 0.01\n",
    "threads = 4\n",
    "upscale_factor = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#img_path_low = '/media/angelo/DATEN/Datasets/Experiment_Masters/300W-3D-crap-56/train'\n",
    "#img_path_ref = '/media/angelo/DATEN/Datasets/Experiment_Masters/300W-3D-low-res-224/train'\n",
    "\n",
    "img_path_low = '/media/angelo/DATEN/Datasets/CelebA/LR_56/test/'\n",
    "img_path_ref = '/media/angelo/DATEN/Datasets/CelebA/HR/test/'\n",
    "\n",
    "#img_path_low = '../dataset/300W-3D-crap-56/train'\n",
    "#img_path_ref = '../dataset/300W-3D-low-res-224/train'\n",
    "\n",
    "train_set = DatasetSuperRes(img_path_low, img_path_ref)\n",
    "training_data_loader = DataLoader(dataset=train_set, num_workers=threads, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===> Building model\n"
     ]
    }
   ],
   "source": [
    "print('===> Building model')\n",
    "model = Net().to(device)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_path = 'results/'\n",
    "out_model_path = 'models/'\n",
    "\n",
    "if not os.path.exists(out_path):\n",
    "    os.makedirs(out_path)    \n",
    "\n",
    "if not os.path.exists(out_model_path):\n",
    "    os.makedirs(out_model_path)   \n",
    "    \n",
    "results = {'avg_loss': [], 'psnr': [], 'ssim': []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    epoch_loss = 0\n",
    "    model.train()\n",
    "    for iteration, batch in enumerate(training_data_loader, 1):\n",
    "        input_, target = batch[0].to(device), batch[1].to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        upsampled_img = model(input_)\n",
    "        loss = criterion(upsampled_img, target)\n",
    "        epoch_loss += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        #print(\"===> Epoch[{}]({}/{}): Loss: {:.4f}\".format(epoch, iteration, len(training_data_loader), loss.item()))\n",
    "    \n",
    "    scheduler.step() # Decrease learning rate after 10 epochs to 10% of its value\n",
    "    \n",
    "    psnr_epoch = 10*log10(1/(epoch_loss / len(training_data_loader)))\n",
    "    ssim_epoch = ssim(upsampled_img, target).item()\n",
    "    avg_loss_batch = epoch_loss/len(training_data_loader)\n",
    "    \n",
    "    results['psnr'].append(psnr_epoch)\n",
    "    results['ssim'].append(ssim_epoch)\n",
    "    results['avg_loss'].append(avg_loss_batch)\n",
    "    \n",
    "    print(\"===> Epoch {} Complete: Avg. Loss: {:.4f} / PSNR: {:.4f} / SSIM {:.4f}\".format(epoch, \n",
    "                                                                                          avg_loss_batch, \n",
    "                                                                                          psnr_epoch,\n",
    "                                                                                          ssim_epoch))\n",
    "    if epoch % (epochs // 2) == 0:\n",
    "    \n",
    "        data_frame = pd.DataFrame(\n",
    "                data={'Avg. Loss': results['avg_loss'], 'PSNR': results['psnr'], 'SSIM': results['ssim']},\n",
    "                index=range(1, epoch + 1))\n",
    "\n",
    "        data_frame.to_csv(out_path + 'SRCNN_coord_x' + str(upscale_factor) + '_train_results.csv', index_label='Epoch')\n",
    "        \n",
    "        checkpoint(epoch)\n",
    "    \n",
    "def checkpoint(epoch):\n",
    "    path = out_model_path + \"SRCNN_coord_x{}_epoch_{}.pth\".format(upscale_factor, epoch)\n",
    "    torch.save(model, path)\n",
    "    print(\"Checkpoint saved to {}\".format(path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===> Epoch 1 Complete: Avg. Loss: 0.4563 / PSNR: 3.4078 / SSIM 0.7507\n",
      "===> Epoch 2 Complete: Avg. Loss: 0.0037 / PSNR: 24.3258 / SSIM 0.7805\n",
      "===> Epoch 3 Complete: Avg. Loss: 0.0029 / PSNR: 25.4447 / SSIM 0.7894\n",
      "===> Epoch 4 Complete: Avg. Loss: 0.0027 / PSNR: 25.6612 / SSIM 0.7952\n",
      "===> Epoch 5 Complete: Avg. Loss: 0.0026 / PSNR: 25.9109 / SSIM 0.7951\n",
      "===> Epoch 6 Complete: Avg. Loss: 0.0020 / PSNR: 26.8866 / SSIM 0.8041\n",
      "===> Epoch 7 Complete: Avg. Loss: 0.0023 / PSNR: 26.3425 / SSIM 0.8084\n",
      "===> Epoch 8 Complete: Avg. Loss: 0.0019 / PSNR: 27.1558 / SSIM 0.8129\n",
      "===> Epoch 9 Complete: Avg. Loss: 0.0019 / PSNR: 27.1963 / SSIM 0.8154\n",
      "===> Epoch 10 Complete: Avg. Loss: 0.0021 / PSNR: 26.7777 / SSIM 0.8149\n",
      "===> Epoch 11 Complete: Avg. Loss: 0.0015 / PSNR: 28.1196 / SSIM 0.8166\n",
      "===> Epoch 12 Complete: Avg. Loss: 0.0015 / PSNR: 28.1856 / SSIM 0.8174\n",
      "===> Epoch 13 Complete: Avg. Loss: 0.0015 / PSNR: 28.2297 / SSIM 0.8180\n",
      "===> Epoch 14 Complete: Avg. Loss: 0.0015 / PSNR: 28.2714 / SSIM 0.8184\n",
      "===> Epoch 15 Complete: Avg. Loss: 0.0015 / PSNR: 28.3109 / SSIM 0.8188\n",
      "===> Epoch 16 Complete: Avg. Loss: 0.0015 / PSNR: 28.3479 / SSIM 0.8191\n",
      "===> Epoch 17 Complete: Avg. Loss: 0.0015 / PSNR: 28.3832 / SSIM 0.8193\n",
      "===> Epoch 18 Complete: Avg. Loss: 0.0014 / PSNR: 28.4154 / SSIM 0.8195\n",
      "===> Epoch 19 Complete: Avg. Loss: 0.0014 / PSNR: 28.4452 / SSIM 0.8195\n",
      "===> Epoch 20 Complete: Avg. Loss: 0.0014 / PSNR: 28.4714 / SSIM 0.8195\n",
      "Checkpoint saved to models/SRCNN_coord_x4_epoch_20.pth\n",
      "===> Epoch 21 Complete: Avg. Loss: 0.0014 / PSNR: 28.4906 / SSIM 0.8195\n",
      "===> Epoch 22 Complete: Avg. Loss: 0.0014 / PSNR: 28.4959 / SSIM 0.8195\n",
      "===> Epoch 23 Complete: Avg. Loss: 0.0014 / PSNR: 28.5010 / SSIM 0.8195\n",
      "===> Epoch 24 Complete: Avg. Loss: 0.0014 / PSNR: 28.5061 / SSIM 0.8195\n",
      "===> Epoch 25 Complete: Avg. Loss: 0.0014 / PSNR: 28.5114 / SSIM 0.8196\n",
      "===> Epoch 26 Complete: Avg. Loss: 0.0014 / PSNR: 28.5168 / SSIM 0.8196\n",
      "===> Epoch 27 Complete: Avg. Loss: 0.0014 / PSNR: 28.5222 / SSIM 0.8196\n",
      "===> Epoch 28 Complete: Avg. Loss: 0.0014 / PSNR: 28.5277 / SSIM 0.8196\n",
      "===> Epoch 29 Complete: Avg. Loss: 0.0014 / PSNR: 28.5333 / SSIM 0.8196\n",
      "===> Epoch 30 Complete: Avg. Loss: 0.0014 / PSNR: 28.5389 / SSIM 0.8196\n",
      "===> Epoch 31 Complete: Avg. Loss: 0.0014 / PSNR: 28.5434 / SSIM 0.8197\n",
      "===> Epoch 32 Complete: Avg. Loss: 0.0014 / PSNR: 28.5447 / SSIM 0.8197\n",
      "===> Epoch 33 Complete: Avg. Loss: 0.0014 / PSNR: 28.5459 / SSIM 0.8197\n",
      "===> Epoch 34 Complete: Avg. Loss: 0.0014 / PSNR: 28.5472 / SSIM 0.8197\n",
      "===> Epoch 35 Complete: Avg. Loss: 0.0014 / PSNR: 28.5485 / SSIM 0.8197\n",
      "===> Epoch 36 Complete: Avg. Loss: 0.0014 / PSNR: 28.5498 / SSIM 0.8197\n",
      "===> Epoch 37 Complete: Avg. Loss: 0.0014 / PSNR: 28.5512 / SSIM 0.8197\n",
      "===> Epoch 38 Complete: Avg. Loss: 0.0014 / PSNR: 28.5526 / SSIM 0.8197\n",
      "===> Epoch 39 Complete: Avg. Loss: 0.0014 / PSNR: 28.5541 / SSIM 0.8197\n",
      "===> Epoch 40 Complete: Avg. Loss: 0.0014 / PSNR: 28.5557 / SSIM 0.8197\n",
      "Checkpoint saved to models/SRCNN_coord_x4_epoch_40.pth\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(1, epochs + 1):\n",
    "    train(epoch)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}