Model_DenseNet_batchnorm2d.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model DenseNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset "
   ]
  },
  {
   "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": 2,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 3, 32, 32])"
      ]
     },
     "execution_count": 3,
     "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. Basic Block & Transition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import torch.nn.functional as F\n",
    "\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",
    "\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 corresponds 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",
    "        :param 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",
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, in_planes, cfg, growth_rate=12, drop_rate=0):\n",
    "        \"\"\"\n",
    "            1 block has 1 convd, and 1 layer have many blocks\n",
    "            :param in_planes: input channel (last layer is the Conv output)\n",
    "            :param cfg: in fact, in_planes = cfg\n",
    "            :param growth_rate: output channel size = `in_planes` + growth_rate\n",
    "            :param drop_rate: dropout rate\n",
    "            \n",
    "            Note: currently we fix the expansion ratio as the default value (expansion=1)\n",
    "\n",
    "        \"\"\"\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.bn1 = nn.BatchNorm2d(in_planes)\n",
    "        self.select = channel_selection(in_planes)\n",
    "        self.conv1 = nn.Conv2d(cfg, growth_rate, kernel_size=3, padding=1, bias=False)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.drop_rate = drop_rate\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.bn1(x)\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "        if self.drop_rate > 0:\n",
    "            out = F.dropout(out, p=self.drop_rate, training=self.training)\n",
    "        out = torch.cat((x, out), 1)\n",
    "        return out\n",
    "\n",
    "    def forward_bn(self, x):\n",
    "        out = self.bn1(x)\n",
    "        bn_value = out.clone()\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "        if self.drop_rate > 0:\n",
    "            out = F.dropout(out, p=self.drop_rate, training=self.training)\n",
    "        out = torch.cat((x, out), 1)\n",
    "        return out, bn_value\n",
    "    \n",
    "    \n",
    "class Transition(nn.Module):\n",
    "    def __init__(self, in_planes, out_planes, cfg):\n",
    "        \"\"\"\n",
    "        :param in_planes: number of the input channel\n",
    "        :param out_planes: number of the output channel\n",
    "        :param cfg: \n",
    "        \"\"\"\n",
    "        super(Transition, self).__init__()\n",
    "        self.bn1 = nn.BatchNorm2d(in_planes)\n",
    "        self.select = channel_selection(in_planes)\n",
    "        self.conv1 = nn.Conv2d(cfg, out_planes, kernel_size=1, bias=False)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.bn1(x)\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "        out = F.avg_pool2d(out, 2)\n",
    "        return out\n",
    "    \n",
    "    def forward_bn(self, x):\n",
    "        out = self.bn1(x)\n",
    "        bn_value = out.clone()\n",
    "        out = self.select(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv1(out)\n",
    "        out = F.avg_pool2d(out, 2)\n",
    "        return out, bn_value\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Densnet "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "\n",
    "class densenet(nn.Module):\n",
    "    def __init__(self, depth=40, drop_rate=0, dataset='cifar10', growth_rate=12, compression_rate=1, cfg=None, block_cfg=None):\n",
    "        \"\"\"\n",
    "        :param depth: 3 (layers) × n (number of conv2ds <- in every layer) + 4 (Conv2ds)\n",
    "        :param drop_rate: dropout rate\n",
    "        :param dataset: cifar10 or cifar100\n",
    "        :param growth_rate: gradually increasing from the `n` conv2d to the `n+1` conv2d <- in every layer\n",
    "        :param cfg:\n",
    "            default cfg is None:\n",
    "                start = 24, len(cfg) = 3, growth_rate = 12, cfg.shape = [3, 12 + 1]\n",
    "                cfg[0] = [24(start), 36, 48, 60, 72, ..., 168]\n",
    "                cfg[1] = [168(start), 170, 182, 194, ..., 312]\n",
    "                cfg[2] = [312(start), 324, 336, 348, ..., 456]\n",
    "\n",
    "        model cfg:\n",
    "            1 conv2d -> (3, start, (3, 3))\n",
    "            1st layer (13 conv2ds)\n",
    "                input_channel = start (= growth_rate * 2)\n",
    "        \"\"\"\n",
    "        super(densenet, self).__init__()\n",
    "        assert (depth - 4) % 3 == 0, 'depth should be 3n+4'  # 3 layers\n",
    "        n = (depth - 4) // 3    # how many Conv2ds in every layer\n",
    "        block = BasicBlock\n",
    "\n",
    "        self.growth_rate = growth_rate\n",
    "        self.drop_rate = drop_rate\n",
    "\n",
    "        if cfg is None:\n",
    "            cfg = []\n",
    "            start = growth_rate * 2\n",
    "            for _ in range(3):\n",
    "                cfg.append([start + growth_rate * i for i in range(n + 1)])\n",
    "                start += growth_rate * n\n",
    "            cfg = [item for sub_list in cfg for item in sub_list]\n",
    "        assert len(cfg) == 3 * (n + 1), 'length of config variable cfg should be 3(n+1)'  # 39 = 3 × (12 + 1)\n",
    "\n",
    "        # Global variable used to across multiple\n",
    "        self.inplanes = growth_rate * 2\n",
    "        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, padding=1, bias=False)\n",
    "\n",
    "        self.dense1 = self._make_dense_block(block, num_block=n, cfg=cfg[0:n])\n",
    "        self.trans1 = self._make_transition(compression_rate, cfg[n])\n",
    "        self.dense2 = self._make_dense_block(block, n, cfg[n + 1:2 * n + 1])\n",
    "        self.trans2 = self._make_transition(compression_rate, cfg[2 * n + 1])\n",
    "        self.dense3 = self._make_dense_block(block, n, cfg[2 * n + 2:3 * n + 2])\n",
    "        self.bn = nn.BatchNorm2d(self.inplanes)\n",
    "        self.select = channel_selection(self.inplanes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.avgpool = nn.AvgPool2d(8)\n",
    "\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",
    "        # Weight initialization\n",
    "        self._initialize_weights()\n",
    "\n",
    "    def _make_dense_block(self, block, num_block, cfg):\n",
    "        \"\"\"\n",
    "        :param block: Basic Block (one block means one conv2d)\n",
    "        :param num_block: number of blocks (n) <- in every layer\n",
    "        :param cfg: all block's channels <- in every layer\n",
    "        \"\"\"\n",
    "        layers = []\n",
    "        assert num_block == len(cfg), 'Length of the cfg parameter is not right.'\n",
    "        for i in range(num_block):\n",
    "            layers.append(block(self.inplanes, cfg=cfg[i], growth_rate=self.growth_rate, drop_rate=self.drop_rate))\n",
    "            self.inplanes += self.growth_rate\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def _make_transition(self, compression_rate, cfg):\n",
    "        # cfg is a number in this case.\n",
    "        inplanes = self.inplanes\n",
    "        outplanes = int(math.floor(self.inplanes // compression_rate))\n",
    "        self.inplanes = outplanes\n",
    "        return Transition(inplanes, outplanes, cfg)\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",
    "        # 1 convd\n",
    "        x = self.conv1(x)\n",
    "\n",
    "        # 3 layers × n convds\n",
    "        x = self.trans1(self.dense1(x))\n",
    "        x = self.trans2(self.dense2(x))\n",
    "        x = self.dense3(x)\n",
    "        x = self.bn(x)\n",
    "        x = self.select(x)\n",
    "        x = self.relu(x)\n",
    "\n",
    "        # 1 pooling\n",
    "        x = self.avgpool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "densenet(\n",
       "  (conv1): Conv2d(3, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (dense1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (bn1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(24, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (bn1): BatchNorm2d(36, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(36, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (bn1): BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(48, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): BasicBlock(\n",
       "      (bn1): BatchNorm2d(60, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(60, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): BasicBlock(\n",
       "      (bn1): BatchNorm2d(72, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(72, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): BasicBlock(\n",
       "      (bn1): BatchNorm2d(84, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(84, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (6): BasicBlock(\n",
       "      (bn1): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(96, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (7): BasicBlock(\n",
       "      (bn1): BatchNorm2d(108, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(108, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (8): BasicBlock(\n",
       "      (bn1): BatchNorm2d(120, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(120, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (9): BasicBlock(\n",
       "      (bn1): BatchNorm2d(132, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(132, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (10): BasicBlock(\n",
       "      (bn1): BatchNorm2d(144, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(144, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (11): BasicBlock(\n",
       "      (bn1): BatchNorm2d(156, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(156, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (trans1): Transition(\n",
       "    (bn1): BatchNorm2d(168, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (select): channel_selection()\n",
       "    (conv1): Conv2d(168, 168, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (relu): ReLU(inplace=True)\n",
       "  )\n",
       "  (dense2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (bn1): BatchNorm2d(168, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(168, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (bn1): BatchNorm2d(180, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(180, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (bn1): BatchNorm2d(192, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(192, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): BasicBlock(\n",
       "      (bn1): BatchNorm2d(204, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(204, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): BasicBlock(\n",
       "      (bn1): BatchNorm2d(216, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(216, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): BasicBlock(\n",
       "      (bn1): BatchNorm2d(228, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(228, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (6): BasicBlock(\n",
       "      (bn1): BatchNorm2d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(240, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (7): BasicBlock(\n",
       "      (bn1): BatchNorm2d(252, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(252, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (8): BasicBlock(\n",
       "      (bn1): BatchNorm2d(264, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(264, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (9): BasicBlock(\n",
       "      (bn1): BatchNorm2d(276, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(276, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (10): BasicBlock(\n",
       "      (bn1): BatchNorm2d(288, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(288, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (11): BasicBlock(\n",
       "      (bn1): BatchNorm2d(300, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(300, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (trans2): Transition(\n",
       "    (bn1): BatchNorm2d(312, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (select): channel_selection()\n",
       "    (conv1): Conv2d(312, 312, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (relu): ReLU(inplace=True)\n",
       "  )\n",
       "  (dense3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (bn1): BatchNorm2d(312, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(312, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (bn1): BatchNorm2d(324, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(324, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (bn1): BatchNorm2d(336, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(336, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): BasicBlock(\n",
       "      (bn1): BatchNorm2d(348, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(348, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): BasicBlock(\n",
       "      (bn1): BatchNorm2d(360, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(360, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): BasicBlock(\n",
       "      (bn1): BatchNorm2d(372, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(372, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (6): BasicBlock(\n",
       "      (bn1): BatchNorm2d(384, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(384, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (7): BasicBlock(\n",
       "      (bn1): BatchNorm2d(396, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(396, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (8): BasicBlock(\n",
       "      (bn1): BatchNorm2d(408, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(408, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (9): BasicBlock(\n",
       "      (bn1): BatchNorm2d(420, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(420, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (10): BasicBlock(\n",
       "      (bn1): BatchNorm2d(432, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(432, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (11): BasicBlock(\n",
       "      (bn1): BatchNorm2d(444, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (select): channel_selection()\n",
       "      (conv1): Conv2d(444, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (bn): BatchNorm2d(456, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (select): channel_selection()\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (avgpool): AvgPool2d(kernel_size=8, stride=8, padding=0)\n",
       "  (classifier): Linear(in_features=456, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = densenet()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9360"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num = 0\n",
    "for k, m in enumerate(model.modules()):\n",
    "    if isinstance(m, nn.BatchNorm2d):\n",
    "        num += m.weight.data.shape[0]\n",
    "num"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9360"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = torch.randn(14, 3, 32, 32)\n",
    "\n",
    "num = 0\n",
    "for k, m in enumerate(model.children()):\n",
    "    if isinstance(m, nn.Sequential):\n",
    "        for l, n in enumerate(m.children()):\n",
    "            data, bn_value = n.forward_bn(data)\n",
    "            num += bn_value.shape[1]\n",
    "    elif isinstance(m, Transition):\n",
    "        data, bn_value = m.forward_bn(data)\n",
    "        num += bn_value.shape[1]\n",
    "    elif isinstance(m, nn.BatchNorm2d):\n",
    "        data = m(data)\n",
    "        num += m.weight.data.shape[0]\n",
    "    elif isinstance(m, nn.Linear):\n",
    "        continue\n",
    "    else:\n",
    "        data = m(data)\n",
    "num"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}