train_SRCNN_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_model import Net\n",
    "\n",
    "torch.manual_seed(1)\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SRCNN parameters\n",
    "\n",
    "batch_size = 64\n",
    "epochs = 10\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-low-res-56/train'\n",
    "#img_path_ref = '/media/angelo/DATEN/Datasets/Experiment_Masters/300W-3D-low-res-224/train'\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=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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=50, gamma=0.1)"
   ]
  },
  {
   "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 100 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_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_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.0221 / PSNR: 16.5473 / SSIM 0.5338\n",
      "===> Epoch 2 Complete: Avg. Loss: 0.0061 / PSNR: 22.1284 / SSIM 0.6462\n",
      "===> Epoch 3 Complete: Avg. Loss: 0.0054 / PSNR: 22.6797 / SSIM 0.7002\n",
      "===> Epoch 4 Complete: Avg. Loss: 0.0048 / PSNR: 23.2195 / SSIM 0.7111\n",
      "===> Epoch 5 Complete: Avg. Loss: 0.0042 / PSNR: 23.7177 / SSIM 0.7193\n",
      "Checkpoint saved to models/SRCNN_x4_epoch_5.pth\n",
      "===> Epoch 6 Complete: Avg. Loss: 0.0041 / PSNR: 23.8610 / SSIM 0.7135\n",
      "===> Epoch 7 Complete: Avg. Loss: 0.0040 / PSNR: 23.9260 / SSIM 0.7159\n",
      "===> Epoch 8 Complete: Avg. Loss: 0.0040 / PSNR: 23.9567 / SSIM 0.7134\n",
      "===> Epoch 9 Complete: Avg. Loss: 0.0041 / PSNR: 23.9194 / SSIM 0.7416\n",
      "===> Epoch 10 Complete: Avg. Loss: 0.0039 / PSNR: 24.0396 / SSIM 0.7049\n",
      "Checkpoint saved to models/SRCNN_x4_epoch_10.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
}