{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "In this tutorial, we give an end-to-end demo of compressing [MobileNetV2](https://arxiv.org/abs/1801.04381) for finegrained classification using [NNI Pruners](https://nni.readthedocs.io/en/stable/Compression/pruning.html). Although MobileNetV2 is already a highly optimized architecture, we show that we can further reduce its size by over 50% with minimal performance loss using iterative pruning and knowledge distillation. To similate a real usage scenario, we use the [Stanford Dogs](http://vision.stanford.edu/aditya86/ImageNetDogs/) dataset as the target task, and show how to implement and optimize the following steps:\n",
    "* Model pre-training\n",
    "* Pruning\n",
    "* Model Speedup\n",
    "* Finetuning the pruned model\n",
    "\n",
    "Also, we will compare our approach with some baseline channel compression schemes defined by the authors of the MobileNets, and show that NNI pruners can provide a superior performance while being easy-to-use. We release this notebook along with our code under the folder `examples/model_compress/pruning/mobilenet_end2end/`.\n",
    "<div>\n",
    "<img src=\"final_performance.png\" width=\"750\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import xml\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "\n",
    "from nni.compression.pytorch import ModelSpeedup\n",
    "from nni.compression.pytorch.utils.counter import count_flops_params\n",
    "\n",
    "from utils import create_model, get_dataloader\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "num_workers = 16\n",
    "torch.set_num_threads(num_workers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Background\n",
    "### Pruning MobileNetV2\n",
    "The main building block of MobileNetV2 is \"inverted residual blocks\", where a pointwise convolution first projects into a feature map with higher channels,following a depthwise convolution, and a pointwise convolution with linear activation that projects into a features map with less channels (thus called \"inverted residuals and linear bottlenecks\"). With 11 such blocks stacked together, the entire model has 3.4M parameters and takes up about 10M storage space (this number is platform-dependent).\n",
    "\n",
    "<div>\n",
    "<img src=\"mobilenet.png\" width=\"300\"/>\n",
    "</div>\n",
    "\n",
    "Now we consider compressing MobileNetV2 by **filter pruning** (also called channel pruning). Recall that in genernal, a $k\\times k$ convolutional kernel has the weight with shape $(out\\_channel, \\frac{in\\_channel}{groups}, k, k)$. If the input has shape $(B, in\\_channel, H, W)$, the convolutional layer's output (with padding) would have shape $(B, out\\_channel, H, W)$. Suppose we remove $M$ filters from this layer, then weight would have shape $(out\\_channel-M, \\frac{in\\_channel}{groups}, k, k)$, and the output would then have shape $(B, out\\_channel - M, H, W)$. Further, we have the following observations:\n",
    "* The model's number of parameters is directly reduced by $M\\times \\frac{in\\_channel}{groups} \\times k \\times k$.\n",
    "* We are performing structured pruning, as each filter's weight elements are adjacent. Compared to unstructured pruning (or fine-grained pruning), structured pruning generally allows us to directly remove weights and their connections from the network, resulting in greater compression and speed-up. For this reason, in this tutorial we solely focus on filter-level pruning. \n",
    "* Since the output channel is shrinked, we can also remove weights from the next layer corresponding to these channel dimensions. In NNI, the pruner prunes the weights by just setting the weight values to zero, and then the [ModelSpeedup](https://nni.readthedocs.io/en/stable/Compression/ModelSpeedup.html) tool infers the weight relations and removes pruned weights and connections, which we will also demonstrate later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/v-diwu4/.cache/torch/hub/pytorch_vision_v0.8.1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MobileNetV2(\n",
      "  (features): Sequential(\n",
      "    (0): ConvBNActivation(\n",
      "      (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU6(inplace=True)\n",
      "    )\n",
      "    (1): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(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",
      "    (2): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), 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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(144, 144, kernel_size=(3, 3), stride=(2, 2), 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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(192, 192, kernel_size=(3, 3), stride=(2, 2), 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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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",
      "    (10): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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",
      "    (11): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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",
      "    (12): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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",
      "    (13): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(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",
      "    (14): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(576, 576, kernel_size=(3, 3), stride=(2, 2), 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(576, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (15): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (16): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(960, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (17): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(160, 960, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(960, 960, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=960, bias=False)\n",
      "          (1): BatchNorm2d(960, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(960, 320, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (18): ConvBNActivation(\n",
      "      (0): Conv2d(320, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU6(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (classifier): Sequential(\n",
      "    (0): Dropout(p=0.2, inplace=False)\n",
      "    (1): Linear(in_features=1280, out_features=1000, bias=True)\n",
      "  )\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ")\n",
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "| Index |         Name         |  Type  |    Weight Shape   |  FLOPs   | #Params |\n",
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "|   0   |     features.0.0     | Conv2d |   (32, 3, 3, 3)   | 10838016 |   864   |\n",
      "|   1   | features.1.conv.0.0  | Conv2d |   (32, 1, 3, 3)   | 3612672  |   288   |\n",
      "|   2   |  features.1.conv.1   | Conv2d |   (16, 32, 1, 1)  | 6422528  |   512   |\n",
      "|   3   | features.2.conv.0.0  | Conv2d |   (96, 16, 1, 1)  | 19267584 |   1536  |\n",
      "|   4   | features.2.conv.1.0  | Conv2d |   (96, 1, 3, 3)   | 2709504  |   864   |\n",
      "|   5   |  features.2.conv.2   | Conv2d |   (24, 96, 1, 1)  | 7225344  |   2304  |\n",
      "|   6   | features.3.conv.0.0  | Conv2d |  (144, 24, 1, 1)  | 10838016 |   3456  |\n",
      "|   7   | features.3.conv.1.0  | Conv2d |   (144, 1, 3, 3)  | 4064256  |   1296  |\n",
      "|   8   |  features.3.conv.2   | Conv2d |  (24, 144, 1, 1)  | 10838016 |   3456  |\n",
      "|   9   | features.4.conv.0.0  | Conv2d |  (144, 24, 1, 1)  | 10838016 |   3456  |\n",
      "|   10  | features.4.conv.1.0  | Conv2d |   (144, 1, 3, 3)  | 1016064  |   1296  |\n",
      "|   11  |  features.4.conv.2   | Conv2d |  (32, 144, 1, 1)  | 3612672  |   4608  |\n",
      "|   12  | features.5.conv.0.0  | Conv2d |  (192, 32, 1, 1)  | 4816896  |   6144  |\n",
      "|   13  | features.5.conv.1.0  | Conv2d |   (192, 1, 3, 3)  | 1354752  |   1728  |\n",
      "|   14  |  features.5.conv.2   | Conv2d |  (32, 192, 1, 1)  | 4816896  |   6144  |\n",
      "|   15  | features.6.conv.0.0  | Conv2d |  (192, 32, 1, 1)  | 4816896  |   6144  |\n",
      "|   16  | features.6.conv.1.0  | Conv2d |   (192, 1, 3, 3)  | 1354752  |   1728  |\n",
      "|   17  |  features.6.conv.2   | Conv2d |  (32, 192, 1, 1)  | 4816896  |   6144  |\n",
      "|   18  | features.7.conv.0.0  | Conv2d |  (192, 32, 1, 1)  | 4816896  |   6144  |\n",
      "|   19  | features.7.conv.1.0  | Conv2d |   (192, 1, 3, 3)  |  338688  |   1728  |\n",
      "|   20  |  features.7.conv.2   | Conv2d |  (64, 192, 1, 1)  | 2408448  |  12288  |\n",
      "|   21  | features.8.conv.0.0  | Conv2d |  (384, 64, 1, 1)  | 4816896  |  24576  |\n",
      "|   22  | features.8.conv.1.0  | Conv2d |   (384, 1, 3, 3)  |  677376  |   3456  |\n",
      "|   23  |  features.8.conv.2   | Conv2d |  (64, 384, 1, 1)  | 4816896  |  24576  |\n",
      "|   24  | features.9.conv.0.0  | Conv2d |  (384, 64, 1, 1)  | 4816896  |  24576  |\n",
      "|   25  | features.9.conv.1.0  | Conv2d |   (384, 1, 3, 3)  |  677376  |   3456  |\n",
      "|   26  |  features.9.conv.2   | Conv2d |  (64, 384, 1, 1)  | 4816896  |  24576  |\n",
      "|   27  | features.10.conv.0.0 | Conv2d |  (384, 64, 1, 1)  | 4816896  |  24576  |\n",
      "|   28  | features.10.conv.1.0 | Conv2d |   (384, 1, 3, 3)  |  677376  |   3456  |\n",
      "|   29  |  features.10.conv.2  | Conv2d |  (64, 384, 1, 1)  | 4816896  |  24576  |\n",
      "|   30  | features.11.conv.0.0 | Conv2d |  (384, 64, 1, 1)  | 4816896  |  24576  |\n",
      "|   31  | features.11.conv.1.0 | Conv2d |   (384, 1, 3, 3)  |  677376  |   3456  |\n",
      "|   32  |  features.11.conv.2  | Conv2d |  (96, 384, 1, 1)  | 7225344  |  36864  |\n",
      "|   33  | features.12.conv.0.0 | Conv2d |  (576, 96, 1, 1)  | 10838016 |  55296  |\n",
      "|   34  | features.12.conv.1.0 | Conv2d |   (576, 1, 3, 3)  | 1016064  |   5184  |\n",
      "|   35  |  features.12.conv.2  | Conv2d |  (96, 576, 1, 1)  | 10838016 |  55296  |\n",
      "|   36  | features.13.conv.0.0 | Conv2d |  (576, 96, 1, 1)  | 10838016 |  55296  |\n",
      "|   37  | features.13.conv.1.0 | Conv2d |   (576, 1, 3, 3)  | 1016064  |   5184  |\n",
      "|   38  |  features.13.conv.2  | Conv2d |  (96, 576, 1, 1)  | 10838016 |  55296  |\n",
      "|   39  | features.14.conv.0.0 | Conv2d |  (576, 96, 1, 1)  | 10838016 |  55296  |\n",
      "|   40  | features.14.conv.1.0 | Conv2d |   (576, 1, 3, 3)  |  254016  |   5184  |\n",
      "|   41  |  features.14.conv.2  | Conv2d |  (160, 576, 1, 1) | 4515840  |  92160  |\n",
      "|   42  | features.15.conv.0.0 | Conv2d |  (960, 160, 1, 1) | 7526400  |  153600 |\n",
      "|   43  | features.15.conv.1.0 | Conv2d |   (960, 1, 3, 3)  |  423360  |   8640  |\n",
      "|   44  |  features.15.conv.2  | Conv2d |  (160, 960, 1, 1) | 7526400  |  153600 |\n",
      "|   45  | features.16.conv.0.0 | Conv2d |  (960, 160, 1, 1) | 7526400  |  153600 |\n",
      "|   46  | features.16.conv.1.0 | Conv2d |   (960, 1, 3, 3)  |  423360  |   8640  |\n",
      "|   47  |  features.16.conv.2  | Conv2d |  (160, 960, 1, 1) | 7526400  |  153600 |\n",
      "|   48  | features.17.conv.0.0 | Conv2d |  (960, 160, 1, 1) | 7526400  |  153600 |\n",
      "|   49  | features.17.conv.1.0 | Conv2d |   (960, 1, 3, 3)  |  423360  |   8640  |\n",
      "|   50  |  features.17.conv.2  | Conv2d |  (320, 960, 1, 1) | 15052800 |  307200 |\n",
      "|   51  |    features.18.0     | Conv2d | (1280, 320, 1, 1) | 20070400 |  409600 |\n",
      "|   52  |     classifier.1     | Linear |    (1000, 1280)   | 1280000  | 1281000 |\n",
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "FLOPs total: 300774272\n",
      "#Params total: 3470760\n",
      "FLOPs: 300774272, params: 3470760\n"
     ]
    }
   ],
   "source": [
    "# check model architecture\n",
    "model = torch.hub.load('pytorch/vision:v0.8.1', 'mobilenet_v2', pretrained=True).to(device)\n",
    "print(model)\n",
    "\n",
    "# check model FLOPs and parameter counts with NNI utils\n",
    "dummy_input = torch.rand([1, 3, 224, 224]).to(device)\n",
    "flops, params, results = count_flops_params(model, dummy_input)\n",
    "print(f\"FLOPs: {flops}, params: {params}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stanford Dogs\n",
    "\n",
    "The [Stanford Dogs](http://vision.stanford.edu/aditya86/ImageNetDogs/) dataset contains images of 120 breeds of dogs from around the world. It is built using images and annotation from ImageNet for the task of fine-grained image classification. We choose this task to simulate a transfer learning scenario, where a model pre-trained on the ImageNet is further transferred to an often simpler downstream task.\n",
    "\n",
    "To download and prepare the data, please run `prepare_data.sh`, which downloads the images and annotations, and preprocesses the images for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "file_list.mat\n",
      "train_list.mat\n",
      "test_list.mat\n",
      "Directory already exists. Nothing done.\n"
     ]
    }
   ],
   "source": [
    "# Run prepare_data.sh\n",
    "!chmod u+x prepare_data.sh\n",
    "!./prepare_data.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Then, you may run following code block, which shows several instances:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 473656 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show several examples\n",
    "# Code adapted from https://www.kaggle.com/mclikmb4/xception-transfer-learning-120-breeds-83-acc\n",
    "image_path = './data/stanford-dogs/Images/'\n",
    "breed_list = sorted(os.listdir(image_path))\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "for i in range(9):\n",
    "    plt.subplot(331 + i)\n",
    "    breed = np.random.choice(breed_list)\n",
    "    dog = np.random.choice(os.listdir('./data/stanford-dogs/Annotation/' + breed))\n",
    "    img = Image.open(image_path + breed + '/' + dog + '.jpg') \n",
    "    tree = xml.etree.ElementTree.parse('./data/stanford-dogs/Annotation/' + breed + '/' + dog)\n",
    "    root = tree.getroot()\n",
    "    objects = root.findall('object')\n",
    "    plt.imshow(img)\n",
    "    for o in objects:\n",
    "        bndbox = o.find('bndbox')\n",
    "        xmin = int(bndbox.find('xmin').text)\n",
    "        ymin = int(bndbox.find('ymin').text)\n",
    "        xmax = int(bndbox.find('xmax').text)\n",
    "        ymax = int(bndbox.find('ymax').text)\n",
    "        plt.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin])\n",
    "        plt.text(xmin, ymin, o.find('name').text, bbox={'ec': None})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Pre-training\n",
    "First, we obtain a MobileNetV2 model on this task, which will serve as the base model for compression. Unfortunately, although this step is often called model \"pre-training\" in the model compression teminologies, we are actually finetuning a model pre-trained on ImageNet. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This script will save the state dict of the pretrained model to \"./pretrained_mobilenet_v2_torchhub/checkpoint_best.pt\"\n",
    "\n",
    "# %run pretrain.py\n",
    "# %run test.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compression via Pruning\n",
    "In this section, we first demonstrate how to perform channel pruning with NNI pruners in three steps: \n",
    "* defining a config list\n",
    "* creating a Pruner instance\n",
    "* calling `pruner.compress` and `pruner.export_model` to calculate and export masks\n",
    "\n",
    "Then, we demonstrate the common practices after pruning:\n",
    "* model speedup\n",
    "* further finetuning (with or without knowledge distillation)\n",
    "* evaluation\n",
    "\n",
    "Finally, we present a grid search example to find the balance between model performance and the final model size. We include some of our results and discuss our observations. \n",
    "\n",
    "Note that the code blocks in this section are taken from the file `pruning_experiments.py`. You can directly run the file by specifying several command line arguments and see the end-to-end process. You can also run the file to reproduce our experiments. We will discuss that in the last section. \n",
    "\n",
    "### Using NNI Pruners"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nni.algorithms.compression.pytorch.pruning import (\n",
    "    LevelPruner,\n",
    "    SlimPruner,\n",
    "    FPGMPruner,\n",
    "    TaylorFOWeightFilterPruner,\n",
    "    L1FilterPruner,\n",
    "    L2FilterPruner,\n",
    "    AGPPruner,\n",
    "    ActivationMeanRankFilterPruner,\n",
    "    ActivationAPoZRankFilterPruner\n",
    ")\n",
    "\n",
    "pruner_name_to_class = {\n",
    "    'level': LevelPruner,\n",
    "    'l1': L1FilterPruner,\n",
    "    'l2': L2FilterPruner,\n",
    "    'slim': SlimPruner,\n",
    "    'fpgm': FPGMPruner,\n",
    "    'taylor': TaylorFOWeightFilterPruner,\n",
    "    'agp': AGPPruner,\n",
    "    'activationmeanrank': ActivationMeanRankFilterPruner,\n",
    "    'apoz': ActivationAPoZRankFilterPruner\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/v-diwu4/.cache/torch/hub/pytorch_vision_v0.8.1\n"
     ]
    }
   ],
   "source": [
    "# load model from the pretrained checkpoint\n",
    "model_type = 'mobilenet_v2_torchhub'\n",
    "checkpoint = \"./pretrained_mobilenet_v2_torchhub/checkpoint_best.pt\"\n",
    "pretrained = True \n",
    "input_size = 224\n",
    "n_classes = 120\n",
    "\n",
    "model = create_model(model_type=model_type, pretrained=pretrained, n_classes=n_classes,\n",
    "                     input_size=input_size, checkpoint=checkpoint).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-08-31 07:17:21] INFO (nni.compression.pytorch.compressor/MainThread) Model state_dict saved to ./pruned_model.pth\n",
      "[2021-08-31 07:17:21] INFO (nni.compression.pytorch.compressor/MainThread) Mask dict saved to ./mask.pth\n"
     ]
    }
   ],
   "source": [
    "# Defining the config list.\n",
    "# Note that here we only prune the depthwise convolution and the last pointwise convolution. \n",
    "# We will let the model speedup tool propagate the sparsity to the first pointwise convolution layer. \n",
    "\n",
    "pruner_name = 'l1'\n",
    "sparsity = 0.5\n",
    "\n",
    "if pruner_name != 'slim':\n",
    "    config_list = [{\n",
    "        'op_names': ['features.{}.conv.1.0'.format(x) for x in range(2, 18)],\n",
    "        'sparsity': sparsity\n",
    "    },{\n",
    "        'op_names': ['features.{}.conv.2'.format(x) for x in range(2, 18)],\n",
    "        'sparsity': sparsity\n",
    "    }]\n",
    "else:\n",
    "    # For slim pruner, we should specify BatchNorm layers instead of the corresponding Conv2d layers\n",
    "    config_list = [{\n",
    "        'op_names': ['features.{}.conv.1.1'.format(x) for x in range(2, 18)],\n",
    "        'sparsity': sparsity\n",
    "    },{\n",
    "        'op_names': ['features.{}.conv.3'.format(x) for x in range(2, 18)],\n",
    "        'sparsity': sparsity\n",
    "    }]\n",
    "\n",
    "# Different pruners require different additional parameters, so we put them together in the kwargs dict. \n",
    "# Please check the docs for detailed information.\n",
    "kwargs = {}                                                                                                                                \n",
    "if pruner_name in ['slim', 'taylor', 'activationmeanrank', 'apoz', 'agp']:\n",
    "    from pruning_experiments import trainer_helper\n",
    "    train_dataloader = get_dataloader('train', './data/stanford-dogs/Processed/train', batch_size=32)\n",
    "    def trainer(model, optimizer, criterion, epoch):\n",
    "        return trainer_helper(model, criterion, optimizer, train_dataloader, device)\n",
    "    kwargs = {\n",
    "        'trainer': trainer,\n",
    "        'optimizer': torch.optim.Adam(model.parameters()),\n",
    "        'criterion': nn.CrossEntropyLoss()\n",
    "    }\n",
    "    if pruner_name == 'agp':\n",
    "        kwargs['pruning_algorithm'] = 'l1'\n",
    "        kwargs['num_iterations'] = 10\n",
    "        kwargs['epochs_per_iteration'] = 1\n",
    "    if pruner_name == 'slim':\n",
    "        kwargs['sparsifying_training_epochs'] = 10\n",
    "\n",
    "# Create pruner, call pruner.compress(), and export the pruned model\n",
    "pruner = pruner_name_to_class[pruner_name](model, config_list, **kwargs)\n",
    "pruner.compress()\n",
    "pruner.export_model('./pruned_model.pth', './mask.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Speedup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: must unwrap the model before speed up\n",
    "pruner._unwrap_model()\n",
    "\n",
    "dummy_input = torch.rand(1,3,224,224).to(device)\n",
    "ms = ModelSpeedup(model, dummy_input, './mask.pth')\n",
    "ms.speedup_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "| Index |         Name         |  Type  |    Weight Shape   |  FLOPs   | #Params |\n",
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "|   0   |     features.0.0     | Conv2d |   (32, 3, 3, 3)   | 10838016 |   864   |\n",
      "|   1   | features.1.conv.0.0  | Conv2d |   (32, 1, 3, 3)   | 3612672  |   288   |\n",
      "|   2   |  features.1.conv.1   | Conv2d |   (16, 32, 1, 1)  | 6422528  |   512   |\n",
      "|   3   | features.2.conv.0.0  | Conv2d |   (48, 16, 1, 1)  | 9633792  |   768   |\n",
      "|   4   | features.2.conv.1.0  | Conv2d |   (48, 1, 3, 3)   | 1354752  |   432   |\n",
      "|   5   |  features.2.conv.2   | Conv2d |   (16, 48, 1, 1)  | 2408448  |   768   |\n",
      "|   6   | features.3.conv.0.0  | Conv2d |   (72, 16, 1, 1)  | 3612672  |   1152  |\n",
      "|   7   | features.3.conv.1.0  | Conv2d |   (72, 1, 3, 3)   | 2032128  |   648   |\n",
      "|   8   |  features.3.conv.2   | Conv2d |   (16, 72, 1, 1)  | 3612672  |   1152  |\n",
      "|   9   | features.4.conv.0.0  | Conv2d |   (72, 16, 1, 1)  | 3612672  |   1152  |\n",
      "|   10  | features.4.conv.1.0  | Conv2d |   (72, 1, 3, 3)   |  508032  |   648   |\n",
      "|   11  |  features.4.conv.2   | Conv2d |   (25, 72, 1, 1)  | 1411200  |   1800  |\n",
      "|   12  | features.5.conv.0.0  | Conv2d |   (96, 25, 1, 1)  | 1881600  |   2400  |\n",
      "|   13  | features.5.conv.1.0  | Conv2d |   (96, 1, 3, 3)   |  677376  |   864   |\n",
      "|   14  |  features.5.conv.2   | Conv2d |   (25, 96, 1, 1)  | 1881600  |   2400  |\n",
      "|   15  | features.6.conv.0.0  | Conv2d |   (96, 25, 1, 1)  | 1881600  |   2400  |\n",
      "|   16  | features.6.conv.1.0  | Conv2d |   (96, 1, 3, 3)   |  677376  |   864   |\n",
      "|   17  |  features.6.conv.2   | Conv2d |   (25, 96, 1, 1)  | 1881600  |   2400  |\n",
      "|   18  | features.7.conv.0.0  | Conv2d |   (96, 25, 1, 1)  | 1881600  |   2400  |\n",
      "|   19  | features.7.conv.1.0  | Conv2d |   (96, 1, 3, 3)   |  169344  |   864   |\n",
      "|   20  |  features.7.conv.2   | Conv2d |   (59, 96, 1, 1)  | 1110144  |   5664  |\n",
      "|   21  | features.8.conv.0.0  | Conv2d |  (192, 59, 1, 1)  | 2220288  |  11328  |\n",
      "|   22  | features.8.conv.1.0  | Conv2d |   (192, 1, 3, 3)  |  338688  |   1728  |\n",
      "|   23  |  features.8.conv.2   | Conv2d |  (59, 192, 1, 1)  | 2220288  |  11328  |\n",
      "|   24  | features.9.conv.0.0  | Conv2d |  (192, 59, 1, 1)  | 2220288  |  11328  |\n",
      "|   25  | features.9.conv.1.0  | Conv2d |   (192, 1, 3, 3)  |  338688  |   1728  |\n",
      "|   26  |  features.9.conv.2   | Conv2d |  (59, 192, 1, 1)  | 2220288  |  11328  |\n",
      "|   27  | features.10.conv.0.0 | Conv2d |  (192, 59, 1, 1)  | 2220288  |  11328  |\n",
      "|   28  | features.10.conv.1.0 | Conv2d |   (192, 1, 3, 3)  |  338688  |   1728  |\n",
      "|   29  |  features.10.conv.2  | Conv2d |  (59, 192, 1, 1)  | 2220288  |  11328  |\n",
      "|   30  | features.11.conv.0.0 | Conv2d |  (192, 59, 1, 1)  | 2220288  |  11328  |\n",
      "|   31  | features.11.conv.1.0 | Conv2d |   (192, 1, 3, 3)  |  338688  |   1728  |\n",
      "|   32  |  features.11.conv.2  | Conv2d |  (87, 192, 1, 1)  | 3273984  |  16704  |\n",
      "|   33  | features.12.conv.0.0 | Conv2d |  (288, 87, 1, 1)  | 4910976  |  25056  |\n",
      "|   34  | features.12.conv.1.0 | Conv2d |   (288, 1, 3, 3)  |  508032  |   2592  |\n",
      "|   35  |  features.12.conv.2  | Conv2d |  (87, 288, 1, 1)  | 4910976  |  25056  |\n",
      "|   36  | features.13.conv.0.0 | Conv2d |  (288, 87, 1, 1)  | 4910976  |  25056  |\n",
      "|   37  | features.13.conv.1.0 | Conv2d |   (288, 1, 3, 3)  |  508032  |   2592  |\n",
      "|   38  |  features.13.conv.2  | Conv2d |  (87, 288, 1, 1)  | 4910976  |  25056  |\n",
      "|   39  | features.14.conv.0.0 | Conv2d |  (288, 87, 1, 1)  | 4910976  |  25056  |\n",
      "|   40  | features.14.conv.1.0 | Conv2d |   (288, 1, 3, 3)  |  127008  |   2592  |\n",
      "|   41  |  features.14.conv.2  | Conv2d |  (134, 288, 1, 1) | 1891008  |  38592  |\n",
      "|   42  | features.15.conv.0.0 | Conv2d |  (480, 134, 1, 1) | 3151680  |  64320  |\n",
      "|   43  | features.15.conv.1.0 | Conv2d |   (480, 1, 3, 3)  |  211680  |   4320  |\n",
      "|   44  |  features.15.conv.2  | Conv2d |  (134, 480, 1, 1) | 3151680  |  64320  |\n",
      "|   45  | features.16.conv.0.0 | Conv2d |  (480, 134, 1, 1) | 3151680  |  64320  |\n",
      "|   46  | features.16.conv.1.0 | Conv2d |   (480, 1, 3, 3)  |  211680  |   4320  |\n",
      "|   47  |  features.16.conv.2  | Conv2d |  (134, 480, 1, 1) | 3151680  |  64320  |\n",
      "|   48  | features.17.conv.0.0 | Conv2d |  (480, 134, 1, 1) | 3151680  |  64320  |\n",
      "|   49  | features.17.conv.1.0 | Conv2d |   (480, 1, 3, 3)  |  211680  |   4320  |\n",
      "|   50  |  features.17.conv.2  | Conv2d |  (160, 480, 1, 1) | 3763200  |  76800  |\n",
      "|   51  |    features.18.0     | Conv2d | (1280, 160, 1, 1) | 10035200 |  204800 |\n",
      "|   52  |     classifier.1     | Linear |    (120, 1280)    |  153600  |  153720 |\n",
      "+-------+----------------------+--------+-------------------+----------+---------+\n",
      "FLOPs total: 139206976\n",
      "#Params total: 1074880\n",
      "MobileNetV2(\n",
      "  (features): Sequential(\n",
      "    (0): ConvBNActivation(\n",
      "      (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU6(inplace=True)\n",
      "    )\n",
      "    (1): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\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): ReLU6(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",
      "    (2): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(16, 48, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(48, 48, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=48, bias=False)\n",
      "          (1): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(48, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (3): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(16, 72, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(72, 72, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=72, bias=False)\n",
      "          (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(72, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (4): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(16, 72, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(72, 72, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=72, bias=False)\n",
      "          (1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(72, 25, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(25, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (5): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(25, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(96, 25, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(25, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (6): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(25, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(96, 25, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(25, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (7): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(25, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(96, 96, kernel_size=(3, 3), stride=(2, 2), 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(96, 59, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(59, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (8): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(59, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(192, 59, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(59, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (9): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(59, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(192, 59, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(59, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (10): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(59, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(192, 59, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(59, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (11): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(59, 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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\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): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(192, 87, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(87, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (12): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(87, 288, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(288, 288, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=288, bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(288, 87, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(87, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (13): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(87, 288, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(288, 288, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=288, bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(288, 87, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(87, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (14): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(87, 288, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(288, 288, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=288, bias=False)\n",
      "          (1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(288, 134, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(134, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (15): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(134, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(480, 134, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(134, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (16): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(134, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(480, 134, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(134, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (17): InvertedResidual(\n",
      "      (conv): Sequential(\n",
      "        (0): ConvBNActivation(\n",
      "          (0): Conv2d(134, 480, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (1): ConvBNActivation(\n",
      "          (0): Conv2d(480, 480, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=480, bias=False)\n",
      "          (1): BatchNorm2d(480, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          (2): ReLU6(inplace=True)\n",
      "        )\n",
      "        (2): Conv2d(480, 160, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "        (3): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (18): ConvBNActivation(\n",
      "      (0): Conv2d(160, 1280, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      (1): BatchNorm2d(1280, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU6(inplace=True)\n",
      "    )\n",
      "  )\n",
      "  (classifier): Sequential(\n",
      "    (0): Dropout(p=0.2, inplace=False)\n",
      "    (1): Linear(in_features=1280, out_features=120, bias=True)\n",
      "  )\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ")\n",
      "FLOPs: 139206976, params: 1074880\n"
     ]
    }
   ],
   "source": [
    "flops, params, results = count_flops_params(model, dummy_input)\n",
    "print(model)\n",
    "print(f\"FLOPs: {flops}, params: {params}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fine-tuning after Pruning\n",
    "\n",
    "Usually, after pruning out some weights from the model, we need further fine-tuning to let the model recover its performance as much as possible. For finetuning, we can either use the same setting during pretraining, or use an additional technique called [**Knowledge Distillation**](https://arxiv.org/pdf/1503.02531.pdf). The key idea is that the model learns on both the original hard labels and the soft labels produced by a teacher model running on the same input. In our setting, **the model before pruning can conveniently serve as the teacher model**. Empirically, we found that using distillation during fine-tuning consistently improves the performance of the pruned model. We will further discuss related experiments in the following section.\n",
    "\n",
    "Note that knowledge distillation can easily be done with the following lines of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sample code: training with knowledge distillation\n",
    "\"\"\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def train_with_distillation(student_model, teacher_model, optimizer, train_dataloader, device, alpha=0.99, temperature=8):\n",
    "    student_model.train()\n",
    "    for i, (inputs, labels) in enumerate(tqdm(train_dataloader)):\n",
    "        optimizer.zero_grad()\n",
    "        inputs, labels = inputs.float().to(device), labels.to(device)\n",
    "        with torch.no_grad():\n",
    "            teacher_preds = teacher_model(inputs)\n",
    "\n",
    "        student_preds = student_model(inputs)\n",
    "        soft_loss = nn.KLDivLoss()(F.log_softmax(student_preds/temperature, dim=1),\n",
    "                                   F.softmax(teacher_preds/temperature, dim=1))\n",
    "        hard_loss = F.cross_entropy(student_preds, labels)\n",
    "        loss = soft_loss * (alpha * temperature * temperature) + hard_loss * (1. - alpha)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finetuning after pruning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/v-diwu4/.cache/torch/hub/pytorch_vision_v0.8.1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start finetuning with distillation epoch 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/338 [00:00<?, ?it/s]/home/v-diwu4/miniconda3/envs/seg/lib/python3.7/site-packages/torch/nn/functional.py:2611: UserWarning: reduction: 'mean' divides the total loss by both the batch size and the support size.'batchmean' divides only by the batch size, and aligns with the KL div math definition.'mean' will be changed to behave the same as 'batchmean' in the next major release.\n",
      "  \"reduction: 'mean' divides the total loss by both the batch size and the support size.\"\n",
      "100%|██████████| 338/338 [00:42<00:00,  7.92it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 22.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0: train loss 0.0597, valid loss 2.6602, valid acc 0.2673\n",
      "Start finetuning with distillation epoch 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:41<00:00,  8.20it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 22.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: train loss 0.0394, valid loss 2.0302, valid acc 0.4400\n",
      "Start finetuning with distillation epoch 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:43<00:00,  7.85it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 23.37it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2: train loss 0.0309, valid loss 1.6968, valid acc 0.5238\n",
      "Start finetuning with distillation epoch 3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:43<00:00,  7.85it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 25.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3: train loss 0.0260, valid loss 1.5152, valid acc 0.5609\n",
      "Start finetuning with distillation epoch 4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:42<00:00,  8.00it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 25.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4: train loss 0.0231, valid loss 1.3936, valid acc 0.6053\n",
      "Start finetuning with distillation epoch 5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:41<00:00,  8.06it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 25.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5: train loss 0.0208, valid loss 1.3012, valid acc 0.6168\n",
      "Start finetuning with distillation epoch 6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:43<00:00,  7.77it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 22.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6: train loss 0.0191, valid loss 1.2575, valid acc 0.6242\n",
      "Start finetuning with distillation epoch 7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:44<00:00,  7.62it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 22.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7: train loss 0.0178, valid loss 1.1819, valid acc 0.6456\n",
      "Start finetuning with distillation epoch 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:43<00:00,  7.76it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 25.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8: train loss 0.0167, valid loss 1.1424, valid acc 0.6620\n",
      "Start finetuning with distillation epoch 9\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 338/338 [00:42<00:00,  7.93it/s]\n",
      "100%|██████████| 38/38 [00:01<00:00, 24.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: train loss 0.0158, valid loss 1.1762, valid acc 0.6464\n",
      "Best validation accuracy: 0.6620065789473685\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 269/269 [00:12<00:00, 21.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 1.2283264702137517\n",
      "Test accuracy: 0.6349907063197026\n"
     ]
    }
   ],
   "source": [
    "from pruning_experiments import run_finetune, run_finetune_distillation, run_eval\n",
    "\n",
    "use_distillation = True\n",
    "n_epochs = 10                  # set for demo purposes; increase this number for your experiments\n",
    "learning_rate = 1e-4\n",
    "weight_decay = 0.0\n",
    "\n",
    "train_dataloader = get_dataloader('train', './data/stanford-dogs/Processed/train', batch_size=32)\n",
    "valid_dataloader = get_dataloader('eval', './data/stanford-dogs/Processed/valid', batch_size=32)\n",
    "test_dataloader = get_dataloader('eval', './data/stanford-dogs/Processed/test', batch_size=32)\n",
    "\n",
    "if not use_distillation:\n",
    "    run_finetune(model, train_dataloader, valid_dataloader, device, \n",
    "                 n_epochs=n_epochs, learning_rate=learning_rate, weight_decay=weight_decay)\n",
    "else:\n",
    "    alpha = 0.99\n",
    "    temperature = 8\n",
    "    # use model with the original checkpoint as the teacher\n",
    "    teacher_model = create_model(model_type=model_type, pretrained=pretrained, n_classes=n_classes,\n",
    "                                 input_size=input_size, checkpoint=checkpoint).to(device)\n",
    "    run_finetune_distillation(model, teacher_model, train_dataloader, valid_dataloader, device, \n",
    "                              alpha, temperature,\n",
    "                              n_epochs=n_epochs, learning_rate=learning_rate, weight_decay=weight_decay)\n",
    "\n",
    "test_loss, test_acc = run_eval(model, test_dataloader, device)\n",
    "print('Test loss: {}\\nTest accuracy: {}'.format(test_loss, test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Steps of Optimizing Pruning Parameters\n",
    "So far, we have shown the end-to-end process of compressing a MobileNetV2 model on Stanford Dogs dataset using NNI pruners. It is crucial to mention that to make sure that the final model has a satisfactory performance, several trials on different sparsity values and pruner settings are necessary. To simplify this process for you, in this section we discuss how we approach the problem and mention some empirical observations. We hope that this section can serve as a good reference of the general process of optimizing the pruning with NNI Pruners.\n",
    "\n",
    "To help you reproduce some of the experiments, we implement `pruning_experiments.py`. Please find examples in the following code blocks for how to run experiments with this script."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: selecting the layer to prune\n",
    "Just as the first step of using the pruning is writing a `config_list`, the first thing you should consider when pruning a model is **which layer to prune**. This is crucial because some layers are not as sensitive to pruning as the others. In our example, we have several candidates for pruning:\n",
    "* the first pointwise convolution in all layers (the `conv 0.0`'s)\n",
    "* the depthwise convolution in all layers (the `conv 1.0`'s)\n",
    "* the second pointwise convolution in all layers (the `conv 2`'s)\n",
    "* some combination of the previous choices\n",
    "\n",
    "The following figure shows our experiment results. We run `L1FilterPruner` to explore some of the previous choices with layer sparsity ranging from 0.1 to 0.9. The x-axis shows the effective global sparsity after `ModelSpeedup`. We observe that jointly pruning the depthwise convolution and the second pointwise convolution often gives higher scores at large global sparsities. \n",
    "\n",
    "<div>\n",
    "<img src=\"step1.png\" width=\"600\"/>\n",
    "</div>\n",
    "\n",
    "Therefore, in the following experiments, we limit the modules to prune to the `conv 1.0`'s and the `conv 2`'s. Thus the config list is always written in the following way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [{\n",
    "    'op_names': ['features.{}.conv.1.0'.format(x) for x in range(2, 18)],\n",
    "    'sparsity': sparsity\n",
    "},{\n",
    "    'op_names': ['features.{}.conv.2'.format(x) for x in range(2, 18)],\n",
    "    'sparsity': sparsity\n",
    "}]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run some experiments for this step, please run `pruning_experiments.py` and specify the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Example shell script: \n",
    "\"\"\"\n",
    "for sparsity in 0.2 0.4 0.6 0.8; do\n",
    "    for pruning_mode in 'conv0' 'conv1' 'conv2' 'conv1andconv2' 'all'; do\n",
    "        python pruning_experiments.py \\\n",
    "            --experiment_dir pretrained_mobilenet_v2_torchhub/ \\\n",
    "            --checkpoint_name 'checkpoint_best.pt' \\\n",
    "            --sparsity $sparsity \\\n",
    "            --pruning_mode $pruning_mode \\\n",
    "            --pruner_name l1 \\\n",
    "            --speed_up \\\n",
    "            --finetune_epochs 30\n",
    "    done\n",
    "done\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: trying one-shot pruners\n",
    "After determining which modules to prune, we consider the next two questions:\n",
    "* **Which global sparsity range should we aim at?**\n",
    "* **Is there any one-shot pruning algorithm outperforming others at a large margin?**\n",
    "\n",
    "The first problem stems from the natural tradeoff between model size and accuracy. As long as we have acceptable performance, we wish the model to be as small as possible. Therefore, in this step, we can run some one-shot pruners with different sparsity settings, and find a range of sparsities that the model seem to maintain acceptable performance. \n",
    "\n",
    "The following figure summarizes our experiments on three pruners. We perform 30 epoch final finetuning for each experiment. Starting from the original model (with accuracy 0.8), we observe that when the sparsity is below 0.4, the pruned model can easily recover, with the performance approaching the model before pruning. On the other hand, when the sparsity is above 0.7, the model's performance drops too much even after finetuning. Therefore, we limit our search space to sparsity settings between 0.4 and 0.7 in the experiments for the following step 3 and step 4.\n",
    "\n",
    "<div>\n",
    "<img src=\"step2.png\" width=\"900\"/>\n",
    "</div>\n",
    "\n",
    "In addition, we observe that the slim pruner has better performance in the one-shot pruning setting. However, as we will show later, when we consider iterative pruning, the importance of choosing base pruning algorithms seem to be overwhelmed by choosing a correct pruning schedule.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run some experiments for this step, please run `pruning_experiments.py` and specify the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example shell script: \n",
    "\"\"\"\n",
    "for sparsity in 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9; do\n",
    "    for pruning_mode in 'conv1', 'conv1andconv2'; do\n",
    "        python pruning_experiments.py \\\n",
    "            --experiment_dir pretrained_mobilenet_v2_torchhub/ \\\n",
    "            --checkpoint_name 'checkpoint_best.pt' \\\n",
    "            --sparsity $sparsity \\\n",
    "            --pruning_mode $pruning_mode \\\n",
    "            --pruner_name l1 \\\n",
    "            --speed_up \\\n",
    "            --finetune_epochs 30\n",
    "    done\n",
    "done\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: determining iterative pruning strategy\n",
    "\n",
    "Now that we have found a good set of modules to prune and a good range of sparsity settings to experiment on, we can shift our focus to iterative pruning. Iterative pruning interleaves pruning with finetuning, and is often shown be more performant than one-shot pruning, which prunes the model once to the target sparsity. The following figure establishes that the superiority of iterative pruning under the same other settings.\n",
    "<img src=\"step3-1.png\" width=\"900\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we consider the following two important hyperparameters for iterative pruning:\n",
    "* the total number of pruning iterations\n",
    "* the number of finetuning epochs between pruning iterations\n",
    "\n",
    "We experiment we 2, 4, and 8 iterations, with 1 or 3 intermediate finetuning epochs. The results are summarized in the following figure. We clearly observe that increasing the number of pruning iterations significantly improves the final performance, while increasing the number of epochs only helps slightly. Therefore, we recommend that you should spend effort in **determining a correct (often large) number of pruning iterations**, while need not to spend a lot of effort tuning the number of finetuning epochs in between. In our case, we found iteration numbers between 64 and 128 gives the best performance. \n",
    "<img src=\"step3-2.png\" width=\"600\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run some experiments for this step, please run `pruning_experiments.py` and specify the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example shell script: \n",
    "\"\"\"\n",
    "for sparsity in 0.4 0.5 0.6 0.7; do\n",
    "    for n_iters in 2 4 8 16; do\n",
    "        python pruning_experiments.py \\\n",
    "            --experiment_dir pretrained_mobilenet_v2_torchhub/ \\\n",
    "            --checkpoint_name 'checkpoint_best.pt' \\\n",
    "            --sparsity $sparsity \\\n",
    "            --pruning_mode 'conv1andconv2' \\\n",
    "            --pruner_name 'agp' \\\n",
    "            --agp_n_iters $n_iters \\\n",
    "            --speed_up \\\n",
    "            --finetune_epochs 30 \\\n",
    "    done\n",
    "done\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: determining finetuning strategy\n",
    "Finally, after pruning the model, we recommend **using knowledge distillation for finetuning**, which only involves changing several lines of code computing the loss (if we reuse the model before pruning as the teacher model). As shown in the following figure, using knowledge distillation during finetuning can bring about 5 percentage performance improvement in our task. \n",
    "<div>\n",
    "<img src=\"step4.png\" width=\"600\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run some experiments for this step, please run `pruning_experiments.py` and specify the following arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example shell script: \n",
    "\"\"\"\n",
    "for sparsity in 0.4 0.5 0.6 0.7; do\n",
    "    python pruning_experiments.py \\\n",
    "        --experiment_dir pretrained_mobilenet_v2_torchhub/ \\\n",
    "        --checkpoint_name 'checkpoint_best.pt' \\\n",
    "        --sparsity $sparsity \\\n",
    "        --pruning_mode 'conv1andconv2' \\\n",
    "        --pruner_name 'agp' \\\n",
    "        --speed_up \\\n",
    "        --finetune_epochs 80\n",
    "done\n",
    "\n",
    "for sparsity in 0.4 0.5 0.6 0.7; do\n",
    "    python pruning_experiments.py \\\n",
    "        --experiment_dir pretrained_mobilenet_v2_torchhub/ \\\n",
    "        --checkpoint_name 'checkpoint_best.pt' \\\n",
    "        --sparsity $sparsity \\\n",
    "        --pruning_mode 'conv1andconv2' \\\n",
    "        --pruner_name 'agp' \\\n",
    "        --speed_up \\\n",
    "        --finetune_epochs 80 \\\n",
    "        -- kd\n",
    "done\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison with Baseline Methods\n",
    "To confirm that using NNI Pruners indeed results in a model with good performance. We implement and compare with the following baseline methods:\n",
    "1. Shrink the number of channel of all layers to half. This is a basic compression method mentioned by the MobileNet authors. We experiment with the following two settings:\n",
    "    * randomly initialize the weights and train with knowledge distillation\n",
    "    * use `L1FilterPruner` to prune the ImageNet pretrained model to 0.5 sparsity, and then train with knowledge distillation. \n",
    "2. Random pruning to 0.5 sparsity. \n",
    "\n",
    "In the first baseline, we observe that the randomly initialized model only has 0.45 test accuracy, while the model initialized with ImageNet weights has 0.7197 test accuracy after training. However, as shown in the table at the beginning of the notebook, using NNI pruners we can achieve 0.7703 test accuracy with the same amount of finetuning with knowledge distillation. This established the superiority of our approach. As a side remark, this observation is also consistent with the AGP authors' claim that \"large sparse\" models obtained by pruning often outperform \"small dense\" models with similar amount of parameters trained from scratch. \n",
    "\n",
    "In the second baseline, we observe that random pruning performs worse than our one-shot baselines, giving 0.7385 validation accuracy and 0.7182 test accuracy for 0.5 sparsity. This establishes that the pruning has its unique values that cannot be replaced by the final knowledge distillation process. \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this end-to-end example, we have shown the process of using NNI Pruners to compress MobileNetV2 on Stanford Dogs. With iterative pruning and knowledge distillation, we have pruned the MobileNetV2 architecture to 1/3 of its size, with 95% accuracy retained. In the last sections, we also introduce our approach to the problem, and wish that it could be a useful reference if you want to solve a similar problem with NNI Pruners. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp",
   "language": "python",
   "name": "nlp"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
