DRT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b96a62bc",
   "metadata": {},
   "source": [
    "# Deep Recursive Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9077649",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6f5699a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#torch imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.cuda as cuda\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from timm.models.layers import trunc_normal_\n",
    "from torch.utils.data import DataLoader\n",
    "from torchsummary import summary\n",
    "\n",
    "#other imports\n",
    "import numpy as np\n",
    "import cv2 as cv\n",
    "from PIL import Image\n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "import time\n",
    "from IPython.display import clear_output\n",
    "\n",
    "#own files import\n",
    "from transformer_block import TransformerBlock, PatchEmbed, PatchUnEmbed\n",
    "from Data.data_loader import Rain800TrainData, Rain800ValData\n",
    "from my_utils import batch_PSNR, batch_SSIM, output_to_image\n",
    "from my_utils import save_ckp, load_ckp, base_path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cadbb221",
   "metadata": {},
   "source": [
    "### Global Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8a442fab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "training_image_size = 56\n",
    "dtype = torch.cuda.FloatTensor\n",
    "batch_size = 5\n",
    "torch.manual_seed(1234)\n",
    "torch.cuda.manual_seed_all(1234)\n",
    "epochs = 4600\n",
    "lr = 0.0001\n",
    "error_plot_freq = 20\n",
    "INT_MAX = 2147483647\n",
    "error_tolerence = 10\n",
    "patch_size = 1\n",
    "\n",
    "#paths\n",
    "base_pth = base_path()\n",
    "ckp_pth = base_pth + \"/CheckPoints\"\n",
    "\n",
    "#miscellaneous\n",
    "warnings.filterwarnings('ignore')\n",
    "plt.rcParams['figure.figsize'] = (16, 9)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59a32126",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3346514e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Rain800\n",
      "Comparing with Hi-Net on Rain100L\n"
     ]
    }
   ],
   "source": [
    "### Prepare Data for Training\n",
    "train_dataset = Rain800TrainData(training_image_size, dataset_dir='/Rain-800/') #/Rain100L-Train/\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, drop_last=False)\n",
    "\n",
    "### Prepare synthetic data for validation\n",
    "val_syn_dataset = Rain800ValData(dataset_dir = '/Test-HiNet/Rain100L/', syn = False)\n",
    "val_syn_dataloader = DataLoader(dataset=val_syn_dataset, batch_size=1, pin_memory=False, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61f2715f",
   "metadata": {},
   "source": [
    "### Model Design"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8eb800",
   "metadata": {},
   "source": [
    "We inherit the swin transformer block with a modification: WSA -> CONV2d for our network. Here we design the recursive network such that every transformer blocks in the same residual units share the same weigh. We simply stack these residual units recursiely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5fb51316",
   "metadata": {},
   "outputs": [],
   "source": [
    "#patch embedding -> transformer -> patch unembedding\n",
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, dim, input_resolution, num_heads, patch_size):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.patch_size = patch_size\n",
    "        self.input_resolution = input_resolution\n",
    "        H, W = self.input_resolution\n",
    "        self.num_heads = num_heads\n",
    "        self.transformer = TransformerBlock(self.dim, (H//self.patch_size, W//self.patch_size),\n",
    "                                            num_heads) #patch size is 4\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x=self.transformer(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a5ad2b8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Basic Block -> basic block (with skip connection)\n",
    "class ResidualLayer(nn.Module):\n",
    "    def __init__(self, dim, input_resolution, residual_depth, patch_size):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.patch_size = patch_size\n",
    "        self.residual_depth = residual_depth\n",
    "        self.input_resolution = input_resolution\n",
    "        self.block1 = BasicBlock(self.dim, self.input_resolution, 2, self.patch_size) #multi-heads: 2\n",
    "        self.block2 = BasicBlock(self.dim, self.input_resolution, 2, self.patch_size)\n",
    "        self.conv_out = nn.Conv2d(self.dim, self.dim, 3, padding = 1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        B, HW, C = x.shape\n",
    "        H, W = self.input_resolution\n",
    "        shortcut = x\n",
    "        for _ in range(self.residual_depth):\n",
    "            x = self.block1(self.block2(x))\n",
    "            x = torch.add(x, shortcut)\n",
    "        #convolution at the end of each residual block\n",
    "        x = x.transpose(1,2).view(B, C, H//self.patch_size, W//self.patch_size)\n",
    "        x = self.conv_out(x).flatten(2).transpose(1,2)#B L C\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "594801ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "#recursive network based on residual units\n",
    "class DeepRecursiveTransformer(nn.Module):\n",
    "    def __init__(self, dim, input_resolution, patch_size, residual_depth, recursive_depth):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.patch_size = patch_size\n",
    "        self.recursive_depth = recursive_depth\n",
    "        self.input_resolution = input_resolution\n",
    "        self.residual_depth = residual_depth\n",
    "        self.H, self.W = self.input_resolution\n",
    "        assert self.H == self.W, \"Input hight and width should be the same\"\n",
    "        self.input_conv1 = nn.Conv2d(3, self.dim, 3, padding=1)\n",
    "        self.patch_embed = PatchEmbed(img_size=self.H, patch_size = self.patch_size,\n",
    "                                      in_chans=3, embed_dim=self.dim)\n",
    "        self.patch_unembed = PatchUnEmbed(img_size=self.H, patch_size = self.patch_size,\n",
    "                                          in_chans=self.dim, unembed_dim=3)\n",
    "        self.recursive_layers = nn.ModuleList()\n",
    "        for i in range(self.recursive_depth):\n",
    "            layer = ResidualLayer(self.dim, self.input_resolution, self.residual_depth, self.patch_size)\n",
    "            self.recursive_layers.append(layer)\n",
    "        self.output_conv1 = nn.Conv2d(self.dim, 3, 3, padding=1)\n",
    "        #use imagenet mean and std for general domain normalisation\n",
    "        self.normalise_layer = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))\n",
    "        self.denormalise_layer = transforms.Normalize((-0.485, -0.456, -0.406), (1./0.229, 1./0.224, 1./0.225))\n",
    "        self.apply(self._init_weights)\n",
    "        self.activation = nn.LeakyReLU()\n",
    "        \n",
    "    #weight initialisation scheme\n",
    "    def _init_weights(self, l):\n",
    "        if isinstance(l, nn.Linear):\n",
    "            trunc_normal_(l.weight, std=.02)\n",
    "            if isinstance(l, nn.Linear) and l.bias is not None:\n",
    "                nn.init.constant_(l.bias, 0)\n",
    "        elif isinstance(l, nn.LayerNorm):\n",
    "            nn.init.constant_(l.bias, 0)\n",
    "            nn.init.constant_(l.weight, 1.0)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        #normalise the data, input shape (B, C, H, W)\n",
    "        x = self.normalise_layer(x)\n",
    "        outer_shortcut = x\n",
    "        x = self.patch_embed(x)\n",
    "        inner_shortcut = x\n",
    "\n",
    "        for i in range(len(self.recursive_layers)):\n",
    "            x = self.recursive_layers[i](x)\n",
    "            \n",
    "        x=torch.add(x, inner_shortcut)\n",
    "        x=self.patch_unembed(x, (self.H//self.patch_size,self.W//self.patch_size))\n",
    "        x=torch.add(x, outer_shortcut)\n",
    "        x=self.denormalise_layer(x)\n",
    "        return x #output shape (B, C, H, W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2278a21f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===============================================================================================\n",
      "Layer (type:depth-idx)                        Output Shape              Param #\n",
      "===============================================================================================\n",
      "├─Normalize: 1-1                              [-1, 3, 56, 56]           --\n",
      "├─PatchEmbed: 1-2                             [-1, 3136, 96]            --\n",
      "|    └─Conv2d: 2-1                            [-1, 96, 56, 56]          384\n",
      "├─ModuleList: 1                               []                        --\n",
      "|    └─ResidualLayer: 2-2                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-1                   [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-2                   [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-3                   [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-4                   [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-5                   [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-6                   [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-7                       [-1, 96, 56, 56]          83,040\n",
      "|    └─ResidualLayer: 2-3                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-8                   [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-9                   [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-10                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-11                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-12                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-13                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-14                      [-1, 96, 56, 56]          83,040\n",
      "|    └─ResidualLayer: 2-4                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-15                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-16                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-17                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-18                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-19                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-20                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-21                      [-1, 96, 56, 56]          83,040\n",
      "|    └─ResidualLayer: 2-5                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-22                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-23                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-24                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-25                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-26                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-27                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-28                      [-1, 96, 56, 56]          83,040\n",
      "|    └─ResidualLayer: 2-6                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-29                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-30                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-31                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-32                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-33                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-34                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-35                      [-1, 96, 56, 56]          83,040\n",
      "|    └─ResidualLayer: 2-7                     [-1, 3136, 96]            --\n",
      "|    |    └─BasicBlock: 3-36                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-37                  [-1, 3136, 96]            56,594\n",
      "|    |    └─BasicBlock: 3-38                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-39                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-40                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─BasicBlock: 3-41                  [-1, 3136, 96]            (recursive)\n",
      "|    |    └─Conv2d: 3-42                      [-1, 96, 56, 56]          83,040\n",
      "├─PatchUnEmbed: 1-3                           [-1, 3, 56, 56]           --\n",
      "|    └─Upsample: 2-8                          [-1, 96, 56, 56]          --\n",
      "|    └─Conv2d: 2-9                            [-1, 3, 56, 56]           2,595\n",
      "├─Normalize: 1-4                              [-1, 3, 56, 56]           --\n",
      "===============================================================================================\n",
      "Total params: 1,180,347\n",
      "Trainable params: 1,180,347\n",
      "Non-trainable params: 0\n",
      "Total mult-adds (G): 1.57\n",
      "===============================================================================================\n",
      "Input size (MB): 0.04\n",
      "Forward/backward pass size (MB): 16.15\n",
      "Params size (MB): 4.50\n",
      "Estimated Total Size (MB): 20.69\n",
      "===============================================================================================\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "===============================================================================================\n",
       "Layer (type:depth-idx)                        Output Shape              Param #\n",
       "===============================================================================================\n",
       "├─Normalize: 1-1                              [-1, 3, 56, 56]           --\n",
       "├─PatchEmbed: 1-2                             [-1, 3136, 96]            --\n",
       "|    └─Conv2d: 2-1                            [-1, 96, 56, 56]          384\n",
       "├─ModuleList: 1                               []                        --\n",
       "|    └─ResidualLayer: 2-2                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-1                   [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-2                   [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-3                   [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-4                   [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-5                   [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-6                   [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-7                       [-1, 96, 56, 56]          83,040\n",
       "|    └─ResidualLayer: 2-3                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-8                   [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-9                   [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-10                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-11                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-12                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-13                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-14                      [-1, 96, 56, 56]          83,040\n",
       "|    └─ResidualLayer: 2-4                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-15                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-16                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-17                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-18                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-19                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-20                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-21                      [-1, 96, 56, 56]          83,040\n",
       "|    └─ResidualLayer: 2-5                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-22                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-23                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-24                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-25                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-26                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-27                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-28                      [-1, 96, 56, 56]          83,040\n",
       "|    └─ResidualLayer: 2-6                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-29                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-30                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-31                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-32                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-33                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-34                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-35                      [-1, 96, 56, 56]          83,040\n",
       "|    └─ResidualLayer: 2-7                     [-1, 3136, 96]            --\n",
       "|    |    └─BasicBlock: 3-36                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-37                  [-1, 3136, 96]            56,594\n",
       "|    |    └─BasicBlock: 3-38                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-39                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-40                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─BasicBlock: 3-41                  [-1, 3136, 96]            (recursive)\n",
       "|    |    └─Conv2d: 3-42                      [-1, 96, 56, 56]          83,040\n",
       "├─PatchUnEmbed: 1-3                           [-1, 3, 56, 56]           --\n",
       "|    └─Upsample: 2-8                          [-1, 96, 56, 56]          --\n",
       "|    └─Conv2d: 2-9                            [-1, 3, 56, 56]           2,595\n",
       "├─Normalize: 1-4                              [-1, 3, 56, 56]           --\n",
       "===============================================================================================\n",
       "Total params: 1,180,347\n",
       "Trainable params: 1,180,347\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (G): 1.57\n",
       "===============================================================================================\n",
       "Input size (MB): 0.04\n",
       "Forward/backward pass size (MB): 16.15\n",
       "Params size (MB): 4.50\n",
       "Estimated Total Size (MB): 20.69\n",
       "==============================================================================================="
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = DeepRecursiveTransformer(96, (training_image_size, training_image_size), patch_size, 3,6)\n",
    "summary(net.cuda(), (3, training_image_size, training_image_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55818e06",
   "metadata": {},
   "source": [
    "### Training Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5f74a6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss().type(dtype)\n",
    "optimiser = optim.Adam(net.parameters(), lr=lr)\n",
    "net = net.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a4edf321",
   "metadata": {},
   "outputs": [],
   "source": [
    "#graph network error\n",
    "def graph_error(error_list, name):\n",
    "    if name[-4:] != \".png\":\n",
    "        if name != \"\":\n",
    "            raise Exception(\"Suffix of file type is needed\")\n",
    "    save_dir = \"Losses/\" + name\n",
    "    x = np.arange(len(error_list))\n",
    "    y = np.asarray(error_list)\n",
    "    plt.plot(x, y)\n",
    "    plt.ylabel(\"Error\")\n",
    "    plt.xlabel(\"Epoches\")\n",
    "    if name != \"\":\n",
    "        plt.savefig(save_dir)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70b41df5",
   "metadata": {},
   "source": [
    "### Network Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "eea27bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def network_training(net, optimiser, criterion, loadCkp = False, loadBest=True, new_dataset=False):\n",
    "    error_list = []\n",
    "    start_epoch = 0\n",
    "    best_model_saved = False\n",
    "    ckp_saved = False\n",
    "    previous_batch_error = INT_MAX #initialise to a large value\n",
    "    best_error = INT_MAX\n",
    "    ###load checkpoint if required\n",
    "    if loadCkp and loadBest:\n",
    "        best_model_saved = True\n",
    "        ckp_saved = True\n",
    "        #when training on a new dataset for the first time, we only load the network itself\n",
    "        if new_dataset:\n",
    "            net, _, _, _, _ = load_ckp(ckp_pth+\"/best_model.pt\", net, optimiser)\n",
    "            print(\"Finished loading the best model, ignored the training history\")\n",
    "        else:\n",
    "            net, optimiser, start_epoch, error_list, best_error = load_ckp(ckp_pth+\"/best_model.pt\", net, optimiser)\n",
    "            print(\"Finished loading the best model\")\n",
    "            previous_batch_error = best_error\n",
    "    elif loadCkp and not loadBest:\n",
    "        ckp_saved = True\n",
    "        if new_dataset:\n",
    "            net, _, _, _, _ = load_ckp(ckp_pth+\"/checkpoint.pt\", net, optimiser)\n",
    "            print(\"Finished loading the checkpoint, ignored the training history\")\n",
    "        else:\n",
    "            net, optimiser, start_epoch, error_list, best_error = load_ckp(ckp_pth+\"/checkpoint.pt\", net, optimiser)\n",
    "            print(\"Finished loading the checkpoint\")\n",
    "            previous_batch_error = best_error\n",
    "    \n",
    "    if best_error == None:\n",
    "        best_error = INT_MAX\n",
    "    \n",
    "    for epoch in range(start_epoch, epochs):\n",
    "        batch_error = 0\n",
    "        epoch_start_time = time.time()\n",
    "        \n",
    "        ### iterate through the batches\n",
    "        for i, data in enumerate(train_loader, 0):\n",
    "            optimiser.zero_grad()\n",
    "            target = data[0].cuda()\n",
    "            net_input = data[1].cuda()\n",
    "            net_output = net(net_input)\n",
    "            loss = criterion(net_output, target)\n",
    "            batch_error += loss.item()\n",
    "            loss.backward()\n",
    "            optimiser.step()\n",
    "        \n",
    "        ### find one epoch training time\n",
    "        one_epoch_time = time.time() - epoch_start_time\n",
    "        print(\"One epoch time: \" + str(one_epoch_time))\n",
    "        \n",
    "        ### process the error information\n",
    "        print('[%d] loss: %.3f' %(epoch + 1, batch_error))\n",
    "        ### if error is too large, roll back, otherwise save and continue\n",
    "        if batch_error > error_tolerence*previous_batch_error and (best_model_saved or ckp_saved):\n",
    "            if ckp_saved:\n",
    "                print(\"Current error is too large, loading the last checkpoint\")\n",
    "                net, optimiser, start_epoch, error_list, best_psnr = \\\n",
    "                    load_ckp(ckp_pth+\"/checkpoint.pt\", net, optimiser)\n",
    "            elif best_model_saved:\n",
    "                print(\"Current error is too large, loading the best model\")\n",
    "                net, optimiser, start_epoch, error_list, best_psnr = \\\n",
    "                    load_ckp(ckp_pth+\"/best_model.pt\", net, optimiser)\n",
    "            else:\n",
    "                raise Exception(\"Error is too large, but no models to load\")\n",
    "        else:\n",
    "            if batch_error > error_tolerence*previous_batch_error:\n",
    "                print(\"Current error is too large, but cannot roll back\")\n",
    "            else:\n",
    "                previous_batch_error = batch_error\n",
    "                \n",
    "            error_list.append(batch_error)\n",
    "            ###save the latest model\n",
    "            checkpoint = {\n",
    "                'epoch': epoch + 1,\n",
    "                'state_dict': net.state_dict(),\n",
    "                'optimizer': optimiser.state_dict(),\n",
    "                'error_list': error_list,\n",
    "                'best_error': best_error\n",
    "            }\n",
    "            save_ckp(checkpoint, False, ckp_pth)\n",
    "            ckp_saved = True\n",
    "            \n",
    "            ###if error is the smallest save it as the best model\n",
    "            if batch_error < best_error:\n",
    "                best_error = batch_error\n",
    "                checkpoint = {\n",
    "                    'epoch': epoch + 1,\n",
    "                    'state_dict': net.state_dict(),\n",
    "                    'optimizer': optimiser.state_dict(),\n",
    "                    'error_list': error_list,\n",
    "                    'best_error': best_error\n",
    "                }\n",
    "                save_ckp(checkpoint, True, ckp_pth)\n",
    "                best_model_saved = True\n",
    "                print(\"New Minimum Error Recorded!\")\n",
    "                \n",
    "            if ((epoch+1) % error_plot_freq) == 0 or epoch == epochs-1:\n",
    "                graph_error(error_list[1:], \"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c94a4116",
   "metadata": {},
   "source": [
    "### Network Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "409e56cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def network_validation(loadBest = True, display_img = False, display_all_img = False, deleted_indices = []):\n",
    "    psnr_list = []\n",
    "    ssim_list = []\n",
    "    local_window_dim = patch_size*7 #patch_size * window_size\n",
    "    deleted_indices = []\n",
    "    if display_img:\n",
    "        if display_all_img:\n",
    "            random_plot_list = np.arange(0,100)\n",
    "            if deleted_indices != []:\n",
    "                random_plot_list = np.delete(random_plot_list, deleted_indices)\n",
    "                print(random_plot_list)\n",
    "        else:\n",
    "            random_plot_list = np.random.randint(0,100, size = 20)\n",
    "    \n",
    "    start_time = time.time()\n",
    "    ### iterate through the validation dataset\n",
    "    for j, val_data in enumerate(val_syn_dataloader, 0):\n",
    "        if(j not in deleted_indices):\n",
    "            print(f\"Evaluating the {j}th sample: \")\n",
    "            ### pad the image to make sure H==W fits the network requirements\n",
    "            net_input = val_data[1]\n",
    "            _, _, h_old, w_old = net_input.size()\n",
    "            h_original = h_old\n",
    "            w_original = w_old\n",
    "            multiplier = max(h_old // local_window_dim + 1, w_old // local_window_dim + 1)\n",
    "            h_pad = (multiplier) * local_window_dim - h_old\n",
    "            w_pad = (multiplier) * local_window_dim - w_old\n",
    "            net_input = torch.cat([net_input, torch.flip(net_input, [2])], 2)[:, :, :h_old + h_pad, :]\n",
    "            net_input = torch.cat([net_input, torch.flip(net_input, [3])], 3)[:, :, :, :w_old + w_pad]\n",
    "            \n",
    "            ## pad again if h/w or w/h ratio is bigger than 2\n",
    "            if h_pad > h_old or w_pad > w_old:\n",
    "                _, _, h_old, w_old = net_input.size()\n",
    "                multiplier = max(h_old // local_window_dim + 1, w_old // local_window_dim + 1)\n",
    "                h_pad = (multiplier) * local_window_dim - h_old\n",
    "                w_pad = (multiplier) * local_window_dim - w_old\n",
    "                net_input = torch.cat([net_input, torch.flip(net_input, [2])], 2)[:, :, :h_old + h_pad, :]\n",
    "                net_input = torch.cat([net_input, torch.flip(net_input, [3])], 3)[:, :, :, :w_old + w_pad]\n",
    "\n",
    "            ### evaluate \n",
    "            _, _, new_h, new_w = net_input.size()\n",
    "            assert new_h == new_w, \"Input image should have square dimension\"\n",
    "            #FIXME: change the fixed parameter below into adaptive variables\n",
    "            eval_net = DeepRecursiveTransformer(96, (new_h, new_w), patch_size, 3,6)\n",
    "            if loadBest:\n",
    "                eval_net,_,_,_,_ = load_ckp(ckp_pth+\"/best_model.pt\", eval_net, optimiser)\n",
    "            else:\n",
    "                eval_net,_,_,_,_ = load_ckp(ckp_pth+\"/checkpoint.pt\", eval_net, optimiser)\n",
    "\n",
    "            net_input = net_input.cuda()\n",
    "            target = val_data[0].cuda()\n",
    "            eval_net = eval_net.to(device)\n",
    "            eval_net.eval()\n",
    "            with torch.no_grad():\n",
    "                net_output = eval_net(net_input)\n",
    "                \n",
    "            ### crop the output\n",
    "            net_output = net_output[:,:,:h_original, :w_original]\n",
    "            output_data = net_output.cpu().detach().numpy() #B C H W\n",
    "            output_data = np.transpose(output_data, (0,2,3,1)) #B H W C\n",
    "            target_data = target.cpu().detach().numpy()\n",
    "            target_data = np.transpose(target_data, (0,2,3,1))\n",
    "            input_data = val_data[1].cpu().detach().numpy()\n",
    "            input_data = np.transpose(input_data, (0,2,3,1))\n",
    "            \n",
    "            #display\n",
    "            if display_img:\n",
    "                if (j in random_plot_list) or j == 0:\n",
    "                    output_to_image(target_data, output_data, input_data, save_img = True, save_index = j)\n",
    "            \n",
    "            \n",
    "            ### calculate the PSNR and SSIM\n",
    "            psnr_list.append(batch_PSNR(target_data, output_data, True).item())\n",
    "            ssim_list.append(batch_SSIM(target_data, output_data).item())\n",
    "            cuda.empty_cache()\n",
    "            \n",
    "    run_time = time.time() - start_time\n",
    "    avg_psnr = sum(psnr_list)/len(psnr_list)\n",
    "    avg_ssim = sum(ssim_list)/len(ssim_list)\n",
    "    print(f\"Average PSNR is {avg_psnr}\")\n",
    "    print(f\"Average SSIM is {avg_ssim}\")\n",
    "    print(f\"Validation time is {run_time}\")\n",
    "    return avg_psnr, avg_ssim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3e976005",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(train_net = True, loadCkp = True, loadBest = True, new_dataset=False):\n",
    "    if train_net:\n",
    "        print(\"Training the network...\")\n",
    "        network_training(net, optimiser, criterion, loadCkp=loadCkp, loadBest=loadBest, new_dataset=new_dataset)\n",
    "    else:\n",
    "        print(\"Evaluating the network...\")\n",
    "        network_validation(display_img = False, display_all_img = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "12ec3353",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluating the network...\n",
      "Evaluating the 0th sample: \n",
      "Evaluating the 1th sample: \n",
      "Evaluating the 2th sample: \n",
      "Evaluating the 3th sample: \n",
      "Evaluating the 4th sample: \n",
      "Evaluating the 5th sample: \n",
      "Evaluating the 6th sample: \n",
      "Evaluating the 7th sample: \n",
      "Evaluating the 8th sample: \n",
      "Evaluating the 9th sample: \n",
      "Evaluating the 10th sample: \n",
      "Evaluating the 11th sample: \n",
      "Evaluating the 12th sample: \n",
      "Evaluating the 13th sample: \n",
      "Evaluating the 14th sample: \n",
      "Evaluating the 15th sample: \n",
      "Evaluating the 16th sample: \n",
      "Evaluating the 17th sample: \n",
      "Evaluating the 18th sample: \n",
      "Evaluating the 19th sample: \n",
      "Evaluating the 20th sample: \n",
      "Evaluating the 21th sample: \n",
      "Evaluating the 22th sample: \n",
      "Evaluating the 23th sample: \n",
      "Evaluating the 24th sample: \n",
      "Evaluating the 25th sample: \n",
      "Evaluating the 26th sample: \n",
      "Evaluating the 27th sample: \n",
      "Evaluating the 28th sample: \n",
      "Evaluating the 29th sample: \n",
      "Evaluating the 30th sample: \n",
      "Evaluating the 31th sample: \n",
      "Evaluating the 32th sample: \n",
      "Evaluating the 33th sample: \n",
      "Evaluating the 34th sample: \n",
      "Evaluating the 35th sample: \n",
      "Evaluating the 36th sample: \n",
      "Evaluating the 37th sample: \n",
      "Evaluating the 38th sample: \n",
      "Evaluating the 39th sample: \n",
      "Evaluating the 40th sample: \n",
      "Evaluating the 41th sample: \n",
      "Evaluating the 42th sample: \n",
      "Evaluating the 43th sample: \n",
      "Evaluating the 44th sample: \n",
      "Evaluating the 45th sample: \n",
      "Evaluating the 46th sample: \n",
      "Evaluating the 47th sample: \n",
      "Evaluating the 48th sample: \n",
      "Evaluating the 49th sample: \n",
      "Evaluating the 50th sample: \n",
      "Evaluating the 51th sample: \n",
      "Evaluating the 52th sample: \n",
      "Evaluating the 53th sample: \n",
      "Evaluating the 54th sample: \n",
      "Evaluating the 55th sample: \n",
      "Evaluating the 56th sample: \n",
      "Evaluating the 57th sample: \n",
      "Evaluating the 58th sample: \n",
      "Evaluating the 59th sample: \n",
      "Evaluating the 60th sample: \n",
      "Evaluating the 61th sample: \n",
      "Evaluating the 62th sample: \n",
      "Evaluating the 63th sample: \n",
      "Evaluating the 64th sample: \n",
      "Evaluating the 65th sample: \n",
      "Evaluating the 66th sample: \n",
      "Evaluating the 67th sample: \n",
      "Evaluating the 68th sample: \n",
      "Evaluating the 69th sample: \n",
      "Evaluating the 70th sample: \n",
      "Evaluating the 71th sample: \n",
      "Evaluating the 72th sample: \n",
      "Evaluating the 73th sample: \n",
      "Evaluating the 74th sample: \n",
      "Evaluating the 75th sample: \n",
      "Evaluating the 76th sample: \n",
      "Evaluating the 77th sample: \n",
      "Evaluating the 78th sample: \n",
      "Evaluating the 79th sample: \n",
      "Evaluating the 80th sample: \n",
      "Evaluating the 81th sample: \n",
      "Evaluating the 82th sample: \n",
      "Evaluating the 83th sample: \n",
      "Evaluating the 84th sample: \n",
      "Evaluating the 85th sample: \n",
      "Evaluating the 86th sample: \n",
      "Evaluating the 87th sample: \n",
      "Evaluating the 88th sample: \n",
      "Evaluating the 89th sample: \n",
      "Evaluating the 90th sample: \n",
      "Evaluating the 91th sample: \n",
      "Evaluating the 92th sample: \n",
      "Evaluating the 93th sample: \n",
      "Evaluating the 94th sample: \n",
      "Evaluating the 95th sample: \n",
      "Evaluating the 96th sample: \n",
      "Evaluating the 97th sample: \n",
      "Evaluating the 98th sample: \n",
      "Evaluating the 99th sample: \n",
      "Average PSNR is 37.605031775301455\n",
      "Average SSIM is 0.947824544404888\n",
      "Validation time is 163.91107869148254\n"
     ]
    }
   ],
   "source": [
    "run(train_net = False, loadCkp = True, loadBest = True, new_dataset = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a4870dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}