ScoreCountingScript_CIFAR100.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculating score for Micronet Challange"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remark: We don't have to account for Batch Norm parameters and FLOPS - https://tehnokv.com/posts/fusing-batchnorm-and-conv/ since they can be fused with the preceding convolutional layer, Dropout is set to 0% so we don't need parameters/FLOPS for that too."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`counting.py` script is taken from https://github.com/google-research/google-research/blob/master/micronet_challenge/counting.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import models\n",
    "from counting import Conv2D, FullyConnected, GlobalAvg, Add, count_ops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CIFAR100 submission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at WideResNet architecture which is a baseline architecture for calculating the score. It is also the architecture that we prune using our proposed method, hence the baseline model and our pruned model use the same architectures - the difference between them are sparsity levels per layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WideResNet(\n",
      "  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (layer1): Sequential(\n",
      "    (0): WideBasic(\n",
      "      (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(16, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential(\n",
      "        (0): Conv2d(16, 160, kernel_size=(1, 1), stride=(1, 1))\n",
      "      )\n",
      "    )\n",
      "    (1): WideBasic(\n",
      "      (bn1): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): WideBasic(\n",
      "      (bn1): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (3): WideBasic(\n",
      "      (bn1): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(160, 160, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): WideBasic(\n",
      "      (bn1): BatchNorm2d(160, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(160, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(320, 320, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (shortcut): Sequential(\n",
      "        (0): Conv2d(160, 320, kernel_size=(1, 1), stride=(2, 2))\n",
      "      )\n",
      "    )\n",
      "    (1): WideBasic(\n",
      "      (bn1): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): WideBasic(\n",
      "      (bn1): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (3): WideBasic(\n",
      "      (bn1): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(320, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): WideBasic(\n",
      "      (bn1): BatchNorm2d(320, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(320, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(640, 640, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "      (shortcut): Sequential(\n",
      "        (0): Conv2d(320, 640, kernel_size=(1, 1), stride=(2, 2))\n",
      "      )\n",
      "    )\n",
      "    (1): WideBasic(\n",
      "      (bn1): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (2): WideBasic(\n",
      "      (bn1): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "    (3): WideBasic(\n",
      "      (bn1): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv1): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (dropout): Dropout(p=0.0)\n",
      "      (bn2): BatchNorm2d(640, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (conv2): Conv2d(640, 640, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (shortcut): Sequential()\n",
      "    )\n",
      "  )\n",
      "  (bn1): BatchNorm2d(640, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)\n",
      "  (linear): Linear(in_features=640, out_features=100, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "num_classes = 100\n",
    "model_cfg = getattr(models, 'WideResNet28x10')\n",
    "model = model_cfg.base(*model_cfg.args, num_classes=num_classes, **model_cfg.kwargs)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ADD_BIT_BASE = 32\n",
    "MUL_BIT_BASE = 32\n",
    "BASELINE_PARAMETER_BITS = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_counts(total_params, total_mults, total_adds, mul_bits, add_bits):\n",
    "    # converting to Mbytes.\n",
    "    total_params = int(total_params) / 8. / 1e6\n",
    "    total_mults = total_mults * mul_bits / MUL_BIT_BASE / 1e6\n",
    "    total_adds = total_adds * add_bits / ADD_BIT_BASE  / 1e6\n",
    "    return total_params, total_mults, total_adds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get all the operations in the WideResNet28x10 architecture - as stated earlier in the remark we can ignore BatchNorm and Dropout layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_image_size = 32\n",
    "#kernel_shape = k_size, _, c_in, c_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_ops = []\n",
    "input_size = start_image_size\n",
    "index = 0\n",
    "input_sizes = [32]*13 + [16]*9 + [8]*7\n",
    "residuals = [3, 5, 7, 9, 12, 14, 16, 18, 21, 23, 25, 27]\n",
    "mask_for_sparsities = []\n",
    "for child in list(model.modules()):\n",
    "    if isinstance(child, nn.Conv2d):\n",
    "        input_size = input_sizes[index]\n",
    "        op = Conv2D(input_size=input_size, \n",
    "                    kernel_shape=child.kernel_size + (child.in_channels, child.out_channels), \n",
    "                    strides=child.stride,\n",
    "                    padding='same' if child.padding==(1,1) else 'valid',\n",
    "                    use_bias=True,\n",
    "                    activation=None if (index-3)%9==0 else 'relu')\n",
    "        model_ops.append(op)\n",
    "        mask_for_sparsities.append(1)\n",
    "        if index in residuals:\n",
    "            op = Add(input_size=input_sizes[index+1],\n",
    "                     n_channels=child.out_channels)\n",
    "            model_ops.append(op)\n",
    "            mask_for_sparsities.append(0)\n",
    "        index += 1\n",
    "    if isinstance(child, nn.Linear):\n",
    "        op = GlobalAvg(input_size=input_sizes[-1],\n",
    "                      n_channels=child.in_features)\n",
    "        model_ops.append(op)\n",
    "        mask_for_sparsities.append(0)\n",
    "        op = FullyConnected(kernel_shape=(child.in_features, child.out_features),\n",
    "                           use_bias=True,\n",
    "                           activation=None)\n",
    "        model_ops.append(op)\n",
    "        mask_for_sparsities.append(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's calculate the sparsities for each layer. Our method prunes the architecture on per-layer basis, hence the sparsity levels across layers differ. We can calculate the sparsity for each layer using `weigth_mask` (bit mask for the model weight). The percentage of 0s is equivalent to the sparsity level of a given layer. (Note: the order of layers in the `weight_mask` is the same as the order of layers in the `model_ops` list)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"weightsmasks_wideresnet.bin\"\n",
    "weight_mask = torch.load(path, map_location='cpu')\n",
    "total_weights_per_layer = [l.numel() for l in weight_mask]\n",
    "zeros_per_layer = [torch.sum(l==0).item() for l in weight_mask]\n",
    "sparsity_per_layer = [zeros/total for (zeros, total) in zip(zeros_per_layer, total_weights_per_layer)]\n",
    "\n",
    "# Pad with zeros for Add and Pool layers\n",
    "sparsity_per_layer_temp = sparsity_per_layer\n",
    "sparsity_per_layer = []\n",
    "index = 0\n",
    "for i in mask_for_sparsities:\n",
    "    if i == 0:\n",
    "        sparsity_per_layer.append(0)\n",
    "    else:\n",
    "        sparsity_per_layer.append(sparsity_per_layer_temp[index])\n",
    "        index += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Counts for the baseline WideResNet28x10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_params_baseline, total_mults_baseline, total_adds_baseline = [0] * 3\n",
    "for op in model_ops:\n",
    "    param_count, flop_mults, flop_adds = count_ops(op=op, sparsity=0, param_bits=BASELINE_PARAMETER_BITS)\n",
    "    total_params_baseline += param_count\n",
    "    total_mults_baseline += flop_mults\n",
    "    total_adds_baseline += flop_adds\n",
    "\n",
    "total_params_baseline, total_mults_baseline, total_adds_baseline = process_counts(\n",
    "                                             total_params = total_params_baseline,\n",
    "                                             total_mults = total_mults_baseline,\n",
    "                                             total_adds = total_adds_baseline, \n",
    "                                             mul_bits = MUL_BIT_BASE,\n",
    "                                             add_bits = ADD_BIT_BASE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11906.21696"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_mults_baseline + total_adds_baseline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Counts for our pruned model of WideResNet28x10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For our pruned model we are using the 'freebie' 16-bit quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMETER_BITS = 16\n",
    "ADD_BITS = 32\n",
    "MULT_BITS = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_params, total_mults, total_adds = [0] * 3\n",
    "for op, sparsity in zip(model_ops, sparsity_per_layer):\n",
    "    param_count, flop_mults, flop_adds = count_ops(op, sparsity, PARAMETER_BITS)\n",
    "    total_params += param_count\n",
    "    total_mults += flop_mults\n",
    "    total_adds += flop_adds\n",
    "    \n",
    "total_params, total_mults, total_adds = process_counts(total_params = total_params,\n",
    "                                                           total_mults = total_mults,\n",
    "                                                           total_adds = total_adds, \n",
    "                                                           mul_bits = MULT_BITS,\n",
    "                                                           add_bits = ADD_BITS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameter storage score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Our parameter score is: 0.06181214293465461.\n"
     ]
    }
   ],
   "source": [
    "param_score = total_params/total_params_baseline\n",
    "print('Our parameter score is: {}.'.format(param_score))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Math operations score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Our math operations score is: 0.04995274547187799.\n"
     ]
    }
   ],
   "source": [
    "total_flops_baseline = total_adds_baseline + total_mults_baseline\n",
    "# We overwrite our baseline number calculations due to discrepencies with counts provided by organizers\n",
    "total_flops_baseline = 10490 #(in millions)\n",
    "total_flops = total_adds + total_mults\n",
    "flops_score = total_flops/total_flops_baseline\n",
    "print('Our math operations score is: {}.'.format(flops_score))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Total score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Our total score is: 0.11176488840653259\n"
     ]
    }
   ],
   "source": [
    "total_score = param_score + flops_score\n",
    "print('Our total score is: {}'.format(total_score))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}