tutorials/03.carnx2_super_resolution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0e043d12-b96b-4007-8af7-3cc0d775b31c",
   "metadata": {},
   "source": [
    "## Tutorial 3. CARN on SuperResolution. \n",
    "\n",
    "In this tutorial, we will show \n",
    "\n",
    "- How to end-to-end train and structurally prune a CARN from scratch on DIV2K to get a compressed CARN.\n",
    "- In this specific run, the compressed (via pruning mode of OTO) could reduce FLOPs and parameters by 81.6% and 82%.\n",
    "- The PSNR on Set14, B100, and Urban100 are 33.05, 31.81, and 30.79, respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd59f7b-1675-49fb-aebd-d69a3cff5d59",
   "metadata": {},
   "source": [
    "### Step 1. Create OTO instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2a35e7ce-64c8-4603-9e06-50599f52f1ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OTO graph constructor\n",
      "graph build\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "from sanity_check.backends import CarnNet\n",
    "from only_train_once import OTO\n",
    "import torch\n",
    "\n",
    "scale = 2\n",
    "model = CarnNet(scale=scale, multi_scale=False, group=1)\n",
    "dummy_input = torch.rand(1, 3, 224, 224)\n",
    "\n",
    "oto = OTO(model.cuda(), (dummy_input.cuda(), scale))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85914b25-9a22-40e4-a4b3-ad6ca81d09d2",
   "metadata": {},
   "source": [
    "#### (Optional) Visualize the pruning dependancy graph of DNN\n",
    "\n",
    "Set `display_params=True` could display parameters and shapes on each node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "27ee22d9-f7ba-40e1-ab65-cca107c2165e",
   "metadata": {},
   "outputs": [],
   "source": [
    "oto.visualize(view=False, out_dir='./cache')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95128b58-b7b7-44fa-b30e-9ae440a49ad2",
   "metadata": {},
   "source": [
    "### Step 2 Set up the second last conv operator as unprunable\n",
    "\n",
    "It was observed having some trouble if that conv included into pruning upon current salience score calculation.\n",
    "\n",
    "It can be done by either mark node group as unprunable via node_ids or param_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f9ce1bac-137a-4384-ba71-37c1894336f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Different torch version may have different node id, use the visualization tool to locate\n",
    "# oto.mark_unprunable_by_node_ids(['node-158']) \n",
    "\n",
    "# Or use param_name to locate the node group to make it as unprunable\n",
    "oto.mark_unprunable_by_param_names(['exit.weight']) \n",
    "\n",
    "# Check the pruning dependency graph after `mark_unprunable`\n",
    "oto.visualize(view=False, out_dir='./cache', display_params=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5780c76-8bfb-4057-9a6f-0f069ff04e04",
   "metadata": {},
   "source": [
    "### Step 3. Dataset Preparation\n",
    "\n",
    "Follow https://github.com/nmhkahn/CARN-pytorch/tree/master?tab=readme-ov-file to prepare train, val datasets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2dc5ed12-edcc-4765-ab89-f061e29c9eed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-02-04 21:02:04--  http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip\n",
      "Resolving data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)... 2001:67c:10ec:36c2::178, 129.132.52.178\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:80... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip [following]\n",
      "--2024-02-04 21:02:05--  https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 3530603713 (3.3G) [application/zip]\n",
      "Saving to: ‘./data/carn_sr/DIV2K_train_HR.zip’\n",
      "\n",
      "DIV2K_train_HR.zip  100%[===================>]   3.29G  18.2MB/s    in 3m 8s   \n",
      "\n",
      "2024-02-04 21:05:14 (17.9 MB/s) - ‘./data/carn_sr/DIV2K_train_HR.zip’ saved [3530603713/3530603713]\n",
      "\n",
      "--2024-02-04 21:05:14--  http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.zip\n",
      "Resolving data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)... 2001:67c:10ec:36c2::178, 129.132.52.178\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:80... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.zip [following]\n",
      "--2024-02-04 21:05:14--  https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.zip\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 925390592 (883M) [application/zip]\n",
      "Saving to: ‘./data/carn_sr/DIV2K_train_LR_bicubic_X2.zip’\n",
      "\n",
      "DIV2K_train_LR_bicu 100%[===================>] 882.52M  18.2MB/s    in 50s     \n",
      "\n",
      "2024-02-04 21:06:05 (17.7 MB/s) - ‘./data/carn_sr/DIV2K_train_LR_bicubic_X2.zip’ saved [925390592/925390592]\n",
      "\n",
      "--2024-02-04 21:06:05--  http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip\n",
      "Resolving data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)... 2001:67c:10ec:36c2::178, 129.132.52.178\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:80... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip [following]\n",
      "--2024-02-04 21:06:05--  https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 448993893 (428M) [application/zip]\n",
      "Saving to: ‘./data/carn_sr/DIV2K_valid_HR.zip’\n",
      "\n",
      "DIV2K_valid_HR.zip  100%[===================>] 428.19M  18.0MB/s    in 25s     \n",
      "\n",
      "2024-02-04 21:06:31 (17.3 MB/s) - ‘./data/carn_sr/DIV2K_valid_HR.zip’ saved [448993893/448993893]\n",
      "\n",
      "--2024-02-04 21:06:31--  http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip\n",
      "Resolving data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)... 2001:67c:10ec:36c2::178, 129.132.52.178\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:80... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip [following]\n",
      "--2024-02-04 21:06:31--  https://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip\n",
      "Connecting to data.vision.ee.ethz.ch (data.vision.ee.ethz.ch)|2001:67c:10ec:36c2::178|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 117763600 (112M) [application/zip]\n",
      "Saving to: ‘./data/carn_sr/DIV2K_valid_LR_bicubic_X2.zip’\n",
      "\n",
      "DIV2K_valid_LR_bicu 100%[===================>] 112.31M  17.8MB/s    in 7.4s    \n",
      "\n",
      "2024-02-04 21:06:40 (15.1 MB/s) - ‘./data/carn_sr/DIV2K_valid_LR_bicubic_X2.zip’ saved [117763600/117763600]\n",
      "\n",
      "\n",
      "4 archives were successfully processed.\n"
     ]
    }
   ],
   "source": [
    "# Download training dataset\n",
    "!wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip -P ./data/carn_sr\n",
    "!wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_LR_bicubic_X2.zip -P ./data/carn_sr\n",
    "!wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_HR.zip -P ./data/carn_sr\n",
    "!wget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_valid_LR_bicubic_X2.zip -P ./data/carn_sr\n",
    "\n",
    "!unzip  -q './data/carn_sr/*.zip' -d ./data/carn_sr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e67db60d-41b0-45de-90f2-e70b76c76241",
   "metadata": {},
   "source": [
    "#### To accelerate training, first convert training images to h5 format as follow (h5py module has to be installed)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a195c9a5-6933-4bae-87a9-9115794b7b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./data/carn_sr/DIV2K_train_HR/0001.png\n",
      "./data/carn_sr/DIV2K_train_HR/0002.png\n",
      "./data/carn_sr/DIV2K_train_HR/0003.png\n",
      "./data/carn_sr/DIV2K_train_HR/0004.png\n",
      "......\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import glob\n",
    "import h5py\n",
    "import cv2\n",
    "import numpy as np\n",
    "from tqdm import tqdm \n",
    "\n",
    "dataset_dir = \"./data/carn_sr/\"\n",
    "dataset_type = \"train\"\n",
    "\n",
    "f = h5py.File(os.path.join(dataset_dir, \"DIV2K_{}.h5\".format(dataset_type)), \"w\")\n",
    "dt = h5py.special_dtype(vlen=np.dtype('uint8'))\n",
    "\n",
    "for subdir in [\"HR\", \"X2\", \"X3\", \"X4\"]:\n",
    "    if subdir in [\"HR\"]:\n",
    "        im_paths = glob.glob(os.path.join(dataset_dir, \n",
    "                                          \"DIV2K_{}_HR\".format(dataset_type), \n",
    "                                          \"*.png\"))\n",
    "\n",
    "    else:\n",
    "        im_paths = glob.glob(os.path.join(dataset_dir, \n",
    "                                          \"DIV2K_{}_LR_bicubic\".format(dataset_type), \n",
    "                                          subdir, \"*.png\"))\n",
    "    im_paths.sort()\n",
    "    grp = f.create_group(subdir)\n",
    "\n",
    "    for i, path in enumerate(im_paths):\n",
    "        im = cv2.imread(path)\n",
    "        print(path)\n",
    "        grp.create_dataset(str(i), data=im)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6eb3c080-7ff8-4445-9ea3-7f6ee47d1c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import h5py\n",
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import torch.utils.data as data\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "def random_crop(hr, lr, size, scale):\n",
    "    h, w = lr.shape[:-1]\n",
    "    x = random.randint(0, w-size)\n",
    "    y = random.randint(0, h-size)\n",
    "\n",
    "    hsize = size*scale\n",
    "    hx, hy = x*scale, y*scale\n",
    "\n",
    "    crop_lr = lr[y:y+size, x:x+size].copy()\n",
    "    crop_hr = hr[hy:hy+hsize, hx:hx+hsize].copy()\n",
    "\n",
    "    return crop_hr, crop_lr\n",
    "\n",
    "\n",
    "def random_flip_and_rotate(im1, im2):\n",
    "    if random.random() < 0.5:\n",
    "        im1 = np.flipud(im1)\n",
    "        im2 = np.flipud(im2)\n",
    "\n",
    "    if random.random() < 0.5:\n",
    "        im1 = np.fliplr(im1)\n",
    "        im2 = np.fliplr(im2)\n",
    "\n",
    "    angle = random.choice([0, 1, 2, 3])\n",
    "    im1 = np.rot90(im1, angle)\n",
    "    im2 = np.rot90(im2, angle)\n",
    "\n",
    "    # have to copy before be called by transform function\n",
    "    return im1.copy(), im2.copy()\n",
    "\n",
    "\n",
    "class TrainDataset(data.Dataset):\n",
    "    def __init__(self, path, size, scale):\n",
    "        super(TrainDataset, self).__init__()\n",
    "\n",
    "        self.size = size\n",
    "        h5f = h5py.File(path, \"r\")\n",
    "        \n",
    "        self.hr = [v[:] for v in h5f[\"HR\"].values()]\n",
    "        # perform multi-scale training\n",
    "        if scale == 0:\n",
    "            self.scale = [2, 3, 4]\n",
    "            self.lr = [[v[:] for v in h5f[\"X{}\".format(i)].values()] for i in self.scale]\n",
    "        else:\n",
    "            self.scale = [scale]\n",
    "            self.lr = [[v[:] for v in h5f[\"X{}\".format(scale)].values()]]\n",
    "        \n",
    "        h5f.close()\n",
    "\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.ToTensor()\n",
    "        ])\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        size = self.size\n",
    "\n",
    "        item = [(self.hr[index], self.lr[i][index]) for i, _ in enumerate(self.lr)]\n",
    "        item = [random_crop(hr, lr, size, self.scale[i]) for i, (hr, lr) in enumerate(item)]\n",
    "        item = [random_flip_and_rotate(hr, lr) for hr, lr in item]\n",
    "        \n",
    "        return [(self.transform(hr), self.transform(lr)) for hr, lr in item]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.hr)\n",
    "        \n",
    "\n",
    "class TestDataset(data.Dataset):\n",
    "    def __init__(self, dirname, scale):\n",
    "        super(TestDataset, self).__init__()\n",
    "\n",
    "        self.name  = dirname.split(\"/\")[-1]\n",
    "        self.scale = scale\n",
    "        \n",
    "        if \"DIV\" in self.name:\n",
    "            self.hr = glob.glob(os.path.join(\"{}_HR\".format(dirname), \"*.png\"))\n",
    "            self.lr = glob.glob(os.path.join(\"{}_LR_bicubic\".format(dirname), \n",
    "                                             \"X{}/*.png\".format(scale)))\n",
    "        else:\n",
    "            all_files = glob.glob(os.path.join(dirname, \"x{}/*.png\".format(scale)))\n",
    "            self.hr = [name for name in all_files if \"HR\" in name]\n",
    "            self.lr = [name for name in all_files if \"LR\" in name]\n",
    "\n",
    "        self.hr.sort()\n",
    "        self.lr.sort()\n",
    "\n",
    "        self.transform = transforms.Compose([\n",
    "            transforms.ToTensor()\n",
    "        ])\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        hr = Image.open(self.hr[index])\n",
    "        lr = Image.open(self.lr[index])\n",
    "\n",
    "        hr = hr.convert(\"RGB\")\n",
    "        lr = lr.convert(\"RGB\")\n",
    "        filename = self.hr[index].split(\"/\")[-1]\n",
    "\n",
    "        return self.transform(hr), self.transform(lr), filename\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.hr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1b79b4a3-7818-42b6-b17e-1e3610827b01",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "train_data = TrainDataset('./data/carn_sr/DIV2K_train.h5', scale=2, size=64)\n",
    "train_loader = DataLoader(train_data, batch_size=64, num_workers=4, shuffle=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c868f22-f415-4d83-aa3d-611ea5cf17fd",
   "metadata": {},
   "source": [
    "#### Download the test datasets following [CARN repo](https://github.com/nmhkahn/CARN-pytorch)\n",
    "\n",
    "\n",
    "To this tutorial, we place them into `./data/SR_benchmark`, which includes `B100`, `Set5`, `Set14`, and `Urban100`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a6b9c8-4591-4171-902c-b8b454686144",
   "metadata": {},
   "source": [
    "### Step 4. Setup HESSO optimizer\n",
    "\n",
    "The following main hyperparameters need to be taken care.\n",
    "\n",
    "- `variant`: The optimizer that is used for training the baseline full model. Currently support `sgd`, `adam` and `adamw`.\n",
    "- `lr`: The initial learning rate.\n",
    "- `weight_decay`: Weight decay as standard DNN optimization.\n",
    "- `target_group_sparsity`: The target group sparsity, typically higher group sparsity refers to more FLOPs and model size reduction, meanwhile may regress model performance more.\n",
    "- `start_pruning_steps`: The number of steps that **starts** to prune.\n",
    "- `pruning_steps`: The number of steps that **finishes** pruning (reach `target_group_sparsity`) after `start_pruning_steps`.\n",
    "- `pruning_periods`:  Incrementally produce the group sparsity equally among pruning periods.\n",
    "\n",
    "We empirically suggest `start_pruning_steps` as 1/10 of total number of training steps. `pruning_steps` until 1/4 or 1/5 of total number of training steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d41591a8-1838-4c6d-b9fc-7defa985c08c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setup HESSO\n",
      "Target redundant groups per period:  [84, 84, 84, 84, 84, 84, 84, 84, 84, 88]\n"
     ]
    }
   ],
   "source": [
    "optimizer = oto.hesso(\n",
    "    variant='adam', \n",
    "    lr=1e-4, \n",
    "    target_group_sparsity=0.6,\n",
    "    start_pruning_step=60000,\n",
    "    pruning_periods=10,\n",
    "    pruning_steps=60000,\n",
    "    importance_score_criteria='default'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "38714afe-6b7c-4f44-b25f-9253edb80c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "from skimage.metrics import peak_signal_noise_ratio\n",
    "from skimage.color import rgb2ycbcr\n",
    "import math\n",
    "\n",
    "# helpers\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from pylab import rcParams\n",
    "rcParams['figure.figsize'] = 15, 10\n",
    "\n",
    "def psnr(im1, im2):\n",
    "    def im2double(im):\n",
    "        min_val, max_val = 0, 255\n",
    "        out = (im.astype(np.float64)-min_val) / (max_val-min_val)\n",
    "        return out\n",
    "\n",
    "    im1 = im2double(im1)\n",
    "    im2 = im2double(im2)\n",
    "    psnr = peak_signal_noise_ratio(im1, im2, data_range=1)\n",
    "    return psnr\n",
    "\n",
    "def display1(img):\n",
    "    plt.imshow(img, interpolation=\"nearest\")\n",
    "    plt.show()\n",
    "\n",
    "def evaluate(model, test_data_dir, scale=2):\n",
    "    shave = 20\n",
    "    mean_psnr = 0\n",
    "    model.eval()\n",
    "    \n",
    "    test_data   = TestDataset(test_data_dir, scale=scale)\n",
    "    test_loader = DataLoader(test_data,\n",
    "                             batch_size=1,\n",
    "                             num_workers=1,\n",
    "                             shuffle=False)\n",
    "\n",
    "    for step, inputs in enumerate(test_loader):\n",
    "        hr = inputs[0].squeeze(0)\n",
    "        lr = inputs[1].squeeze(0)\n",
    "        name = inputs[2][0]\n",
    "\n",
    "        h, w = lr.size()[1:]\n",
    "        h_half, w_half = int(h/2), int(w/2)\n",
    "        h_chop, w_chop = h_half + shave, w_half + shave\n",
    "\n",
    "        # split large image to 4 patch to avoid OOM error\n",
    "        lr_patch = torch.FloatTensor(4, 3, h_chop, w_chop)\n",
    "        lr_patch[0].copy_(lr[:, 0:h_chop, 0:w_chop])\n",
    "        lr_patch[1].copy_(lr[:, 0:h_chop, w-w_chop:w])\n",
    "        lr_patch[2].copy_(lr[:, h-h_chop:h, 0:w_chop])\n",
    "        lr_patch[3].copy_(lr[:, h-h_chop:h, w-w_chop:w])\n",
    "        lr_patch = lr_patch.cuda()\n",
    "        \n",
    "        # run refine process in here!\n",
    "        sr = model(lr_patch, scale).data\n",
    "        \n",
    "        h, h_half, h_chop = h*scale, h_half*scale, h_chop*scale\n",
    "        w, w_half, w_chop = w*scale, w_half*scale, w_chop*scale\n",
    "        \n",
    "        # merge splited patch images\n",
    "        result = torch.FloatTensor(3, h, w).cuda()\n",
    "        result[:, 0:h_half, 0:w_half].copy_(sr[0, :, 0:h_half, 0:w_half])\n",
    "        result[:, 0:h_half, w_half:w].copy_(sr[1, :, 0:h_half, w_chop-w+w_half:w_chop])\n",
    "        result[:, h_half:h, 0:w_half].copy_(sr[2, :, h_chop-h+h_half:h_chop, 0:w_half])\n",
    "        result[:, h_half:h, w_half:w].copy_(sr[3, :, h_chop-h+h_half:h_chop, w_chop-w+w_half:w_chop])\n",
    "        sr = result\n",
    "\n",
    "        hr = hr.cpu().mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()\n",
    "        sr = sr.cpu().mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()\n",
    "        \n",
    "        # evaluate PSNR\n",
    "        bnd = scale\n",
    "        im1 = hr[bnd:-bnd, bnd:-bnd]\n",
    "        im2 = sr[bnd:-bnd, bnd:-bnd]\n",
    "\n",
    "        # change to evaluate y-channel, based on a reproduction open issue in CARN \n",
    "        im1_y = rgb2ycbcr(im1)[..., 0]\n",
    "        im2_y = rgb2ycbcr(im2)[..., 0]\n",
    "\n",
    "        mean_psnr += psnr(im1_y, im2_y) / len(test_data)\n",
    "    return mean_psnr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb65ed8e",
   "metadata": {},
   "source": [
    "### Step 4. Train and prune via OTO HESSO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "adb1267d-1792-43a3-97c1-649b1ea737b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step: 0, loss: 0.0002, norm_all:952.65, grp_sparsity: 0.00, norm_import: 952.65, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 12.6870, Set14: 11.7356, Urban100: 11.0549.\n",
      "Step: 1000, loss: 0.0433, norm_all:997.35, grp_sparsity: 0.00, norm_import: 997.35, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 27.9681, Set14: 27.8497, Urban100: 25.0936.\n",
      "Step: 2000, loss: 0.0204, norm_all:1019.24, grp_sparsity: 0.00, norm_import: 1019.24, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 28.9096, Set14: 28.8339, Urban100: 26.0179.\n",
      "Step: 3000, loss: 0.0188, norm_all:1030.98, grp_sparsity: 0.00, norm_import: 1030.98, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.0512, Set14: 29.0506, Urban100: 26.1827.\n",
      "Step: 4000, loss: 0.0182, norm_all:1040.63, grp_sparsity: 0.00, norm_import: 1040.63, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.1222, Set14: 29.2128, Urban100: 26.3156.\n",
      "Step: 5000, loss: 0.0178, norm_all:1049.91, grp_sparsity: 0.00, norm_import: 1049.91, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.2313, Set14: 29.3998, Urban100: 26.4735.\n",
      "Step: 6000, loss: 0.0175, norm_all:1059.47, grp_sparsity: 0.00, norm_import: 1059.47, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.3695, Set14: 29.5718, Urban100: 26.6359.\n",
      "Step: 7000, loss: 0.0171, norm_all:1068.58, grp_sparsity: 0.00, norm_import: 1068.58, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.4573, Set14: 29.6965, Urban100: 26.7836.\n",
      "Step: 8000, loss: 0.0169, norm_all:1078.04, grp_sparsity: 0.00, norm_import: 1078.04, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.5872, Set14: 29.8749, Urban100: 26.9799.\n",
      "Step: 9000, loss: 0.0166, norm_all:1088.24, grp_sparsity: 0.00, norm_import: 1088.24, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.6178, Set14: 29.9242, Urban100: 27.1206.\n",
      "Step: 10000, loss: 0.0163, norm_all:1098.21, grp_sparsity: 0.00, norm_import: 1098.21, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.7323, Set14: 30.0821, Urban100: 27.3241.\n",
      "Step: 11000, loss: 0.0161, norm_all:1107.49, grp_sparsity: 0.00, norm_import: 1107.49, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.8584, Set14: 30.2452, Urban100: 27.5413.\n",
      "Step: 12000, loss: 0.0160, norm_all:1115.67, grp_sparsity: 0.00, norm_import: 1115.67, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.9285, Set14: 30.3140, Urban100: 27.6706.\n",
      "Step: 13000, loss: 0.0159, norm_all:1123.71, grp_sparsity: 0.00, norm_import: 1123.71, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.9543, Set14: 30.3813, Urban100: 27.7880.\n",
      "Step: 14000, loss: 0.0157, norm_all:1131.30, grp_sparsity: 0.00, norm_import: 1131.30, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 29.9931, Set14: 30.4252, Urban100: 27.8702.\n",
      "Step: 15000, loss: 0.0155, norm_all:1138.27, grp_sparsity: 0.00, norm_import: 1138.27, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.0467, Set14: 30.5025, Urban100: 27.9945.\n",
      "Step: 16000, loss: 0.0155, norm_all:1145.33, grp_sparsity: 0.00, norm_import: 1145.33, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.0483, Set14: 30.4854, Urban100: 28.0387.\n",
      "Step: 17000, loss: 0.0154, norm_all:1152.26, grp_sparsity: 0.00, norm_import: 1152.26, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1199, Set14: 30.5970, Urban100: 28.1424.\n",
      "Step: 18000, loss: 0.0152, norm_all:1159.20, grp_sparsity: 0.00, norm_import: 1159.20, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.0787, Set14: 30.5648, Urban100: 28.2031.\n",
      "Step: 19000, loss: 0.0152, norm_all:1165.93, grp_sparsity: 0.00, norm_import: 1165.93, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.0928, Set14: 30.5753, Urban100: 28.2311.\n",
      "Step: 20000, loss: 0.0152, norm_all:1172.60, grp_sparsity: 0.00, norm_import: 1172.60, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1276, Set14: 30.6350, Urban100: 28.3137.\n",
      "Step: 21000, loss: 0.0150, norm_all:1178.71, grp_sparsity: 0.00, norm_import: 1178.71, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1964, Set14: 30.6976, Urban100: 28.4022.\n",
      "Step: 22000, loss: 0.0151, norm_all:1184.96, grp_sparsity: 0.00, norm_import: 1184.96, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1852, Set14: 30.6860, Urban100: 28.4411.\n",
      "Step: 23000, loss: 0.0150, norm_all:1191.08, grp_sparsity: 0.00, norm_import: 1191.08, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1839, Set14: 30.6829, Urban100: 28.4668.\n",
      "Step: 24000, loss: 0.0149, norm_all:1197.10, grp_sparsity: 0.00, norm_import: 1197.10, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.2625, Set14: 30.7774, Urban100: 28.5556.\n",
      "Step: 25000, loss: 0.0149, norm_all:1203.05, grp_sparsity: 0.00, norm_import: 1203.05, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.2071, Set14: 30.7191, Urban100: 28.5328.\n",
      "Step: 26000, loss: 0.0149, norm_all:1208.94, grp_sparsity: 0.00, norm_import: 1208.94, norm_redund: 0.00, num_grp_import: 1408, num_grp_redund: 0\n",
      "Val PSNR: B100: 30.1811, Set14: 30.6557, Urban100: 28.4756.\n",
      "......\n",
      "Step: 597000, loss: 0.0142, norm_all:973.37, grp_sparsity: 0.60, norm_import: 973.37, norm_redund: 0.00, num_grp_import: 564, num_grp_redund: 844\n",
      "Val PSNR: B100: 31.8130, Set14: 33.0487, Urban100: 30.7883.\n",
      "Step: 598000, loss: 0.0141, norm_all:973.37, grp_sparsity: 0.60, norm_import: 973.37, norm_redund: 0.00, num_grp_import: 564, num_grp_redund: 844\n",
      "Val PSNR: B100: 31.8130, Set14: 33.0487, Urban100: 30.7883.\n",
      "Step: 599000, loss: 0.0141, norm_all:973.37, grp_sparsity: 0.60, norm_import: 973.37, norm_redund: 0.00, num_grp_import: 564, num_grp_redund: 844\n",
      "Val PSNR: B100: 31.8130, Set14: 33.0487, Urban100: 30.7883.\n",
      "Step: 600000, loss: 0.0141, norm_all:973.37, grp_sparsity: 0.60, norm_import: 973.37, norm_redund: 0.00, num_grp_import: 564, num_grp_redund: 844\n",
      "Val PSNR: B100: 31.8130, Set14: 33.0487, Urban100: 30.7883.\n"
     ]
    }
   ],
   "source": [
    "# max_step and lr_decay_step are the same as carn official repo. \n",
    "max_step = 600000\n",
    "lr_decay_step = 400000 \n",
    "print_interval = 1000\n",
    "loss_fn = torch.nn.L1Loss()\n",
    "step = 0\n",
    "f_avg_val = 0.0\n",
    "learning_rate = optimizer.get_learning_rate()\n",
    "\n",
    "while True:\n",
    "    for inputs in train_loader:\n",
    "        model.train()\n",
    "        hr, lr = inputs[-1][0], inputs[-1][1]\n",
    "        hr, lr = hr.cuda(), lr.cuda()\n",
    "        sr = model(lr, scale)\n",
    "        loss = loss_fn(sr, hr)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        torch.nn.utils.clip_grad_norm(model.parameters(), 10.0)\n",
    "        optimizer.step()\n",
    "\n",
    "        learning_rate = learning_rate * (0.5 ** (step // lr_decay_step))\n",
    "        optimizer.set_learning_rate(learning_rate)\n",
    "        f_avg_val += loss.item()  \n",
    "        if step % print_interval == 0:\n",
    "            group_sparsity, param_norm, _ = optimizer.compute_group_sparsity_param_norm()\n",
    "            norm_important, norm_redundant, num_grps_important, num_grps_redundant = optimizer.compute_norm_groups()\n",
    "            print(\"Step: {step}, loss: {f:.4f}, norm_all:{param_norm:.2f}, grp_sparsity: {gs:.2f}, norm_import: {norm_import:.2f}, norm_redund: {norm_redund:.2f}, num_grp_import: {num_grps_import}, num_grp_redund: {num_grps_redund}\"\\\n",
    "                 .format(step=step, f=f_avg_val/print_interval, param_norm=param_norm, gs=group_sparsity, norm_import=norm_important, \\\n",
    "                         norm_redund=norm_redundant, num_grps_import=num_grps_important, num_grps_redund=num_grps_redundant\n",
    "                ))\n",
    "            f_avg_val = 0.0\n",
    "            psnr_B100 = evaluate(model, './data/SR_benchmark/B100', scale=2)\n",
    "            psnr_Set14 = evaluate(model, './data/SR_benchmark/Set14', scale=2)\n",
    "            psnr_Urban100 = evaluate(model, './data/SR_benchmark/Urban100', scale=2)\n",
    "            print(\"Val PSNR: B100: {p_b:.4f}, Set14: {p_s:.4f}, Urban100: {p_u:.4f}.\".format(p_b=psnr_B100, p_s=psnr_Set14, p_u=psnr_Urban100))\n",
    "        step += 1\n",
    "    if step > max_step: \n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa0a8927-0aec-4a5a-b342-a31df759b481",
   "metadata": {},
   "source": [
    "### (Optional) get FLOPs and number of parameters for full model\n",
    "\n",
    "It must be excuted before oto.construct_subnet(). Otherwise, the API would calculate based on the pruned model.\n",
    "\n",
    "`oto.compute_flops()` returns a dictionary with each node group' FLOPs, use `total` can get the total FLOPs for the whole DNN. \n",
    "\n",
    "`in_million` and `in_billion` argument could scale the numbers in the unit of million or billion. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c97c2f0f-5b27-4163-afe2-4a795e0859fc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Full FLOPs(M) 48527.66822400001 full Number of parameters (M) 0.964187\n"
     ]
    }
   ],
   "source": [
    "full_flops = oto.compute_flops(in_million=True)['total'] \n",
    "full_num_params = oto.compute_num_params(in_million=True)\n",
    "\n",
    "print(\"Full FLOPs(M)\", full_flops, \"full Number of parameters (M)\", full_num_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdd6b7a5-7afa-49f9-a2e6-bf63ac6eb7ac",
   "metadata": {},
   "source": [
    "### Step 5. Get pruned model in torch format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1bed11cf-2dc1-4120-afc4-f7ad62cd6f83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# By default OTO will construct subnet by the last checkpoint. If intermedia ckpt reaches the best performance,\n",
    "# need to reinitialize OTO instance\n",
    "# oto = OTO(torch.load(ckpt_path), dummy_input)\n",
    "# then construct subnetwork\n",
    "oto.construct_subnet(out_dir='./cache')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a18bca-58b9-49f2-9a8e-1e5a159cce82",
   "metadata": {},
   "source": [
    "### (Optional) Check the compressed model size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "363333ae-9042-4f6b-8438-fe5bfc566593",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of full model     :  0.0036377040669322014 GBs\n",
      "Size of compress model :  0.0006903214380145073 GBs\n"
     ]
    }
   ],
   "source": [
    "full_model_size = os.stat(oto.full_group_sparse_model_path)\n",
    "compressed_model_size = os.stat(oto.compressed_model_path)\n",
    "print(\"Size of full model     : \", full_model_size.st_size / (1024 ** 3), \"GBs\")\n",
    "print(\"Size of compress model : \", compressed_model_size.st_size / (1024 ** 3), \"GBs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c911fbd-23a9-4632-ae77-f47af9056a4d",
   "metadata": {},
   "source": [
    "### (Optional) get FLOPs and number of parameters for compressed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7e0ceb2d-9899-427e-b906-08fcace81a2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pruned FLOPs(M) 8913.716224, pruned Number of parameters (M) 0.173537\n"
     ]
    }
   ],
   "source": [
    "pruned_flops = oto.compute_flops(in_million=True)['total'] \n",
    "pruned_num_params = oto.compute_num_params(in_million=True)\n",
    "\n",
    "print(\"Pruned FLOPs(M)\", pruned_flops, \"pruned Number of parameters (M)\", pruned_num_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd16afa6-8c77-4b00-a5c6-b53e5e36cfb6",
   "metadata": {},
   "source": [
    "### (Optional) Check the pruned model and full model difference. \n",
    "#### # Both full and pruned model should return the exact same output given the same input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e1df449a-3210-4b71-b582-e1bc0f770bc7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum output difference  5.960464477539062e-07\n"
     ]
    }
   ],
   "source": [
    "full_model = torch.load(oto.full_group_sparse_model_path).cpu()\n",
    "compressed_model = torch.load(oto.compressed_model_path).cpu()\n",
    "\n",
    "full_output = full_model(dummy_input, scale)\n",
    "compressed_output = compressed_model(dummy_input, scale)\n",
    "\n",
    "max_output_diff = torch.max(torch.abs(full_output - compressed_output))\n",
    "print(\"Maximum output difference \", str(max_output_diff.item()))"
   ]
  }
 ],
 "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}