SISR_Mv2_DNS_3c.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "GeForce RTX 2080 Ti\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "from torchvision.datasets import MNIST\n",
    "import torchvision.transforms.functional as TF\n",
    "from torch.optim import lr_scheduler\n",
    "import time\n",
    "import os\n",
    "import cv2\n",
    "import pdb\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt \n",
    "cuda = torch.cuda.get_device_name(0)\n",
    "print(cuda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "/home/ppant/Desktop/Pranshu/Lab/LES-DNS/SR-DNS_Net/k_21\n"
     ]
    }
   ],
   "source": [
    "print(os.getcwd())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('save_3c'):\n",
    "    os.mkdir('save_3c')\n",
    "\n",
    "if not os.path.exists('model_3c'):\n",
    "    os.mkdir('model_3c')\n",
    "\n",
    "if not os.path.exists('test_3c'):\n",
    "    os.mkdir('test_3c')\n",
    "\n",
    "def to_img(x):\n",
    "    x = 0.5 * (x + 1)\n",
    "    x = x.clamp(0, 1)\n",
    "    x = x.view(x.size(0), -1, 64, 64)\n",
    "    return x\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "10000\n10000\n40000\nDNS-LES_3C/les_3c/10.0/464.png\n"
     ]
    }
   ],
   "source": [
    "nums = 100*100\n",
    "image_paths = [\"DNS-LES_3C/les_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(nums)]\n",
    "target_paths = [\"DNS-LES_3C/dns_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(nums)] \n",
    "print(len(image_paths))\n",
    "#print(image_paths[170099])\n",
    "nums = 300*100\n",
    "image_paths_1 = [\"DNS-LES_3C/les_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(30000, 30000+nums)]\n",
    "target_paths_1 = [\"DNS-LES_3C/dns_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(30000, 30000+nums)] \n",
    "print(len(image_paths))\n",
    "image_paths.extend(image_paths_1)\n",
    "target_paths.extend(target_paths_1)\n",
    "print(len(image_paths))\n",
    "print(image_paths[39999])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dev Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1000\nDNS-LES_3C/les_3c/1.70/361.png\n"
     ]
    }
   ],
   "source": [
    "nums = 1000\n",
    "image_paths_dev = [\"DNS-LES_3C/les_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(10000, 10000+nums)]\n",
    "target_paths_dev = [\"DNS-LES_3C/dns_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(10000, 10000+nums)] \n",
    "print(len(image_paths_dev))\n",
    "#print(image_paths[170099])\n",
    "print(image_paths_dev[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1000\nDNS-LES_3C/les_3c/3.50/155.png\n"
     ]
    }
   ],
   "source": [
    "nums = 1000\n",
    "image_paths_test = [\"DNS-LES_3C/les_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(20000, 20000+nums)]\n",
    "target_paths_test = [\"DNS-LES_3C/dns_3c/\"+str(i/5670)[:3]+\"0/\"+str(i%567)+\".png\" for i in range(20000, 20000+nums)] \n",
    "print(len(image_paths_test))\n",
    "#print(image_paths[170099])\n",
    "print(image_paths_test[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image Transforms (Trial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# img_transform = transforms.Compose([\n",
    "#     transforms.RandomHorizontalFlip(p=0.5),\n",
    "#     transforms.RandomResizedCrop(size = 32, scale=(0.75, 1.0), ratio=(0.8, 1.25)),\n",
    "#     transforms.RandomPerspective(distortion_scale=0.25, p=0.3, interpolation=3),\n",
    "#     transforms.ColorJitter(brightness=0.35, contrast=0.25, saturation=0, hue=0),\n",
    "#     torchvision.transforms.RandomRotation(7),\n",
    "#     transforms.ToTensor(),\n",
    "#     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
    "# ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(data.Dataset):\n",
    "    def __init__(self, image_paths, target_paths, transform=None):\n",
    "\n",
    "        self.image_paths = image_paths\n",
    "        self.target_paths = target_paths\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        ip = Image.open(self.image_paths[index])\n",
    "        op = Image.open(self.target_paths[index])\n",
    "        x = self.transform(ip)\n",
    "        y = self.transform(op)\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model for SR-DNS_Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvBNReLU(nn.Sequential):\n",
    "    def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1):\n",
    "        padding = (kernel_size - 1) // 2\n",
    "        super(ConvBNReLU, self).__init__(\n",
    "            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),\n",
    "            nn.BatchNorm2d(out_planes),\n",
    "            nn.LeakyReLU(inplace=True)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InvertedResidual(nn.Module):\n",
    "    def __init__(self, inp, oup, stride, expand_ratio):\n",
    "        super(InvertedResidual, self).__init__()\n",
    "        self.stride = stride\n",
    "        assert stride in [1, 2]\n",
    "\n",
    "        hidden_dim = int(round(inp * expand_ratio))\n",
    "        self.use_res_connect = self.stride == 1 and inp == oup\n",
    "\n",
    "        layers = []\n",
    "        if expand_ratio != 1:\n",
    "            # pw\n",
    "            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))\n",
    "        layers.extend([\n",
    "            # dw\n",
    "            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),\n",
    "            # pw-linear\n",
    "            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),\n",
    "            nn.BatchNorm2d(oup),\n",
    "        ])\n",
    "        self.conv = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.use_res_connect:\n",
    "            return x + self.conv(x)\n",
    "        else:\n",
    "            return self.conv(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MobileNetv2_SISR(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MobileNetv2_SISR, self).__init__()\n",
    "\n",
    "        self.conv1 = nn.Sequential(\n",
    "        nn.Conv2d(3, 32, 3, stride=1, padding=1),  # b, 32, 64, 64\n",
    "        nn.BatchNorm2d(32),\n",
    "        nn.LeakyReLU()\n",
    "        )\n",
    "        self.bottlenecks1 = nn.Sequential(\n",
    "            InvertedResidual(32, 16, 1, 1),\n",
    "            InvertedResidual(16, 16, 1, 1),\n",
    "            InvertedResidual(16, 24, 1, 6),\n",
    "            InvertedResidual(24, 24, 1, 6),\n",
    "            InvertedResidual(24, 32, 1, 6),\n",
    "            InvertedResidual(32, 32, 1, 6),\n",
    "            InvertedResidual(32, 32, 1, 6),\n",
    "            InvertedResidual(32, 64, 1, 6),\n",
    "            InvertedResidual(64, 64, 1, 6),\n",
    "            InvertedResidual(64, 64, 1, 6),\n",
    "        )\n",
    "\n",
    "        self.bottlenecks2 = nn.Sequential(\n",
    "            InvertedResidual(64, 64, 1, 6),\n",
    "            InvertedResidual(64, 96, 2, 6),\n",
    "            InvertedResidual(96, 96, 1, 6),\n",
    "            InvertedResidual(96, 96, 1, 6),\n",
    "            InvertedResidual(96, 128, 1, 6),\n",
    "            InvertedResidual(128, 128, 1, 6),\n",
    "            InvertedResidual(128, 128, 1, 6),\n",
    "            InvertedResidual(128, 256, 1, 6),\n",
    "            #InvertedResidual(256, 256, 1, 6),\n",
    "        )\n",
    "        self.deconv1 = nn.Sequential(\n",
    "        nn.ConvTranspose2d(64, 3, 3, stride=1, padding=1),  # b, 1, 64, 64\n",
    "        #nn.Tanh()\n",
    "        )\n",
    "        self.pix_shuffle = nn.PixelShuffle(2)\n",
    "        self.tanh = nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        \n",
    "        out = self.bottlenecks1(out)\n",
    "        int1 = out\n",
    "        out = self.bottlenecks2(out)\n",
    "        out = self.pix_shuffle(out)\n",
    "        out = torch.add(out, int1)\n",
    "\n",
    "        out = self.deconv1(out)\n",
    "        out1 = self.tanh(out)+0.0001*out\n",
    "        return out1\n",
    "\n",
    "        #return out\n",
    "        \n",
    "def init_weights(m):\n",
    "    if isinstance(m, nn.Conv2d):\n",
    "        nn.init.kaiming_normal_(m.weight, mode='fan_out')\n",
    "        if m.bias is not None:\n",
    "            nn.init.zeros_(m.bias)\n",
    "    elif isinstance(m, nn.BatchNorm2d):\n",
    "        nn.init.ones_(m.weight)\n",
    "        nn.init.zeros_(m.bias)\n",
    "    elif isinstance(m, nn.Linear):\n",
    "        nn.init.normal_(m.weight, 0, 0.01)\n",
    "        nn.init.zeros_(m.bias)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Metrics\n",
    "* PSNR\n",
    "* SSIM\n",
    "* KE\n",
    "* Vorticity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def PSNR(op, t, batch_size): \n",
    "    mse = torch.sum((t - op) ** 2) \n",
    "    #print(mse.size())\n",
    "    mse /= (batch_size*64*64)\n",
    "\n",
    "    max_pixel = 1.\n",
    "    psnr = 20 * torch.log10(max_pixel / torch.sqrt(mse))\n",
    "    #print(psnr.size())\n",
    "    return psnr "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def KE(img, op, t): \n",
    "    ke_recon = torch.sum(op ** 2)/op.size()[0] \n",
    "    ke_dns = torch.sum(t ** 2)/t.size()[0] \n",
    "    ke_les = torch.sum(img ** 2)/img.size()[0] \n",
    " \n",
    "    return ke_les, ke_recon, ke_dns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Avg_KE(img, op, t): \n",
    "    # pdb.set_trace()\n",
    "\n",
    "    op = np.squeeze(op)\n",
    "    img = np.squeeze(img)\n",
    "    t = np.squeeze(t)\n",
    "    # pdb.set_trace()\n",
    "\n",
    "    ke_recon = torch.mean(torch.abs(op - torch.mean(op)))\n",
    "    ke_dns = torch.mean(torch.abs(t - torch.mean(t)))\n",
    "    ke_les = torch.mean(torch.abs(img - torch.mean(img)))\n",
    " \n",
    "    return ke_les, ke_recon, ke_dns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.measure import compare_ssim\n",
    "def SSIM(op, t, batch_size):\n",
    "    ssim = 0 \n",
    "    #print(op.size(), t.size())\n",
    "    for i in range(op.size()[0]):\n",
    "        tar = to_img(t[i])\n",
    "        out = to_img(op[i])\n",
    "        # print(out[0,0].size())\n",
    "        (score, diff) = compare_ssim(out[0,0].cpu().numpy(), tar[0,0].cpu().numpy(), full=True)\n",
    "        ssim+=score/batch_size\n",
    "    \n",
    "        #print(\"SSIM: {}\".format(score))\n",
    "    return ssim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(model, data_loader, criterion, optimizer):\n",
    "    model.train()\n",
    "    \n",
    "    running_loss = 0\n",
    "    avg_psnr = 0\n",
    "    avg_psnr_les = 0\n",
    "\n",
    "    start_time = time.time()\n",
    "    print('Train Loop')\n",
    "    for batch_num, (img, target) in tqdm(enumerate(data_loader), total=len(data_loader), ascii=True):\n",
    "\n",
    "        img = img.to(device)\n",
    "        target = target.to(device)\n",
    "\n",
    "        output = model(img)\n",
    "        loss = criterion(output, target)\n",
    "        psnr = PSNR(output, target, batch_size)\n",
    "        psnr_les = PSNR(img, target, batch_size)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()/len(data_loader)\n",
    "        avg_psnr += psnr.item()/len(data_loader)\n",
    "        avg_psnr_les += psnr_les.item()/len(data_loader)\n",
    "\n",
    "    #print(' ')\n",
    "    print('Train_Loss:{:.6f}, PSNR_DNS:{:.4f}, PSNR_LES:{:.4f}'.format(running_loss, psnr, psnr_les))\n",
    "    torch.cuda.empty_cache()\n",
    "    end_time = time.time()\n",
    "    del img\n",
    "    del target\n",
    "    del loss\n",
    "    print(\"Train Time: {:.2f} s\".format(end_time-start_time))\n",
    "\n",
    "    return running_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dev_epoch(model, data_loader, criterion):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "\n",
    "        running_loss = 0\n",
    "        avg_psnr = 0\n",
    "        avg_psnr_les = 0\n",
    "        avg_ssim_dns = 0\n",
    "        avg_ssim_les = 0\n",
    "\n",
    "        start_time = time.time()\n",
    "        print('Dev Loop')\n",
    "        print(' ')\n",
    "\n",
    "        for batch_num, (img, target) in tqdm(enumerate(data_loader), total=len(data_loader), ascii=True):\n",
    "\n",
    "            img = img.to(device)\n",
    "            target = target.to(device)\n",
    "\n",
    "            output = model(img)\n",
    "            loss = criterion(output, target)\n",
    "            psnr = PSNR(output, target,batch_size)\n",
    "            psnr_les = PSNR(img, target,batch_size)\n",
    "            ssim_dns = SSIM(output, target, batch_size)\n",
    "            ssim_les = SSIM(img, target, batch_size)\n",
    "\n",
    "            running_loss += loss.item()/len(data_loader)\n",
    "            avg_psnr += psnr.item()/len(data_loader)\n",
    "            avg_psnr_les += psnr_les.item()/len(data_loader)\n",
    "            avg_ssim_dns += ssim_dns.item()/len(test_loader)\n",
    "            avg_ssim_les += ssim_les.item()/len(test_loader)\n",
    "\n",
    "        torch.cuda.empty_cache()\n",
    "        end_time = time.time()\n",
    "\n",
    "        del img\n",
    "        del target\n",
    "        del loss\n",
    "        del psnr\n",
    "        del psnr_les\n",
    "\n",
    "        print(\"Dev Time: {:.2f} s\".format(end_time-start_time))\n",
    "\n",
    "        return running_loss, avg_psnr_les, avg_psnr, avg_ssim_les, avg_ssim_dns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "40000\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 200\n",
    "batch_size = 32\n",
    "learning_rate = 1e-3\n",
    "\n",
    "img_transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(p=0.05),\n",
    "    transforms.RandomVerticalFlip(p=0.05),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "])\n",
    "\n",
    "test_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.5], [0.5])\n",
    "])\n",
    "\n",
    "train_dataset = MyDataset(image_paths, target_paths, transform=img_transform)\n",
    "train_loader_args = dict(batch_size=batch_size, shuffle=True, num_workers=8)\n",
    "train_loader = data.DataLoader(train_dataset, **train_loader_args)\n",
    "print(train_dataset.__len__())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "dev_dataset = MyDataset(image_paths_dev, target_paths_dev, transform=test_transform)\n",
    "dev_loader_args = dict(batch_size=batch_size, shuffle=False, num_workers=8)\n",
    "dev_loader = data.DataLoader(dev_dataset, **dev_loader_args)\n",
    "print(dev_dataset.__len__())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "test_dataset = MyDataset(image_paths_test, target_paths_test, transform=test_transform)\n",
    "test_loader_args = dict(batch_size=1, shuffle=False, num_workers=8)\n",
    "test_loader = data.DataLoader(test_dataset, **test_loader_args)\n",
    "print(test_dataset.__len__())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "MobileNetv2_SISR(\n  (conv1): Sequential(\n    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n    (2): LeakyReLU(negative_slope=0.01)\n  )\n  (bottlenecks1): Sequential(\n    (0): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (1): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=16, bias=False)\n          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (2): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=96, bias=False)\n          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (3): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (4): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(144, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (5): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (6): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (7): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (8): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (9): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n  )\n  (bottlenecks2): Sequential(\n    (0): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (1): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=384, bias=False)\n          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(384, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (2): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (3): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (4): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(576, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (5): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(768, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (6): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(768, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n    (7): InvertedResidual(\n      (conv): Sequential(\n        (0): ConvBNReLU(\n          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (1): ConvBNReLU(\n          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n        )\n        (2): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n        (3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n      )\n    )\n  )\n  (deconv1): Sequential(\n    (0): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n  )\n  (pix_shuffle): PixelShuffle(upscale_factor=2)\n  (tanh): Tanh()\n)\n"
     ]
    }
   ],
   "source": [
    "model = MobileNetv2_SISR()\n",
    "model.apply(init_weights)\n",
    "device = torch.device(\"cuda\")\n",
    "model.eval()\n",
    "model.to(device)        \n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_epochs = 200\n",
    "Train_Loss = []\n",
    "Dev_Loss = []\n",
    "Dev_Acc = []\n",
    "criterion = nn.L1Loss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode = 'min', factor = 0.1, patience=3, threshold=5e-4, eps=1e-6)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "\n",
    "    train_loss = train_epoch(model, train_loader, criterion, optimizer)\n",
    "    dev_loss, psnr_les, psnr, ssim_les, ssim_dns = dev_epoch(model, dev_loader, criterion)\n",
    "    Train_Loss.append(train_loss)\n",
    "    Dev_Loss.append(dev_loss)\n",
    "\n",
    "    scheduler.step(dev_loss)\n",
    "\n",
    "    print(' ')\n",
    "    print('epoch [{}/{}], Train_Loss:{:.6f}, Dev_Loss:{:.6f}'.format(epoch+1, num_epochs, train_loss, dev_loss))\n",
    "    print('PSNR_DNS:{:.4f}, PSNR_LES:{:.4f}, SSIM_LES:{:.4f}, SSIM_DNS:{:.4f}'.format(psnr, psnr_les, ssim_les, ssim_dns))\n",
    "\n",
    "    torch.save(model.state_dict(), 'model_3c/SISR_mv2f_{}.pth'.format(epoch))\n",
    "    scheduler.step(train_loss)\n",
    "    print(optimizer)\n",
    "    print('='*100)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_MAE(L, R, D):\n",
    "\n",
    "    MAE = np.abs((np.array(R) - np.array(D))/np.array(D))\n",
    "    print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg = np.mean(MAE)\n",
    "    plt.xticks([])\n",
    "    plt.title('Reconstruction')\n",
    "    plt.ylabel('MAE')\n",
    "    plt.ylim(0,1.2)\n",
    "    plt.plot(MAE, 'ko', fillstyle = 'none')\n",
    "    plt.axhline(avg, color = 'r')\n",
    "    plt.savefig('MAE.eps')\n",
    "    plt.close()\n",
    "\n",
    "    MAE_L = np.abs((np.array(L) - np.array(D))/np.array(D))\n",
    "    print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg_l = np.mean(MAE_L)\n",
    "    plt.xticks([])\n",
    "    plt.title('LES')\n",
    "    plt.ylabel('MAE')\n",
    "    # plt.plot(MAE, 'yo', fillstyle = 'none')\n",
    "    plt.plot(MAE_L, 'ko', fillstyle = 'none')\n",
    "    # plt.axhline(avg, color = 'r')\n",
    "    plt.axhline(avg_l, color = 'r')\n",
    "    plt.savefig('MAE_L.eps')  # plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_Avg_MAE(L, R, D):\n",
    "\n",
    "    # MAE = np.abs((np.array(R) - np.array(D))/np.array(D))\n",
    "    # print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg = np.mean(np.abs(R))\n",
    "    # plt.xticks([])\n",
    "    plt.title('Reconstruction')\n",
    "    plt.ylabel('Average Turbulent Velocity')\n",
    "    # plt.ylim(0,1.2)\n",
    "    plt.plot(np.abs(R), 'ko', fillstyle = 'none')\n",
    "    plt.axhline(avg, color = 'r')\n",
    "    plt.savefig('AM.eps')\n",
    "    plt.close()\n",
    "\n",
    "    # MAE_L = np.abs((np.array(L) - np.array(D))/np.array(D))\n",
    "    # print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg_l = np.mean(np.abs(D))\n",
    "    # plt.xticks([])\n",
    "    plt.title('DNS')\n",
    "    plt.ylabel('Average Turbulent Velocity')\n",
    "    # plt.plot(MAE, 'yo', fillstyle = 'none')\n",
    "    plt.plot(np.abs(D), 'ko', fillstyle = 'none')\n",
    "    # plt.axhline(avg, color = 'r')\n",
    "    plt.axhline(avg_l, color = 'r')\n",
    "    plt.savefig('AM_L.eps')  # plt.show()\n",
    "    plt.close()\n",
    "\n",
    "    error = np.abs(np.array(D) - np.array(R))\n",
    "    avg_e = np.mean(np.abs(error))\n",
    "    # plt.xticks([])\n",
    "    plt.title('Error')\n",
    "    plt.ylim(0, 0.35)\n",
    "    plt.ylabel('Turbulent Velocity Error')\n",
    "    # plt.plot(MAE, 'yo', fillstyle = 'none')\n",
    "    plt.plot(error, 'ko', fillstyle = 'none')\n",
    "    # plt.axhline(avg, color = 'r')\n",
    "    plt.axhline(avg_e, color = 'r')\n",
    "    plt.savefig('error.eps')  # plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "#combined plot for turbulent velocity\n",
    "def plot_Avg_MAE(L, R, D):\n",
    "\n",
    "    # MAE = np.abs((np.array(R) - np.array(D))/np.array(D))\n",
    "    # print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg = np.mean(np.abs(R))\n",
    "    # plt.xticks([])\n",
    "    # plt.title('Reconstruction')\n",
    "    plt.ylabel('Average Turbulent Velocity')\n",
    "    plt.xlabel('Samples')\n",
    "    # plt.ylim(0,1.2)\n",
    "    plt.plot(np.abs(R),  marker = '^', c = 'dodgerblue', label='Recon',ls=' ', ms='3.5')\n",
    "\n",
    "\n",
    "    # MAE_L = np.abs((np.array(L) - np.array(D))/np.array(D))\n",
    "    # print(MAE.shape)\n",
    "    # print(MAE)\n",
    "    avg_l = np.mean(np.abs(D))\n",
    "    # plt.xticks([])\n",
    "    # plt.title('DNS')\n",
    "    # plt.ylabel('Average Turbulent Velocity')\n",
    "    # plt.plot(MAE, 'yo', fillstyle = 'none')\n",
    "    plt.plot(np.abs(D), marker = '+', c = 'darkorange', label='DNS',ls=' ', ms='4')\n",
    "    # plt.axhline(avg, color = 'r')\n",
    "    plt.axhline(avg_l, color = 'red', ls='-.',label='Avg. DNS', lw='1.5')\n",
    "    plt.axhline(avg, color = 'k', ls='-.', label='Avg. Recon', lw='1.5')\n",
    "    plt.legend(loc='best')\n",
    "    plt.savefig('combined.eps')  # plt.show()\n",
    "    plt.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metric Predictions (Test and Train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_predictions(model, test_loader):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        avg_psnr = 0\n",
    "        avg_psnr_les = 0\n",
    "        avg_dns_ke = 0\n",
    "        avg_les_ke = 0\n",
    "        avg_recon_ke = 0\n",
    "        avg_ssim_les = 0\n",
    "        avg_ssim_dns = 0\n",
    "\n",
    "        P = []\n",
    "        L = []\n",
    "        R = []\n",
    "        D = []\n",
    "\n",
    "        L_ = []\n",
    "        R_ = []\n",
    "        D_ = []\n",
    "\n",
    "        for batch_idx, (img, target) in enumerate(test_loader):   \n",
    "            \n",
    "            img = img.to(device)\n",
    "            target = target.to(device)\n",
    "            \n",
    "            out = model(img)\n",
    "            psnr = PSNR(out, target, batch_size)\n",
    "            psnr_les = PSNR(img, target, batch_size)\n",
    "\n",
    "            les_ke, recon_ke, dns_ke = KE(img, out, target)\n",
    "            les_ake, recon_ake, dns_ake = Avg_KE(img, out, target)\n",
    "\n",
    "            ssim_dns = SSIM(out, target, batch_size)\n",
    "            ssim_les = SSIM(img, target, batch_size)\n",
    "            L.append(les_ke.cpu().numpy())\n",
    "            R.append(recon_ke.cpu().numpy())\n",
    "            D.append(dns_ke.cpu().numpy())\n",
    "\n",
    "            L_.append(les_ake.cpu().numpy())\n",
    "            R_.append(recon_ake.cpu().numpy())\n",
    "            D_.append(dns_ake.cpu().numpy())\n",
    "\n",
    "            pic = to_img(out.cpu()) #Only the first image of the batch.\n",
    "            save_image(pic, 'test_3c/image_{}.png'.format(batch_idx))\n",
    "            t = to_img(target) #Only the first image of the batch.\n",
    "            save_image(t, 'test_3c/image_{}_t.png'.format(batch_idx))\n",
    "            i = to_img(img) #Only the first image of the batch.\n",
    "            save_image(i, 'test_3c/image_{}_og.png'.format(batch_idx))\n",
    "\n",
    "            del img\n",
    "            del target\n",
    "\n",
    "            avg_psnr += psnr.item()/len(test_loader)\n",
    "            avg_psnr_les += psnr_les.item()/len(test_loader)\n",
    "\n",
    "            avg_dns_ke += dns_ke.item()/len(test_loader)\n",
    "            avg_recon_ke += recon_ke.item()/len(test_loader)\n",
    "            avg_les_ke += les_ke.item()/len(test_loader)\n",
    "\n",
    "            avg_ssim_les += ssim_les.item()/len(test_loader)\n",
    "            avg_ssim_dns += ssim_dns.item()/len(test_loader)\n",
    "\n",
    "    plot_MAE(L, R, D)\n",
    "    plot_Avg_MAE(L_, R_, D_)\n",
    "    return avg_psnr, avg_psnr_les, avg_les_ke, avg_recon_ke, avg_dns_ke, avg_ssim_les, avg_ssim_dns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n",
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n",
      "(1000,)\n",
      "(1000,)\n",
      "14.990708944320705\n",
      "11.733874299049377\n",
      "1341.4863097534167\n",
      "1519.8945644531243\n",
      "2152.8065732421906\n",
      "0.8105047876127578\n",
      "0.8720655256391564\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "avg_psnr, avg_psnr_les, avg_les_ke, avg_recon_ke, avg_dns_ke, avg_ssim_les, avg_ssim_dns = test_predictions(model, test_loader)\n",
    "print(avg_psnr)\n",
    "print(avg_psnr_les)\n",
    "print(avg_recon_ke)\n",
    "print(avg_dns_ke)\n",
    "print(avg_les_ke)\n",
    "print(avg_ssim_les)\n",
    "print(avg_ssim_dns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_predictions(model, train_loader):\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        avg_psnr = 0\n",
    "        avg_psnr_les = 0\n",
    "        avg_dns_ke = 0\n",
    "        avg_les_ke = 0\n",
    "        avg_recon_ke = 0\n",
    "        avg_ssim_les = 0\n",
    "        avg_ssim_dns = 0\n",
    "\n",
    "        P = []\n",
    "\n",
    "        for batch_idx, (img, target) in enumerate(train_loader):  \n",
    "\n",
    "            # if batch_idx > 3:\n",
    "            #     break \n",
    "            \n",
    "            img = img.to(device)\n",
    "            target = target.to(device)\n",
    "            \n",
    "            out = model(img)\n",
    "            psnr = PSNR(out, target, batch_size)\n",
    "            psnr_les = PSNR(img, target, batch_size)\n",
    "            les_ke, recon_ke, dns_ke = KE(img, out, target)\n",
    "            ssim_dns = SSIM(out, target, batch_size)\n",
    "            ssim_les = SSIM(img, target, batch_size)\n",
    "\n",
    "            pic = to_img(out.cpu()) #Only the first image of the batch.\n",
    "            save_image(pic, 'test_3c/image_tr_{}.png'.format(batch_idx))\n",
    "            t = to_img(target) #Only the first image of the batch.\n",
    "            save_image(t, 'test_3c/image_tr_{}_t.png'.format(batch_idx))\n",
    "            i = to_img(img) #Only the first image of the batch.\n",
    "            save_image(i, 'test_3c/image_tr_{}_og.png'.format(batch_idx))\n",
    "\n",
    "            del img\n",
    "            del target\n",
    "\n",
    "            avg_psnr += psnr.item()/len(train_loader)\n",
    "            avg_psnr_les += psnr_les.item()/len(train_loader)\n",
    "\n",
    "            avg_dns_ke += dns_ke.item()/len(train_loader)\n",
    "            avg_recon_ke += recon_ke.item()/len(train_loader)\n",
    "            avg_les_ke += les_ke.item()/len(train_loader)\n",
    "\n",
    "            avg_ssim_les += ssim_les.item()/len(train_loader)\n",
    "            avg_ssim_dns += ssim_dns.item()/len(train_loader)\n",
    "\n",
    "    return avg_psnr, avg_psnr_les, avg_les_ke, avg_recon_ke, avg_dns_ke, avg_ssim_les, avg_ssim_dns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "avg_psnr, avg_psnr_les, avg_les_ke, avg_recon_ke, avg_dns_ke, avg_ssim_les, avg_ssim_dns = train_predictions(model, train_loader)\n",
    "print(avg_psnr)\n",
    "print(avg_psnr_les)\n",
    "print(avg_recon_ke)\n",
    "print(avg_dns_ke)\n",
    "print(avg_les_ke)\n",
    "print(avg_ssim_les)\n",
    "print(avg_ssim_dns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use this model for loading..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "MobileNetv2_SISR(\n",
       "  (conv1): Sequential(\n",
       "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): LeakyReLU(negative_slope=0.01)\n",
       "  )\n",
       "  (bottlenecks1): Sequential(\n",
       "    (0): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)\n",
       "          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=16, bias=False)\n",
       "          (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (2): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(16, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=96, bias=False)\n",
       "          (1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(96, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (3): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n",
       "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(144, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (4): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(24, 144, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=144, bias=False)\n",
       "          (1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(144, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (5): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (6): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(192, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (7): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(32, 192, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=192, bias=False)\n",
       "          (1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(192, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (8): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (9): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (bottlenecks2): Sequential(\n",
       "    (0): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=384, bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(384, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(64, 384, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(384, 384, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=384, bias=False)\n",
       "          (1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(384, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (2): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (3): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(576, 96, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (4): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(96, 576, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=576, bias=False)\n",
       "          (1): BatchNorm2d(576, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(576, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (5): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(768, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (6): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(768, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (7): InvertedResidual(\n",
       "      (conv): Sequential(\n",
       "        (0): ConvBNReLU(\n",
       "          (0): Conv2d(128, 768, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (1): ConvBNReLU(\n",
       "          (0): Conv2d(768, 768, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=768, bias=False)\n",
       "          (1): BatchNorm2d(768, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (2): LeakyReLU(negative_slope=0.01, inplace=True)\n",
       "        )\n",
       "        (2): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (deconv1): Sequential(\n",
       "    (0): ConvTranspose2d(64, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  )\n",
       "  (pix_shuffle): PixelShuffle(upscale_factor=2)\n",
       "  (tanh): Tanh()\n",
       ")"
      ]
     },
     "metadata": {},
     "execution_count": 20
    }
   ],
   "source": [
    "model = MobileNetv2_SISR()\n",
    "PATH = 'model_3c/SISR_mv2f_48.pth'\n",
    "model.load_state_dict(torch.load(PATH))\n",
    "device = torch.device(\"cuda\")\n",
    "model.eval()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "version": "3.7.6-final"
  },
  "orig_nbformat": 2,
  "file_extension": ".py",
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3,
  "kernelspec": {
   "name": "python37664bitpytorchenvcondac65bc6c7ab474553ae2b82822152b36f",
   "display_name": "Python 3.7.6 64-bit ('pytorch_env': conda)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}