Model_ResNet_modules_params.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Resnet: modules & parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if torch.cuda.is_available() else {}\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "            datasets.CIFAR10('./data.cifar10', train=False, download=True, transform=transforms.Compose([\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])),\n",
    "            batch_size=8, shuffle=True, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 3, 32, 32])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "idx, data = next(enumerate(test_loader))\n",
    "one_batch = data[0]\n",
    "one_batch.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Channel Selection & Bottleneck"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class channel_selection(nn.Module):\n",
    "    \"\"\"\n",
    "    Select channels from the output of BatchNorm2d layer. It should be put directly after BatchNorm2d layer.\n",
    "    The output shape of this layer is determined by the number of 1 in `self.indexes`.\n",
    "    \"\"\"\n",
    "    def __init__(self, num_channels):\n",
    "        \"\"\"\n",
    "        Initialize the `indexes` with all one vector with the length same as the number of channels.\n",
    "        During pruning, the places in `indexes` which correpond to the channels to be pruned will be set to 0.\n",
    "        \"\"\"\n",
    "        super(channel_selection, self).__init__()\n",
    "        self.indexes = nn.Parameter(torch.ones(num_channels))\n",
    "\n",
    "    def forward(self, input_tensor):\n",
    "        \"\"\"\n",
    "        Parameter\n",
    "        ---------\n",
    "        input_tensor: (N,C,H,W). It should be the output of BatchNorm2d layer.\n",
    "        \"\"\"\n",
    "        selected_index = np.squeeze(np.argwhere(self.indexes.data.cpu().numpy()))\n",
    "        if selected_index.size == 1:\n",
    "            selected_index = np.resize(selected_index, (1,)) \n",
    "        output = input_tensor[:, selected_index, :, :]\n",
    "        return output\n",
    "\n",
    "    \n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, cfg, stride=1, down_sample=None):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.bn1 = nn.BatchNorm2d(inplanes)\n",
    "        self.select = channel_selection(inplanes)\n",
    "        self.conv1 = nn.Conv2d(cfg[0], cfg[1], kernel_size=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(cfg[1])\n",
    "        self.conv2 = nn.Conv2d(cfg[1], cfg[2], kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(cfg[2])\n",
    "        self.conv3 = nn.Conv2d(cfg[2], planes * 4, kernel_size=1, bias=False)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.down_sample = down_sample\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "\n",
    "        # group1\n",
    "        out = self.bn1(x)\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "\n",
    "        # group2\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        # group3\n",
    "        out = self.bn3(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv3(out)\n",
    "\n",
    "        # down sample\n",
    "        if self.down_sample is not None:\n",
    "            residual = self.down_sample(x)\n",
    "\n",
    "        out += residual\n",
    "        return out\n",
    "\n",
    "    def mask(self, index, cfg_mask):\n",
    "        if index == 0:\n",
    "            self.bn1.weight.data.mul_(cfg_mask)\n",
    "            self.bn1.bias.data.mul_(cfg_mask)\n",
    "        elif index == 1:\n",
    "            self.bn2.weight.data.mul_(cfg_mask)\n",
    "            self.bn2.bias.data.mul_(cfg_mask)\n",
    "        elif index == 2:\n",
    "            self.bn3.weight.data.mul_(cfg_mask)\n",
    "            self.bn3.bias.data.mul_(cfg_mask)\n",
    "        else:\n",
    "            raise ValueError(\"Index is not including.\")\n",
    "\n",
    "    def forward_bn(self, x):\n",
    "        bn_value = []\n",
    "        residual = x\n",
    "\n",
    "        # group1\n",
    "        out = self.bn1(x)\n",
    "        bn_value.append(out.clone())\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "\n",
    "        # group2\n",
    "        out = self.bn2(out)\n",
    "        bn_value.append(out.clone())\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        # group3\n",
    "        out = self.bn3(out)\n",
    "        bn_value.append(out.clone())\n",
    "        out = self.relu(out)\n",
    "        out = self.conv3(out)\n",
    "\n",
    "        # down sample\n",
    "        if self.down_sample is not None:\n",
    "            residual = self.down_sample(x)\n",
    "\n",
    "        out += residual\n",
    "        return out, bn_value\n",
    "\n",
    "    def forward_ge(self, x):\n",
    "        bn_value, bn_weight = [], []\n",
    "        residual = x\n",
    "\n",
    "        # group1\n",
    "        out = self.bn1(x)\n",
    "        bn_value.append(out.clone())\n",
    "        bn_weight.append(self.bn1.weight.data)\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "\n",
    "        # group2\n",
    "        out = self.bn2(out)\n",
    "        bn_value.append(out.clone())\n",
    "        bn_weight.append(self.bn2.weight.data)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        # group3\n",
    "        out = self.bn3(out)\n",
    "        bn_value.append(out.clone())\n",
    "        bn_weight.append(self.bn3.weight.data)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv3(out)\n",
    "\n",
    "        # down sample\n",
    "        if self.down_sample is not None:\n",
    "            residual = self.down_sample(x)\n",
    "\n",
    "        out += residual\n",
    "        return out, bn_value, bn_weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. ResNet Model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "resnet(\n",
       "  (feature): Sequential(\n",
       "    (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (1): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (down_sample): Sequential(\n",
       "          (0): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (4): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (5): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (6): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (7): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (8): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (9): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (10): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (11): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (12): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (13): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (14): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (15): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (16): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (17): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(16, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (2): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (down_sample): Sequential(\n",
       "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (4): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (5): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (6): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (7): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (8): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (9): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (10): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (11): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (12): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (13): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (14): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (15): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (16): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (17): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (3): Sequential(\n",
       "      (0): Bottleneck(\n",
       "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "        (down_sample): Sequential(\n",
       "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        )\n",
       "      )\n",
       "      (1): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (2): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (3): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (4): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (5): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (6): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (7): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (8): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (9): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (10): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (11): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (12): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (13): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (14): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (15): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (16): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "      (17): Bottleneck(\n",
       "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (select): channel_selection()\n",
       "        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn3): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (relu): ReLU(inplace=True)\n",
       "      )\n",
       "    )\n",
       "    (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (5): channel_selection()\n",
       "    (6): ReLU(inplace=True)\n",
       "    (7): AvgPool2d(kernel_size=8, stride=8, padding=0)\n",
       "  )\n",
       "  (classifier): Linear(in_features=256, out_features=100, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import math\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "preactivation resnet with bottleneck design.\n",
    "\"\"\"\n",
    "\n",
    "class resnet(nn.Module):\n",
    "    def __init__(self, depth=164, dataset='cifar10', cfg=None, block_cfg=None):\n",
    "        \"\"\"\n",
    "        :param depth:\n",
    "            164 layers => 1 conv2d + 3 layers × 18 blocks (every layer)  × 3 conv2ds (every block)  + 1 avgPool2d\n",
    "            param n = (depth - 2) // 9:\n",
    "                n means how many blocks in every layer\n",
    "                9 = 3 layers × 3 conv2d (every block)\n",
    "        :param cfg:\n",
    "            if depth = 164, then len(cfg) = 164\n",
    "        :param block_cfg:\n",
    "            every layer block indexes: [6, 12, 18], [9, 18], [18]\n",
    "            every `conv_cfg` value should <= `n` (if depth = 164, n = 18)\n",
    "        \"\"\"\n",
    "        super(resnet, self).__init__()\n",
    "        assert (depth - 2) % 9 == 0, 'depth should be 9n+2'\n",
    "\n",
    "        # model value\n",
    "        n = (depth - 2) // 9  # depth = 164, n = 18\n",
    "        block = Bottleneck\n",
    "        self.block_cfg = block_cfg\n",
    "        self.inplanes = 16\n",
    "\n",
    "        # model config\n",
    "        if cfg is None:\n",
    "            cfg = [[16, 16, 16], [64, 16, 16] * (n - 1),\n",
    "                   [64, 32, 32], [128, 32, 32] * (n - 1),\n",
    "                   [128, 64, 64], [256, 64, 64] * (n - 1), [256]]\n",
    "            cfg = [item for sub_list in cfg for item in sub_list]\n",
    "\n",
    "        # model feature\n",
    "        conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1, bias=False)\n",
    "        layer1 = self._make_layer(block, 16, n, cfg=cfg[0:3 * n])                # 32 × 32\n",
    "        layer2 = self._make_layer(block, 32, n, cfg=cfg[3 * n:6 * n], stride=2)  # 16 × 16\n",
    "        layer3 = self._make_layer(block, 64, n, cfg=cfg[6 * n:9 * n], stride=2)  # 8 × 8\n",
    "        bn = nn.BatchNorm2d(64 * block.expansion)\n",
    "        select = channel_selection(64 * block.expansion)\n",
    "        relu = nn.ReLU(inplace=True)\n",
    "        avg_pool = nn.AvgPool2d(8)\n",
    "\n",
    "        feature = [conv1, layer1, layer2, layer3, bn, select, relu, avg_pool]\n",
    "        self.feature = nn.Sequential(*feature)\n",
    "\n",
    "        # model classifier\n",
    "        if dataset == 'cifar10':\n",
    "            num_classes = 10\n",
    "        elif dataset == 'cifar100':\n",
    "            num_classes = 100\n",
    "        else:\n",
    "            raise ValueError('Model `dataset` parameter is Error!')\n",
    "        self.classifier = nn.Linear(cfg[-1], num_classes)\n",
    "\n",
    "        # model initialize weight\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def _make_layer(self, block, planes, num_block, cfg, stride=1):\n",
    "        \"\"\"\n",
    "        :param block: Bottleneck item\n",
    "        :param planes: record the layer's output channel size\n",
    "        :param num_block: how many blocks in every layer\n",
    "        :param cfg:\n",
    "        \"\"\"\n",
    "        down_sample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            down_sample = nn.Sequential(nn.Conv2d(\n",
    "                self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False))\n",
    "\n",
    "        layers = [block(self.inplanes, planes, cfg[0:3], stride, down_sample)]\n",
    "        self.inplanes = planes * block.expansion\n",
    "        for i in range(1, num_block):\n",
    "            layers.append(block(self.inplanes, planes, cfg[3 * i: 3 * (i + 1)]))\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def _initialize_weights(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n",
    "                m.weight.data.normal_(0, math.sqrt(2. / n))\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                m.weight.data.fill_(0.5)\n",
    "                m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                m.weight.data.normal_(0, 0.01)\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # model feature\n",
    "        block_value = []\n",
    "        if self.block_cfg:\n",
    "            for k, m in enumerate(self.feature):\n",
    "                if isinstance(m, nn.Sequential):\n",
    "                    for j, block_item in enumerate(m):\n",
    "                        x = block_item(x)\n",
    "                        if j + 1 in self.block_cfg:\n",
    "                            block_value.append(x)\n",
    "                else:\n",
    "                    x = m(x)\n",
    "        else:\n",
    "            x = self.feature(x)\n",
    "\n",
    "        # model classifier\n",
    "        x = x.view(x.size(0), -1)\n",
    "        y = self.classifier(x)\n",
    "\n",
    "        if len(block_value):\n",
    "            return y, block_value\n",
    "        return y\n",
    "\n",
    "model_resnet = resnet(dataset='cifar100')\n",
    "model_resnet.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5 tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,\n",
      "        0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000]) <bound method Module.type of BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)>\n",
      "8 tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,\n",
      "        0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000]) <bound method Module.type of BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)>\n"
     ]
    }
   ],
   "source": [
    "for k, m in enumerate(model_resnet.modules()):\n",
    "    if k < 10 and isinstance(m, nn.BatchNorm2d):\n",
    "        print(k, m.weight.data, m.type)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(86, 0.7687)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint = torch.load(r'G:\\lizulin\\network-slimming\\logs\\sparsity_resnet164_cifar100_s_1e_4\\\\model_best.pth.tar')\n",
    "best_prec1 = checkpoint['best_prec1']\n",
    "epoch1 = checkpoint['epoch']\n",
    "if checkpoint['state_dict'] is not None:\n",
    "    model_resnet.load_state_dict(checkpoint['state_dict'])\n",
    "    \n",
    "epoch1, best_prec1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### enumerate()迭代访问model.modules()，返回type如下：\n",
    "- __整个模型的type__。即这里的resnet，其中`feature`和`classifier`仅为对层的重命名，不具有具体类型或不参与到模型遍历中。\n",
    "- __依次遍历模型__。对Sequential类访问一次，然后对Sequential内的模块再访问。\n",
    "- 此处列举出`model_resnet.modules()`，__粗体__表示在feature过程有效的模块：\n",
    "    1. resnet (model type)\n",
    "    2. __Conv2d (input_channel=3, output_channel=16, kernel_size=3)__\n",
    "    3. Bottleneck\n",
    "    4. Bottleneck的items：__BatchNorm2d__、(__channel_selection__仅三个一组中的第一个有)、__ReLU__、__Conv2d__\n",
    "    5. AvgPool\n",
    "    6. Linear\n",
    "- 总的来说，数据流动的四个type：*Conv2d*、*BatchNorm2d*、*channel_selection*、*ReLU*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, m in enumerate(model_resnet.modules()):\n",
    "    print(k, \"-->\", m.type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### enumerate()迭代访问Model.children()，返回type如下：\n",
    "- feature命名的Sequential\n",
    "- classifier命名的Linear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, m in enumerate(model_resnet.children()):\n",
    "    print(k, \"-->\", m.type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Model modules and parameters\n",
    "\n",
    "1. model.modules() 返回模型的所有`子层`和`子层的子层`\n",
    "2. model.named_modules() 以字符串的形式返回\n",
    "3. model.children() 返回模型的子层\n",
    "4. model.named_children() 以字符串的形式返回\n",
    "5. model.parameters() 返回模型的所有参数，使用list()迭代\n",
    "6. model.named_parameters() 以字符串的形式返回\n",
    "7. model.state_dict() 返回模型的字典形式"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1 model.modules() 返回模型的所有`nn.Module`的子层以及子层的子层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resnet.modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, m in enumerate(model_resnet.modules()):\n",
    "    print(k, m.type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.3 model.children() 返回模型的子层，常见的比如`feature`和`classifier`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resnet.children"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.5 model.parameters()  返回模型的所有参数，使用list()迭代"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(model_resnet.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.7 model.state_dict() 返回模型的字典形式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resnet.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Access Model feature\n",
    "\n",
    "#### 4.1 module -> feature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resnet.feature"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.2 module -> feature -> layer (18 bottlenecks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_resnet.feature[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.3 module -> feature -> layer -> bottleneck "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_resnet.feature[1][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.4 module -> feature -> layer -> bottleneck -> batchNorm2d "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, m in enumerate(model_resnet.feature[1].children()):\n",
    "    for j, n in enumerate(m.children()):\n",
    "        print(k, \"->\", n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.5 foreach: module -> feature -> layer -> bottleneck -> batchNorm2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "layer_idx = 0\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        layer_idx += 1\n",
    "        print(\"layer {}\".format(layer_idx))\n",
    "        block_idx = 0\n",
    "        for j, n in enumerate(m.children()):     # layers -> layer (many Bottlencks)\n",
    "            block_idx += 1\n",
    "            print(\"\\tblock {}\".format(block_idx))\n",
    "            bn_idx = 0\n",
    "            for i, o in enumerate(n.children()): # Bottlencks -> batchNorm2d\n",
    "                if isinstance(o, nn.BatchNorm2d):\n",
    "                    bn_idx += 1\n",
    "                    print(\"\\t\\tbn {}\".format(bn_idx))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Access Model forward\n",
    "\n",
    "#### 5.1 model block config\n",
    "\n",
    "- block_cfg = [6, 12, 18], block value: 3 × 3 = 9\n",
    "- block_cfg = [9, 18], block value: 3 × 2 = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "block_value = []\n",
    "block_cfg = [6, 12, 18]\n",
    "for k, m in enumerate(model_resnet.feature):\n",
    "    if isinstance(m, nn.Sequential):\n",
    "        for j, block_item in enumerate(m.children()):\n",
    "            if j+1 in block_cfg:\n",
    "                block_value.append(\"layer {} block {} value\".format(k, j+1))\n",
    "block_value, len(block_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.2 `forward_bn(x)`: model > bottlenck  > batchNorm2d => Feature map \n",
    "\n",
    "- 实现方式：`forward_bn(x)` retrun `data`, `bn_value`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "data = torch.randn(32, 3, 32, 32)\n",
    "\n",
    "layer_idx = 0\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        layer_idx += 1\n",
    "        print(\"layer {}\".format(layer_idx))\n",
    "        block_idx = 0\n",
    "        for j, n in enumerate(model_resnet.feature[k]):     # layers -> layer (many Bottlencks)\n",
    "            block_idx += 1\n",
    "            print(\"\\tblock {}\".format(block_idx))\n",
    "            bn_idx = 0\n",
    "            data, bn_value = n.forward_bn(data)\n",
    "            print(\"\\t\\tbn_value: \", list(bn_value[0].shape))\n",
    "    else:\n",
    "        data = m(data)\n",
    "        print(k, '->', list(data.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.3 `bn_value()`: model > bottlenck > batchNorm2d => Feature map\n",
    "\n",
    "- 实现方式：bn_value() return bn_data\n",
    "- Model访问方式：\n",
    "    - model_resnet.feature => k, m => model_resnet.feature[k]\n",
    "    - model_resnet.feature => k, m => m.children()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "data = torch.randn(32, 3, 32, 32)\n",
    "data1 = data.clone()\n",
    "\n",
    "record1 = []\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        print(\"layer {}\".format(k))\n",
    "        for j, n in enumerate(m.children()):     # layers -> layer (many Bottlencks)\n",
    "            print(\"\\tblock {}\".format(j))\n",
    "            \n",
    "            data = n(data)\n",
    "            bn_value = n.bn_value()\n",
    "            \n",
    "            record1.append(bn_value)\n",
    "            print(\"\\t\\tbn_value: \", list(bn_value[0].shape))\n",
    "\n",
    "    else:\n",
    "        data = m(data)\n",
    "        print(k, '->', data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.4 相同的输入data，验证两种方式的结果是相同的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "data = data1.clone()\n",
    "\n",
    "record2 = []\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        print(\"layer {}\".format(k))\n",
    "        for j, n in enumerate(m.children()):     # layers -> layer (many Bottlencks)\n",
    "            print(\"\\tblock {}\".format(j))\n",
    "            \n",
    "            data, bn_value = n.forward_bn(data)\n",
    "            record2.append(bn_value)\n",
    "            \n",
    "            print(\"\\t\\tbn_value: \", list(bn_value[0].shape))\n",
    "\n",
    "    else:\n",
    "        data = m(data)\n",
    "        print(k, '->', data.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "record1[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "record2[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.5 Bottlenck => BatchNorm2d\n",
    "\n",
    "    由于Bootlenck的forward后的`out`需要与输入的数据`x`残差连接，所以使用`data=m(data)`遍历Bootlenck会出现问题"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch \n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "\n",
    "data = torch.randn(32, 3, 32, 32)\n",
    "\n",
    "layer_idx = 0\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        layer_idx += 1\n",
    "        print(\"layer {}\".format(layer_idx))\n",
    "        block_idx = 0\n",
    "        for j, n in enumerate(model_resnet.feature[k].children()):     # layers -> layer (many Bottlencks)\n",
    "            block_idx += 1\n",
    "            print(\"\\tblock {}\".format(block_idx))\n",
    "            bn_idx = 0\n",
    "            for i, o in enumerate(n.children()): # Bottlencks -> batchNorm2d\n",
    "                if isinstance(o, nn.Sequential):\n",
    "                    print(o.type)\n",
    "                    for l, p in enumerate(o):\n",
    "                        print(p.type)\n",
    "                        data = p(data)\n",
    "                        print('\\t\\t-', k, '->', p.type)\n",
    "                else:\n",
    "                    data = o(data)\n",
    "                    print('\\t\\t', k, '->', o.type)\n",
    "                if isinstance(o, nn.BatchNorm2d):\n",
    "                    bn_idx += 1\n",
    "                    print(\"\\t\\tbn {}\".format(bn_idx))\n",
    "    else:\n",
    "        data = m(data)\n",
    "        print(k, '->', data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Access Model BatchNorm2d  Number: 163\n",
    "\n",
    "$$163 = 162 (3 layers × 18 Bottlenecks × 3 BatchNorm2ds) + 1 BatchNorm2d$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bn_count, bn_size = 0, 0\n",
    "for k, m in enumerate(model_resnet.modules()):\n",
    "    if isinstance(m, nn.BatchNorm2d):\n",
    "        bn_count += 1\n",
    "        bn_size += m.weight.data.shape[0]\n",
    "bn_count, bn_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.randn(32, 3, 32, 32)\n",
    "\n",
    "bn_count, bn_size = 0, 0\n",
    "for k, m in enumerate(model_resnet.feature):     # modules (-classifier) -> feature\n",
    "    if isinstance(m, nn.Sequential):             # feature (-conv2d) -> layers\n",
    "        for j, n in enumerate(m.children()):     # layers -> layer (many Bottlencks)\n",
    "            data = n(data)\n",
    "            bn_value = n.bn_value()\n",
    "            \n",
    "            bn_count += 3\n",
    "            for i, bn_item in enumerate(bn_value):\n",
    "                bn_size += bn_item.shape[1]\n",
    "    else:\n",
    "        data = m(data)\n",
    "        if isinstance(m, nn.BatchNorm2d):\n",
    "            bn_count += 1\n",
    "            bn_size += data.shape[1]\n",
    "bn_count, bn_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def gammas(feature_map):\n",
    "    \"\"\"\n",
    "    :param feature_map: [B, C, H, W]\n",
    "    :return: [C]\n",
    "    \"\"\"\n",
    "    b, c, h, w = feature_map.shape\n",
    "    gamma = torch.zeros(c)\n",
    "    for j in range(c):\n",
    "        feature_map_j = feature_map[:, torch.arange(c) != j, :, :]\n",
    "        gamma[j] = at_loss(feature_map, feature_map_j)\n",
    "    return gamma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b, c, h, w = 64, 3, 32, 32\n",
    "gamma = torch.randn(b, c, h, w)\n",
    "value = gamma[0:23, :, :, :]\n",
    "value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "value = gamma.view(0:23, -1, -1, -1)\n",
    "value.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:python37] *",
   "language": "python",
   "name": "conda-env-python37-py"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}