aimet_zoo_torch/superres/evaluators/notebooks/superres_quanteval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b6189f92",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8181df87",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os, sys\n",
    "sys.path.append(os.path.dirname(os.getcwd()))\n",
    "import glob\n",
    "import urllib.request\n",
    "import tarfile\n",
    "from matplotlib import pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from aimet_torch.quantsim import QuantizationSimModel\n",
    "from aimet_torch.qc_quantize_op import QuantScheme\n",
    "from utils.imresize import imresize\n",
    "from utils.models import *\n",
    "from utils.helpers import *\n",
    "from utils.downloader import get_tar_path\n",
    "from utils.inference import load_model, run_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87ad0177",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Global Constants"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5acd1f05",
   "metadata": {},
   "source": [
    "These are all the constant variables that we use throughout this notebook. It also specifies the different model variations that were trained. You can select the necessary model at the end of this notebook while running inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8a89a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "''' Set the following variable to the path of your dataset (parent directory of actual images) '''\n",
    "DATA_DIR = '/<path to parent>/'\n",
    "DATASET_NAME = 'Set14' # Tested on Set5, Set14 and BSDS100\n",
    "\n",
    "# Directory to store downloaded checkpoints\n",
    "CHECKPOINT_DIR = './checkpoints/'\n",
    "if not os.path.exists(CHECKPOINT_DIR):\n",
    "    os.makedirs(CHECKPOINT_DIR)\n",
    "    \n",
    "use_cuda = True  # Whether to use CUDA or CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86b3209",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "MODEL_ARGS = {\n",
    "    'ABPNRelease': {\n",
    "        'abpn_28_2x': {\n",
    "            'num_channels': 28,\n",
    "            'scaling_factor': 2,\n",
    "        },\n",
    "        'abpn_28_3x': {\n",
    "            'num_channels': 28,\n",
    "            'scaling_factor': 3,\n",
    "        },\n",
    "        'abpn_28_4x': {\n",
    "            'num_channels': 28,\n",
    "            'scaling_factor': 4,\n",
    "        },\n",
    "        'abpn_32_2x': {\n",
    "            'num_channels': 32,\n",
    "            'scaling_factor': 2,\n",
    "        },\n",
    "        'abpn_32_3x': {\n",
    "            'num_channels': 32,\n",
    "            'scaling_factor': 3,\n",
    "        },\n",
    "        'abpn_32_4x': {\n",
    "            'num_channels': 32,\n",
    "            'scaling_factor': 4,\n",
    "        }\n",
    "    },\n",
    "    'XLSRRelease': {\n",
    "        'xlsr_2x': {\n",
    "            'scaling_factor': 2,\n",
    "        },\n",
    "        'xlsr_3x': {\n",
    "            'scaling_factor': 3,\n",
    "        },\n",
    "        'xlsr_4x': {\n",
    "            'scaling_factor': 4,\n",
    "        }\n",
    "    },\n",
    "    'SESRRelease_M3': {\n",
    "        'sesr_m3_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'sesr_m3_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'sesr_m3_4x': {\n",
    "            'scaling_factor': 4\n",
    "        },\n",
    "    },\n",
    "    'SESRRelease_M5': {\n",
    "        'sesr_m5_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'sesr_m5_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'sesr_m5_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'SESRRelease_M7': {\n",
    "        'sesr_m7_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'sesr_m7_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'sesr_m7_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'SESRRelease_M11': {\n",
    "        'sesr_m11_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'sesr_m11_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'sesr_m11_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'SESRRelease_XL': {\n",
    "         'sesr_xl_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'sesr_xl_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'sesr_xl_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'QuickSRNetSmall': {\n",
    "        'quicksrnet_small_1.5x': {\n",
    "            'scaling_factor': 1.5\n",
    "        },\n",
    "        'quicksrnet_small_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'quicksrnet_small_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'quicksrnet_small_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'QuickSRNetMedium': {\n",
    "        'quicksrnet_medium_1.5x': {\n",
    "            'scaling_factor': 1.5\n",
    "        },\n",
    "        'quicksrnet_medium_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'quicksrnet_medium_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'quicksrnet_medium_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    },\n",
    "    'QuickSRNetLarge': {\n",
    "        'quicksrnet_large_1.5x': {\n",
    "            'scaling_factor': 1.5\n",
    "        },\n",
    "        'quicksrnet_large_2x': {\n",
    "            'scaling_factor': 2\n",
    "        },\n",
    "        'quicksrnet_large_3x': {\n",
    "            'scaling_factor': 3\n",
    "        },\n",
    "        'quicksrnet_large_4x': {\n",
    "            'scaling_factor': 4\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "MODELS = list(MODEL_ARGS.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30bcbe92",
   "metadata": {},
   "source": [
    "# Select a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33868e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_DICT = {}\n",
    "for idx in range(len(MODELS)):\n",
    "    MODEL_DICT[idx] = MODELS[idx]\n",
    "\n",
    "MODEL_DICT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05484950",
   "metadata": {},
   "source": [
    "Select one of the models printed above by selecting the corresponding index in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d479cd0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "''' Set this variable'''\n",
    "model_index = 0  # Model index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c7fd63",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = MODELS[model_index]  # Selected model type\n",
    "\n",
    "MODEL_SPECS = list(MODEL_ARGS.get(MODEL_NAME).keys())\n",
    "\n",
    "MODEL_SPECS_DICT = {}\n",
    "for idx in range(len(MODEL_SPECS)):\n",
    "    MODEL_SPECS_DICT[idx] = MODEL_SPECS[idx]\n",
    "\n",
    "MODEL_SPECS_DICT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "913b88e8",
   "metadata": {},
   "source": [
    "Select one of the models printed above by selecting the corresponding index in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5388b4b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "''' Set this variable'''\n",
    "model_spec_index = 0 # Model specification index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e95616f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose model\n",
    "MODEL_CONFIG = MODEL_SPECS[model_spec_index]\n",
    "print(f'{MODEL_CONFIG} will be used')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ff2460",
   "metadata": {},
   "source": [
    "Automatically download model weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a5137f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define paths to download filenames\n",
    "# Do not change the file names\n",
    "FILENAME_FP32 = 'checkpoint_float32.pth.tar'  # full precision model\n",
    "FILENAME_INT8 = 'checkpoint_int8.pth' # quantized model\n",
    "ENCODINGS = 'checkpoint_int8.encodings'\n",
    "if MODEL_CONFIG.startswith('quicksrnet'):\n",
    "    ENCODINGS = 'checkpoint_int8_params_only.encodings'  # encodings of the quantized models\n",
    "\n",
    "# Path to desired model weights and encodings (if necessary)\n",
    "ENCODING_PATH = os.path.join(CHECKPOINT_DIR, f'release_{MODEL_CONFIG}', ENCODINGS)\n",
    "\n",
    "# Path to model checkpoint and encodings (if necessary)\n",
    "MODEL_PATH_INT8 = os.path.join(CHECKPOINT_DIR, f'release_{MODEL_CONFIG}', FILENAME_INT8)\n",
    "MODEL_PATH_FP32 = os.path.join(CHECKPOINT_DIR, f'release_{MODEL_CONFIG}', FILENAME_FP32)\n",
    "\n",
    "# AIMET Config\n",
    "QUANTSIM_CONFIG_FILENAME = 'default_config_per_channel.json'\n",
    "CONFIG_PATH = os.path.join(CHECKPOINT_DIR, 'default_config_per_channel.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fae0a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(MODEL_PATH_INT8) or not os.path.exists(MODEL_PATH_FP32) or not os.path.exists(ENCODING_PATH):\n",
    "    print('Downloading model weights')\n",
    "    tar_path = get_tar_path(model_index, model_spec_index) # path to model weights .tar\n",
    "    urllib.request.urlretrieve(tar_path, \n",
    "                               CHECKPOINT_DIR+tar_path.split('/')[-1])\n",
    "    with tarfile.open(CHECKPOINT_DIR+tar_path.split('/')[-1]) as pth_weights:\n",
    "          pth_weights.extractall(CHECKPOINT_DIR)\n",
    "            \n",
    "if not os.path.exists(CONFIG_PATH):\n",
    "    QUANTSIM_CONFIG_URL = 'https://raw.githubusercontent.com/quic/aimet/1.23.0/TrainingExtensions/common/src/python/aimet_common/quantsim_config/default_config_per_channel.json'\n",
    "    urllib.request.urlretrieve(QUANTSIM_CONFIG_URL, \"default_config_per_channel.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acac4223",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d94bb80b",
   "metadata": {},
   "source": [
    "Load test-set images (low-res and high-res pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c392f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to test images\n",
    "TEST_IMAGES_DIR = os.path.join(DATA_DIR, DATASET_NAME)\n",
    "\n",
    "# Get test images\n",
    "IMAGES_LR, IMAGES_HR = load_dataset(TEST_IMAGES_DIR, MODEL_ARGS[MODEL_NAME].get(MODEL_CONFIG)['scaling_factor'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dd1faa3",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Create model instance and load weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca47bf9",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Load the model\n",
    "model_original_fp32 = load_model(MODEL_PATH_FP32, MODEL_NAME, MODEL_ARGS[MODEL_NAME].get(MODEL_CONFIG), \n",
    "                   use_quant_sim_model=False, encoding_path=ENCODING_PATH, quantsim_config_path=CONFIG_PATH,\n",
    "                   calibration_data=IMAGES_LR, use_cuda=use_cuda)\n",
    "\n",
    "model_original_int8 = load_model(MODEL_PATH_FP32, MODEL_NAME, MODEL_ARGS[MODEL_NAME].get(MODEL_CONFIG), \n",
    "                   use_quant_sim_model=True, encoding_path=ENCODING_PATH, quantsim_config_path=CONFIG_PATH,\n",
    "                   calibration_data=IMAGES_LR, use_cuda=use_cuda)\n",
    "\n",
    "model_optimized_fp32 = load_model(MODEL_PATH_INT8, MODEL_NAME, MODEL_ARGS[MODEL_NAME].get(MODEL_CONFIG), \n",
    "                   use_quant_sim_model=False, encoding_path=ENCODING_PATH, quantsim_config_path=CONFIG_PATH,\n",
    "                   calibration_data=IMAGES_LR, use_cuda=use_cuda, before_quantization=True)\n",
    "\n",
    "model_optimized_int8 = load_model(MODEL_PATH_INT8, MODEL_NAME, MODEL_ARGS[MODEL_NAME].get(MODEL_CONFIG), \n",
    "                   use_quant_sim_model=True, encoding_path=ENCODING_PATH, quantsim_config_path=CONFIG_PATH,\n",
    "                   calibration_data=IMAGES_LR, use_cuda=use_cuda, before_quantization=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3783302e",
   "metadata": {},
   "source": [
    "# Model Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cddda678",
   "metadata": {},
   "source": [
    "Run inference to get the respective super-resolved images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8819afb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run model inference on test images and get super-resolved images\n",
    "IMAGES_SR_original_fp32 = run_model(model_original_fp32, IMAGES_LR, use_cuda)\n",
    "IMAGES_SR_original_int8 = run_model(model_original_int8.model, IMAGES_LR, use_cuda)\n",
    "IMAGES_SR_optimized_fp32 = run_model(model_optimized_fp32, IMAGES_LR, use_cuda)\n",
    "IMAGES_SR_optimized_int8 = run_model(model_optimized_int8.model, IMAGES_LR, use_cuda)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f1b17d3",
   "metadata": {},
   "source": [
    "Calculate average-PSNR between the test-set high-res and super-resolved images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f63d02",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the average PSNR for all test-images\n",
    "avg_psnr = evaluate_average_psnr(IMAGES_SR_original_fp32, IMAGES_HR)\n",
    "print(f'Original Model | FP32 Environment | Avg. PSNR: {avg_psnr:.3f}')\n",
    "avg_psnr = evaluate_average_psnr(IMAGES_SR_original_int8, IMAGES_HR)\n",
    "print(f'Original Model | INT8 Environment | Avg. PSNR: {avg_psnr:.3f}')\n",
    "avg_psnr = evaluate_average_psnr(IMAGES_SR_optimized_fp32, IMAGES_HR)\n",
    "print(f'Optimized Model | FP32 Environment | Avg. PSNR: {avg_psnr:.3f}')\n",
    "avg_psnr = evaluate_average_psnr(IMAGES_SR_optimized_int8, IMAGES_HR)\n",
    "print(f'Optimized Model | INT8 Environment | Avg. PSNR: {avg_psnr:.3f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}