CNN_archs/cnn_archs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CNN Architectures and How to Use Them with fastai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Update__: \n",
    "\n",
    "* Added [an example](cnn_archs_more.ipynb) for how to use the [pytorchcv](https://github.com/osmr/imgclsmob/tree/master/pytorch) models. `pytorchcv` provides a much more comprehensive list of models.\n",
    "\n",
    "* [Torchvision models](https://pytorch.org/docs/stable/torchvision/index.html) are all included in fastai v1 in [this PR](https://github.com/fastai/fastai/pull/1523), except for \"inception_v3\" (check the \"Inception\" section to see why it's left out and a workaround). \n",
    "\n",
    "* Most of the [Cadene pretrained models](https://github.com/Cadene/pretrained-models.pytorch) are included in fastai v1 in [this PR](https://github.com/fastai/fastai/pull/1753). You can directly import the models, just like ResNets. See [this file](PR_test.ipynb) for some examples.\n",
    "\n",
    "* EfficientNet is added."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In an image classification project, I have tried various CNN architectures and compared their performances, with the help of the great deep learning library [fastai](https://docs.fast.ai/). Fastai didn't include a lot of architectures by itself, but its flexible API allows us to use pretrained models from [torchvision](https://pytorch.org/docs/stable/torchvision/index.html) or [Cadene](https://github.com/Cadene/pretrained-models.pytorch) for transfer learning, with a little bit work of customization. Since I have already done these customizations to make the models work, I'd like to share it in this notebook. \n",
    "\n",
    "We'll use the pretrained models from torchvision and Cadene:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import *\n",
    "import pretrainedmodels\n",
    "\n",
    "from fastai.vision import *\n",
    "from fastai.vision.models import *\n",
    "from fastai.vision.learner import model_meta\n",
    "\n",
    "from utils import *\n",
    "import sys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The fastai version used in here is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.53.dev0'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are all models provided by torchvision: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['alexnet',\n",
       " 'AlexNet',\n",
       " 'ResNet',\n",
       " 'resnet18',\n",
       " 'resnet34',\n",
       " 'resnet50',\n",
       " 'resnet101',\n",
       " 'resnet152',\n",
       " 'VGG',\n",
       " 'vgg11',\n",
       " 'vgg11_bn',\n",
       " 'vgg13',\n",
       " 'vgg13_bn',\n",
       " 'vgg16',\n",
       " 'vgg16_bn',\n",
       " 'vgg19_bn',\n",
       " 'vgg19',\n",
       " 'SqueezeNet',\n",
       " 'squeezenet1_0',\n",
       " 'squeezenet1_1',\n",
       " 'Inception3',\n",
       " 'inception_v3',\n",
       " 'DenseNet',\n",
       " 'densenet121',\n",
       " 'densenet169',\n",
       " 'densenet201',\n",
       " 'densenet161']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[k for k,v in sys.modules['torchvision.models'].__dict__.items() if callable(v)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are all models provided by Cadene: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['fbresnet152',\n",
       " 'bninception',\n",
       " 'resnext101_32x4d',\n",
       " 'resnext101_64x4d',\n",
       " 'inceptionv4',\n",
       " 'inceptionresnetv2',\n",
       " 'alexnet',\n",
       " 'densenet121',\n",
       " 'densenet169',\n",
       " 'densenet201',\n",
       " 'densenet161',\n",
       " 'resnet18',\n",
       " 'resnet34',\n",
       " 'resnet50',\n",
       " 'resnet101',\n",
       " 'resnet152',\n",
       " 'inceptionv3',\n",
       " 'squeezenet1_0',\n",
       " 'squeezenet1_1',\n",
       " 'vgg11',\n",
       " 'vgg11_bn',\n",
       " 'vgg13',\n",
       " 'vgg13_bn',\n",
       " 'vgg16',\n",
       " 'vgg16_bn',\n",
       " 'vgg19_bn',\n",
       " 'vgg19',\n",
       " 'nasnetamobile',\n",
       " 'nasnetalarge',\n",
       " 'dpn68',\n",
       " 'dpn68b',\n",
       " 'dpn92',\n",
       " 'dpn98',\n",
       " 'dpn131',\n",
       " 'dpn107',\n",
       " 'xception',\n",
       " 'senet154',\n",
       " 'se_resnet50',\n",
       " 'se_resnet101',\n",
       " 'se_resnet152',\n",
       " 'se_resnext50_32x4d',\n",
       " 'se_resnext101_32x4d',\n",
       " 'cafferesnet101',\n",
       " 'pnasnet5large',\n",
       " 'polynet']"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pretrainedmodels.model_names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## ResNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We first look at ResNet, which is proposed in \"[Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)\". In the torchvision implementation, the ResNet-34 model has the following structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 15  layers (total: 19)\n",
      "(5) Sequential  : 22  layers (total: 41)\n",
      "(6) Sequential  : 32  layers (total: 73)\n",
      "(7) Sequential  : 17  layers (total: 90)\n",
      "(8) AvgPool2d   : 1   layers (total: 91)\n",
      "(9) Linear      : 1   layers (total: 92)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnet34)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The `arch_summary` prints out the module name of the _direct_ children, and count how many submodules in each of them. E.g., the (0) contains only one module (itself), the (4) contains 15 submodules. The direct children gives us a big picture of the structure. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can see after the `Conv2d`, `BatchNorm2d`, `ReLU` and `MaxPool2d` layers, we have four `BasicBlock`. Finally, we have an `AvgPool2d` layer and `Linear` layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For transfer learning, we cut out layer (8) and (9) and add our own custom head depending on the problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "One of greatest techniques fastai uses heavily is called [discriminative layer training](https://docs.fast.ai/basic_train.html#Discriminative-layer-training), which is to use apply different learning rates to different layers. The idea is, what the first few layers learn is more task-independent, so we don't want to change them too much from the pretrained weights; we'd like to train more for the last few layers added by us. \n",
    "\n",
    "To group the layers for discriminative learning rates, we can set (0) ~ (5) as group 1, set (6) and (7) as group 2, and our own custom head as group 3. In fastai, the cut and split is defined as:\n",
    "```\n",
    "{'cut':-2, 'split':_resnet_split }\n",
    "```\n",
    "\n",
    "The \"cut\" part means we'll cut out the last two layers from the original model, then add a custom head. Now the new model has two parts: a \"body\" and a \"head\". The \"split\" part:\n",
    "\n",
    "```\n",
    "def _resnet_split(m:nn.Module): return (m[0][6],m[1])\n",
    "```\n",
    "\n",
    "It means we split the whole model at the (6) of the \"body\" part (`m[0][6]`), and split at the \"head\" part (`m[1]`), so finally we got three layer groups (\"body\" cut to two, plus the \"head\"), and we can apply three different learning rates during training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Since we're only testing the models and don't care about the data, so we can mock the data up by using `FakeData`. Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), resnet34, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "When you use it to train you'll want to set `pretrained=True` to use the pretrained weights, but in here since we're testing the model structures so we don't care about the weights. \n",
    "\n",
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', 'Sequential', 'Sequential']\n",
      "Group 2: ['Sequential', 'Sequential']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "It's split as we expected. \n",
    "\n",
    "There's quite a few versions of ResNet, we now take a look ResNet-50: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 23  layers (total: 27)\n",
      "(5) Sequential  : 30  layers (total: 57)\n",
      "(6) Sequential  : 44  layers (total: 101)\n",
      "(7) Sequential  : 23  layers (total: 124)\n",
      "(8) AvgPool2d   : 1   layers (total: 125)\n",
      "(9) Linear      : 1   layers (total: 126)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnet50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can see it has the same structure as ResNet-34, but the four `BasicBlock` are replaced by four `Bottleneck` blocks. Same for ResNet-101 and ResNet-152:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 23  layers (total: 27)\n",
      "(5) Sequential  : 30  layers (total: 57)\n",
      "(6) Sequential  : 163 layers (total: 220)\n",
      "(7) Sequential  : 23  layers (total: 243)\n",
      "(8) AvgPool2d   : 1   layers (total: 244)\n",
      "(9) Linear      : 1   layers (total: 245)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnet101)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 23  layers (total: 27)\n",
      "(5) Sequential  : 58  layers (total: 85)\n",
      "(6) Sequential  : 254 layers (total: 339)\n",
      "(7) Sequential  : 23  layers (total: 362)\n",
      "(8) AvgPool2d   : 1   layers (total: 363)\n",
      "(9) Linear      : 1   layers (total: 364)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnet152)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "This means we can use the same code in fastai for all of these architectures."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## ResNeXt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "ResNeXt is proposed in \"[Aggregated Residual Transformations for Deep Neural Networks](https://arxiv.org/abs/1611.05431)\". We first take a look at the Cadene's implementation of \"resnext101_32x4d\". Wrap it to adapt the PyTorch models' API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def resnext101_32x4d(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.resnext101_32x4d(pretrained=pretrained)\n",
    "    all_layers = list(model.children())\n",
    "    return nn.Sequential(*all_layers[0], *all_layers[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 34  layers (total: 38)\n",
      "(5) Sequential  : 45  layers (total: 83)\n",
      "(6) Sequential  : 254 layers (total: 337)\n",
      "(7) Sequential  : 34  layers (total: 371)\n",
      "(8) AvgPool2d   : 1   layers (total: 372)\n",
      "(9) Linear      : 1   layers (total: 373)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnext101_32x4d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "After the `Conv2d`, `BatchNorm2d`, `ReLU` and `MaxPool2d` layers, we have four blocks. Finally, we have an `AvgPool2d` layer and `Linear` layer as usual. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For transfer learning, we cut out layer (8) and (9) and add our own custom head depending on the problem.  \n",
    "\n",
    "To group the layers for discriminative learning rates, we can set (0) ~ (5) as group 1, set (6) and (7) as group 2, and our own custom head as group 3. In fastai, we can specify the cut and split by:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), resnext101_32x4d, pretrained=False,\n",
    "                    cut=-2, split_on=lambda m: (m[0][6], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', 'Sequential', 'Sequential']\n",
      "Group 2: ['Sequential', 'Sequential']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can see the original model's last two layers are removed, the rest layers are split into group 1 and 2. Group 3 is the custom head added by fastai. So the customization works as expected.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Alternatively, you can add the cut and split to the model metadata before calling `create_cnn`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_resnext_meta = {'cut': -2, 'split': lambda m: (m[0][6], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[resnext101_32x4d] = _resnext_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "\"resnext101_64x4d\" has similar structures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def resnext101_64x4d(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.resnext101_64x4d(pretrained=pretrained)\n",
    "    all_layers = list(model.children())\n",
    "    return nn.Sequential(*all_layers[0], *all_layers[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) Sequential  : 34  layers (total: 38)\n",
      "(5) Sequential  : 45  layers (total: 83)\n",
      "(6) Sequential  : 254 layers (total: 337)\n",
      "(7) Sequential  : 34  layers (total: 371)\n",
      "(8) AvgPool2d   : 1   layers (total: 372)\n",
      "(9) Linear      : 1   layers (total: 373)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(resnext101_64x4d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So we can use the same cut and split to create a learner. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## SENet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "SENet is proposed in \"[Squeeze-and-Excitation Networks](https://arxiv.org/pdf/1709.01507.pdf)\". We first take a look at the Cadene's implementation of \"se_resnet50\". Wrap it to adapt the PyTorch models' API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def se_resnet50(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.se_resnet50(pretrained=pretrained)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 4   layers (total: 4)\n",
      "(1) Sequential  : 38  layers (total: 42)\n",
      "(2) Sequential  : 50  layers (total: 92)\n",
      "(3) Sequential  : 74  layers (total: 166)\n",
      "(4) Sequential  : 38  layers (total: 204)\n",
      "(5) AvgPool2d   : 1   layers (total: 205)\n",
      "(6) Linear      : 1   layers (total: 206)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(se_resnet50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "It contains five `SEResNetBottleneck` blocks, followed by an `AvgPool2d` layer and `Linear` layer. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For transfer learning, we cut out the last two layers and add our own custom head depending on the problem.  \n",
    "\n",
    "To group the layers for discriminative learning rates, we can set (0) ~ (2) as group 1, set (3) and (4) as group 2, and our own custom head as group 3. In fastai, we can specify the cut and split by:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), se_resnet50, pretrained=False,\n",
    "                    cut=-2, split_on=lambda m: (m[0][3], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Sequential', 'Sequential', 'Sequential']\n",
      "Group 2: ['Sequential', 'Sequential']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Lambda', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Alternatively, you can add the cut and split to the model metadata before calling `create_cnn`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_se_resnet_meta = {'cut': -2, 'split': lambda m: (m[0][3], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[se_resnet50] = _se_resnet_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "What about other SENet? Similarly, they all have the same structure. The structure of \"se_resnet101\": "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 4   layers (total: 4)\n",
      "(1) Sequential  : 38  layers (total: 42)\n",
      "(2) Sequential  : 50  layers (total: 92)\n",
      "(3) Sequential  : 278 layers (total: 370)\n",
      "(4) Sequential  : 38  layers (total: 408)\n",
      "(5) AvgPool2d   : 1   layers (total: 409)\n",
      "(6) Linear      : 1   layers (total: 410)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: pretrainedmodels.se_resnet101(pretrained=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The structure of \"se_resnext50_32x4d\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 4   layers (total: 4)\n",
      "(1) Sequential  : 38  layers (total: 42)\n",
      "(2) Sequential  : 50  layers (total: 92)\n",
      "(3) Sequential  : 74  layers (total: 166)\n",
      "(4) Sequential  : 38  layers (total: 204)\n",
      "(5) AvgPool2d   : 1   layers (total: 205)\n",
      "(6) Linear      : 1   layers (total: 206)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: pretrainedmodels.se_resnext50_32x4d(pretrained=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So the same cut and split can be used for the above models. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "\"senet154\" has one more dropout layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 10  layers (total: 10)\n",
      "(1) Sequential  : 38  layers (total: 48)\n",
      "(2) Sequential  : 98  layers (total: 146)\n",
      "(3) Sequential  : 434 layers (total: 580)\n",
      "(4) Sequential  : 38  layers (total: 618)\n",
      "(5) AvgPool2d   : 1   layers (total: 619)\n",
      "(6) Dropout     : 1   layers (total: 620)\n",
      "(7) Linear      : 1   layers (total: 621)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: pretrainedmodels.senet154(pretrained=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can use the same split but set `cut=-3`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Densenet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Densenet is proposed in \"[Densely Connected Convolutional Networks](https://arxiv.org/pdf/1608.06993.pdf)\". \n",
    "\n",
    "We first look at the torchvision implementation of \"densenet121\". It has the following structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 365 layers (total: 365)\n",
      "(1) Linear      : 1   layers (total: 366)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(densenet121)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We expand and look into the first module: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) _DenseBlock : 36  layers (total: 40)\n",
      "(5) _Transition : 4   layers (total: 44)\n",
      "(6) _DenseBlock : 72  layers (total: 116)\n",
      "(7) _Transition : 4   layers (total: 120)\n",
      "(8) _DenseBlock : 144 layers (total: 264)\n",
      "(9) _Transition : 4   layers (total: 268)\n",
      "(10) _DenseBlock : 96  layers (total: 364)\n",
      "(11) BatchNorm2d : 1   layers (total: 365)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(densenet121(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "After the usual `Conv2d`, `BatchNorm2d`, `ReLU` and `MaxPool2d` layers, we got some `_DenseBlock` and `_Transition` blocks. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For transfer learning, we can cut out the last `Linear` layer and add our own custom head depending on the problem.  \n",
    "\n",
    "To group the layers for discriminative learning rates, we can set (0) ~ (6) as group 1, set (7) and (11) as group 2, and our own custom head as group 3. In fastai, we can specify the cut and split by:\n",
    "\n",
    "```\n",
    "{'cut':-1, 'split':_densenet_split}\n",
    "```\n",
    "\n",
    "where:\n",
    "```\n",
    "def _densenet_split(m:nn.Module): return (m[0][0][7],m[1])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), densenet121, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', '_DenseBlock', '_Transition', '_DenseBlock']\n",
      "Group 2: ['_Transition', '_DenseBlock', '_Transition', '_DenseBlock', 'BatchNorm2d']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Lambda', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Other Densenet has similar structures, e.g.:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 605 layers (total: 605)\n",
      "(1) Linear      : 1   layers (total: 606)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(densenet201)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) MaxPool2d   : 1   layers (total: 4)\n",
      "(4) _DenseBlock : 36  layers (total: 40)\n",
      "(5) _Transition : 4   layers (total: 44)\n",
      "(6) _DenseBlock : 72  layers (total: 116)\n",
      "(7) _Transition : 4   layers (total: 120)\n",
      "(8) _DenseBlock : 288 layers (total: 408)\n",
      "(9) _Transition : 4   layers (total: 412)\n",
      "(10) _DenseBlock : 192 layers (total: 604)\n",
      "(11) BatchNorm2d : 1   layers (total: 605)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(densenet201(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Inception"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We first look at Inception-v4, proposed in \"[Inception-ResNet and the Impact of Residual Connections on Learning](https://arxiv.org/pdf/1602.07261.pdf)\".\n",
    "\n",
    "We wrap the Cadene implementation into the torchvision model API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def inceptionv4(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.inceptionv4(pretrained=pretrained)\n",
    "    all_layers = list(model.children())\n",
    "    return nn.Sequential(*all_layers[0], *all_layers[1:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "It has the following structures: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) BasicConv2d : 3   layers (total: 3)\n",
      "(1) BasicConv2d : 3   layers (total: 6)\n",
      "(2) BasicConv2d : 3   layers (total: 9)\n",
      "(3) Mixed_3a    : 4   layers (total: 13)\n",
      "(4) Mixed_4a    : 18  layers (total: 31)\n",
      "(5) Mixed_5a    : 4   layers (total: 35)\n",
      "(6) Inception_A : 22  layers (total: 57)\n",
      "(7) Inception_A : 22  layers (total: 79)\n",
      "(8) Inception_A : 22  layers (total: 101)\n",
      "(9) Inception_A : 22  layers (total: 123)\n",
      "(10) Reduction_A : 13  layers (total: 136)\n",
      "(11) Inception_B : 31  layers (total: 167)\n",
      "(12) Inception_B : 31  layers (total: 198)\n",
      "(13) Inception_B : 31  layers (total: 229)\n",
      "(14) Inception_B : 31  layers (total: 260)\n",
      "(15) Inception_B : 31  layers (total: 291)\n",
      "(16) Inception_B : 31  layers (total: 322)\n",
      "(17) Inception_B : 31  layers (total: 353)\n",
      "(18) Reduction_B : 19  layers (total: 372)\n",
      "(19) Inception_C : 31  layers (total: 403)\n",
      "(20) Inception_C : 31  layers (total: 434)\n",
      "(21) Inception_C : 31  layers (total: 465)\n",
      "(22) AvgPool2d   : 1   layers (total: 466)\n",
      "(23) Linear      : 1   layers (total: 467)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(inceptionv4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "For transfer learning, we can cut out the last two layers and add our own custom head depending on the problem.  \n",
    "\n",
    "To group the layers for discriminative learning rates, we can set (0) ~ (10) as group 1, set (11) and (21) as group 2, and our own custom head as group 3. Now we specify the cut and split and create the learner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), inceptionv4, pretrained=False,\n",
    "                    cut=-2, split_on=lambda m: (m[0][11], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['BasicConv2d', 'BasicConv2d', 'BasicConv2d', 'Mixed_3a', 'Mixed_4a', 'Mixed_5a', 'Inception_A', 'Inception_A', 'Inception_A', 'Inception_A', 'Reduction_A']\n",
      "Group 2: ['Inception_B', 'Inception_B', 'Inception_B', 'Inception_B', 'Inception_B', 'Inception_B', 'Inception_B', 'Reduction_B', 'Inception_C', 'Inception_C', 'Inception_C']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Alternatively, you can add the cut and split to the model metadata before calling `create_cnn`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_inception_4_meta = { 'cut': -2, 'split': lambda m: (m[0][11], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[inceptionv4] = _inception_4_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Cadene's Inception ResNet V2 has similar implementation. Define:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def inceptionresnetv2(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.inceptionresnetv2(pretrained=pretrained)\n",
    "    return nn.Sequential(*model.children())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) BasicConv2d : 3   layers (total: 3)\n",
      "(1) BasicConv2d : 3   layers (total: 6)\n",
      "(2) BasicConv2d : 3   layers (total: 9)\n",
      "(3) MaxPool2d   : 1   layers (total: 10)\n",
      "(4) BasicConv2d : 3   layers (total: 13)\n",
      "(5) BasicConv2d : 3   layers (total: 16)\n",
      "(6) MaxPool2d   : 1   layers (total: 17)\n",
      "(7) Mixed_5b    : 22  layers (total: 39)\n",
      "(8) Sequential  : 200 layers (total: 239)\n",
      "(9) Mixed_6a    : 13  layers (total: 252)\n",
      "(10) Sequential  : 280 layers (total: 532)\n",
      "(11) Mixed_7a    : 22  layers (total: 554)\n",
      "(12) Sequential  : 126 layers (total: 680)\n",
      "(13) Block8      : 13  layers (total: 693)\n",
      "(14) BasicConv2d : 3   layers (total: 696)\n",
      "(15) AvgPool2d   : 1   layers (total: 697)\n",
      "(16) Linear      : 1   layers (total: 698)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(inceptionresnetv2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), inceptionresnetv2, pretrained=False,\n",
    "                    cut=-2, split_on=lambda m: (m[0][9], m[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['BasicConv2d', 'BasicConv2d', 'BasicConv2d', 'MaxPool2d', 'BasicConv2d', 'BasicConv2d', 'MaxPool2d', 'Mixed_5b', 'Sequential']\n",
      "Group 2: ['Mixed_6a', 'Sequential', 'Mixed_7a', 'Sequential', 'Block8', 'BasicConv2d']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Although we used it with almost same code, the underlying structures are completely different for different versions of Inception, as the models are redesigned in each generation. The torchvision's implementation of Inception V3 has the structure: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) BasicConv2d : 2   layers (total: 2)\n",
      "(1) BasicConv2d : 2   layers (total: 4)\n",
      "(2) BasicConv2d : 2   layers (total: 6)\n",
      "(3) BasicConv2d : 2   layers (total: 8)\n",
      "(4) BasicConv2d : 2   layers (total: 10)\n",
      "(5) InceptionA  : 14  layers (total: 24)\n",
      "(6) InceptionA  : 14  layers (total: 38)\n",
      "(7) InceptionA  : 14  layers (total: 52)\n",
      "(8) InceptionB  : 8   layers (total: 60)\n",
      "(9) InceptionC  : 20  layers (total: 80)\n",
      "(10) InceptionC  : 20  layers (total: 100)\n",
      "(11) InceptionC  : 20  layers (total: 120)\n",
      "(12) InceptionC  : 20  layers (total: 140)\n",
      "(13) InceptionAux: 5   layers (total: 145)\n",
      "(14) InceptionD  : 12  layers (total: 157)\n",
      "(15) InceptionE  : 18  layers (total: 175)\n",
      "(16) InceptionE  : 18  layers (total: 193)\n",
      "(17) Linear      : 1   layers (total: 194)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(inception_v3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "__Note__: some layers are not shown in here, e.g., between (2) and (3) there's a max pooling layer. This is because torchvision implementation uses `F.max_pool2d` in `forward`, which won't be included as children. So we can't cut and split this one like what we did for other models. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "If you'd really like to use this model, here's a workaround. Now, because of these `F.max_pool2d` functions, we must pass the model as a whole to fastai; at the same time, we'll want to cut off layers after `F.avg_pool2d`. How do you use it as a whole but only want part of it? The only way (that I can come up with) to solve this contradiction is to monkey patch the `forward` function. Copy it without layers after `F.avg_pool2d`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def forward_inception_v3(self, x):\n",
    "    x = self.Conv2d_1a_3x3(x)\n",
    "    x = self.Conv2d_2a_3x3(x)\n",
    "    x = self.Conv2d_2b_3x3(x)\n",
    "    x = F.max_pool2d(x, kernel_size=3, stride=2)\n",
    "    x = self.Conv2d_3b_1x1(x)\n",
    "    x = self.Conv2d_4a_3x3(x)\n",
    "    x = F.max_pool2d(x, kernel_size=3, stride=2)\n",
    "    x = self.Mixed_5b(x)\n",
    "    x = self.Mixed_5c(x)\n",
    "    x = self.Mixed_5d(x)\n",
    "    x = self.Mixed_6a(x)\n",
    "    x = self.Mixed_6b(x)\n",
    "    x = self.Mixed_6c(x)        \n",
    "    x = self.Mixed_6d(x)        \n",
    "    x = self.Mixed_6e(x)        \n",
    "    if self.training and self.aux_logits:\n",
    "        aux = self.AuxLogits(x)\n",
    "    x = self.Mixed_7a(x)        \n",
    "    x = self.Mixed_7b(x)        \n",
    "    x = self.Mixed_7c(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "And override it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "Inception3.forward = forward_inception_v3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Wrap the model into a `nn.Sequential` as a whole:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def inception_v3_cut(pretrained=False):    \n",
    "    model = inception_v3(pretrained)\n",
    "    return nn.Sequential(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Add the `cut` and `split` to `model_meta`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[inception_v3_cut] =  { 'cut': noop, \n",
    "                                  'split': lambda m: (list(m[0][0].children())[9], m[1]) }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Note: it _won't_ work if you specify `cut` and `split` in `create_cnn`, the `cut` will be overriden by the `_default_meta`. \n",
    "\n",
    "Now we can create the learner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), inception_v3_cut, pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['BasicConv2d', 'BasicConv2d', 'BasicConv2d', 'BasicConv2d', 'BasicConv2d', 'InceptionA', 'InceptionA', 'InceptionA', 'InceptionB']\n",
      "Group 2: ['InceptionC', 'InceptionC', 'InceptionC', 'InceptionC', 'InceptionAux', 'InceptionD', 'InceptionE', 'InceptionE', 'Linear']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0].children(), *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "It works! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Wide ResNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Wide ResNet (WRN) is proposed in \"[Wide Residual Networks](https://arxiv.org/pdf/1605.07146.pdf)\". \n",
    "\n",
    "We take a look at the fastai's implementation: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 54  layers (total: 54)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: wrn_22())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We look into this `Sequential` module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BasicBlock  : 6   layers (total: 7)\n",
      "(2) BasicBlock  : 5   layers (total: 12)\n",
      "(3) BasicBlock  : 5   layers (total: 17)\n",
      "(4) BasicBlock  : 6   layers (total: 23)\n",
      "(5) BasicBlock  : 5   layers (total: 28)\n",
      "(6) BasicBlock  : 5   layers (total: 33)\n",
      "(7) BasicBlock  : 6   layers (total: 39)\n",
      "(8) BasicBlock  : 5   layers (total: 44)\n",
      "(9) BasicBlock  : 5   layers (total: 49)\n",
      "(10) BatchNorm2d : 1   layers (total: 50)\n",
      "(11) ReLU        : 1   layers (total: 51)\n",
      "(12) AdaptiveAvgPool2d: 1   layers (total: 52)\n",
      "(13) Flatten     : 1   layers (total: 53)\n",
      "(14) Linear      : 1   layers (total: 54)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(wrn_22().children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Fastai didn't implement the `wrn_22` in the torchvision model API, probably because the lack of pretrained models for WRN. It'll not be very useful if we don't have pretrained model, we'll have to train it from the beginning. But for completeness, we wrap it and test our cut and split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def w_rn_22(pretrained=False):\n",
    "    return nn.Sequential(*list(next(wrn_22().children())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_wrn_22_meta = { 'cut': None, 'split': lambda m: (m[0][5], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[w_rn_22] = _wrn_22_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), w_rn_22, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BasicBlock', 'BasicBlock', 'BasicBlock', 'BasicBlock']\n",
      "Group 2: ['BasicBlock', 'BasicBlock', 'BasicBlock', 'BasicBlock', 'BasicBlock', 'BatchNorm2d', 'ReLU', 'AdaptiveAvgPool2d', 'Flatten', 'Linear']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## SqueezeNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "SqueezeNet is proposed in \"[SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://arxiv.org/abs/1602.07360)\". \n",
    "\n",
    "We take a look at the torchvision implementation: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 53  layers (total: 53)\n",
      "(1) Sequential  : 4   layers (total: 57)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(squeezenet1_0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The first `Sequential` module contains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) ReLU        : 1   layers (total: 2)\n",
      "(2) MaxPool2d   : 1   layers (total: 3)\n",
      "(3) Fire        : 6   layers (total: 9)\n",
      "(4) Fire        : 6   layers (total: 15)\n",
      "(5) Fire        : 6   layers (total: 21)\n",
      "(6) MaxPool2d   : 1   layers (total: 22)\n",
      "(7) Fire        : 6   layers (total: 28)\n",
      "(8) Fire        : 6   layers (total: 34)\n",
      "(9) Fire        : 6   layers (total: 40)\n",
      "(10) Fire        : 6   layers (total: 46)\n",
      "(11) MaxPool2d   : 1   layers (total: 47)\n",
      "(12) Fire        : 6   layers (total: 53)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(squeezenet1_0(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The second `Sequential` module contains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Dropout     : 1   layers (total: 1)\n",
      "(1) Conv2d      : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) AvgPool2d   : 1   layers (total: 4)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: list(squeezenet1_0(False).children())[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So we want to cut the second `Sequential` module out. We can split the first `Sequential` module on the layer (7) (_note_: this will override the current split defined in the fastai library):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_squeezenet_meta = { 'cut': -1, 'split': lambda m: (m[0][0][7], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[squeezenet1_0] = _squeezenet_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "hidden": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), squeezenet1_0, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'ReLU', 'MaxPool2d', 'Fire', 'Fire', 'Fire', 'MaxPool2d']\n",
      "Group 2: ['Fire', 'Fire', 'Fire', 'Fire', 'MaxPool2d', 'Fire']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Xception"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Xception is proposed in [Xception: Deep Learning with Depthwise Separable Convolutions](https://arxiv.org/pdf/1610.02357.pdf). Wrap the Cadene implementation into torchvision model API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def xception(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.xception(pretrained=pretrained)\n",
    "    return nn.Sequential(*list(model.children()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) Conv2d      : 1   layers (total: 4)\n",
      "(4) BatchNorm2d : 1   layers (total: 5)\n",
      "(5) Block       : 11  layers (total: 16)\n",
      "(6) Block       : 12  layers (total: 28)\n",
      "(7) Block       : 12  layers (total: 40)\n",
      "(8) Block       : 12  layers (total: 52)\n",
      "(9) Block       : 12  layers (total: 64)\n",
      "(10) Block       : 12  layers (total: 76)\n",
      "(11) Block       : 12  layers (total: 88)\n",
      "(12) Block       : 12  layers (total: 100)\n",
      "(13) Block       : 12  layers (total: 112)\n",
      "(14) Block       : 12  layers (total: 124)\n",
      "(15) Block       : 12  layers (total: 136)\n",
      "(16) Block       : 12  layers (total: 148)\n",
      "(17) SeparableConv2d: 2   layers (total: 150)\n",
      "(18) BatchNorm2d : 1   layers (total: 151)\n",
      "(19) SeparableConv2d: 2   layers (total: 153)\n",
      "(20) BatchNorm2d : 1   layers (total: 154)\n",
      "(21) Linear      : 1   layers (total: 155)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(xception)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We can cut out the last layer, and split at (11):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), xception, pretrained=False,\n",
    "                    cut=-1, split_on=lambda m: (m[0][11], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Alternatively, you can add the cut and split to the model metadata before calling `create_cnn`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_xception_meta = { 'cut': -1, 'split': lambda m: (m[0][11], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[xception] = _xception_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'Block', 'Block', 'Block', 'Block', 'Block', 'Block']\n",
      "Group 2: ['Block', 'Block', 'Block', 'Block', 'Block', 'Block', 'SeparableConv2d', 'BatchNorm2d', 'SeparableConv2d', 'BatchNorm2d']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## DPN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Dual Path Networks (DPN) is proposed in \"[Dual Path Networks](https://arxiv.org/abs/1707.01629)\". Wrap the Cadene implementation into torchvision model API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def dpn92(pretrained=False):\n",
    "    pretrained = 'imagenet+5k' if pretrained else None\n",
    "    model = pretrainedmodels.dpn92(pretrained=pretrained)\n",
    "    return nn.Sequential(*list(model.children()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 288 layers (total: 288)\n",
      "(1) Conv2d      : 1   layers (total: 289)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(dpn92)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The first `Sequential` contains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) InputBlock  : 4   layers (total: 4)\n",
      "(1) DualPathBlock: 12  layers (total: 16)\n",
      "(2) DualPathBlock: 9   layers (total: 25)\n",
      "(3) DualPathBlock: 9   layers (total: 34)\n",
      "(4) DualPathBlock: 12  layers (total: 46)\n",
      "(5) DualPathBlock: 9   layers (total: 55)\n",
      "(6) DualPathBlock: 9   layers (total: 64)\n",
      "(7) DualPathBlock: 9   layers (total: 73)\n",
      "(8) DualPathBlock: 12  layers (total: 85)\n",
      "(9) DualPathBlock: 9   layers (total: 94)\n",
      "(10) DualPathBlock: 9   layers (total: 103)\n",
      "(11) DualPathBlock: 9   layers (total: 112)\n",
      "(12) DualPathBlock: 9   layers (total: 121)\n",
      "(13) DualPathBlock: 9   layers (total: 130)\n",
      "(14) DualPathBlock: 9   layers (total: 139)\n",
      "(15) DualPathBlock: 9   layers (total: 148)\n",
      "(16) DualPathBlock: 9   layers (total: 157)\n",
      "(17) DualPathBlock: 9   layers (total: 166)\n",
      "(18) DualPathBlock: 9   layers (total: 175)\n",
      "(19) DualPathBlock: 9   layers (total: 184)\n",
      "(20) DualPathBlock: 9   layers (total: 193)\n",
      "(21) DualPathBlock: 9   layers (total: 202)\n",
      "(22) DualPathBlock: 9   layers (total: 211)\n",
      "(23) DualPathBlock: 9   layers (total: 220)\n",
      "(24) DualPathBlock: 9   layers (total: 229)\n",
      "(25) DualPathBlock: 9   layers (total: 238)\n",
      "(26) DualPathBlock: 9   layers (total: 247)\n",
      "(27) DualPathBlock: 9   layers (total: 256)\n",
      "(28) DualPathBlock: 12  layers (total: 268)\n",
      "(29) DualPathBlock: 9   layers (total: 277)\n",
      "(30) DualPathBlock: 9   layers (total: 286)\n",
      "(31) CatBnAct    : 2   layers (total: 288)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(dpn92(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "There're 30 `DualPathBlock`, we can split them into two groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), dpn92, pretrained=False,\n",
    "                    cut=-1, split_on=lambda m: (m[0][0][16], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['InputBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock']\n",
      "Group 2: ['DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'DualPathBlock', 'CatBnAct']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The customization works as expected. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Alternatively, you can add the cut and split to the model metadata before calling `create_cnn`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "_dpn_meta = {'cut': -1, 'split': lambda m: (m[0][0][16], m[1]) }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[dpn92] = _dpn_meta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## NASNetAMobile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "NASNet is proposed in \"[NASNet](https://arxiv.org/abs/1707.07012)\". Wrap the Cadene implementation into torchvision model API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def identity(x): return x\n",
    "\n",
    "def nasnetamobile(pretrained=False):\n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.nasnetamobile(pretrained=pretrained, num_classes=1000)  \n",
    "    model.logits = identity\n",
    "    return nn.Sequential(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Note the `logits` part is replaced by identity because we'll use fastai later to replace this part, which contains the last few layers. Also, if you use lambda function for the identity, the model won't be picklable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 2   layers (total: 2)\n",
      "(1) CellStem0   : 47  layers (total: 49)\n",
      "(2) CellStem1   : 57  layers (total: 106)\n",
      "(3) FirstCell   : 53  layers (total: 159)\n",
      "(4) NormalCell  : 49  layers (total: 208)\n",
      "(5) NormalCell  : 49  layers (total: 257)\n",
      "(6) NormalCell  : 49  layers (total: 306)\n",
      "(7) ReductionCell0: 58  layers (total: 364)\n",
      "(8) FirstCell   : 53  layers (total: 417)\n",
      "(9) NormalCell  : 49  layers (total: 466)\n",
      "(10) NormalCell  : 49  layers (total: 515)\n",
      "(11) NormalCell  : 49  layers (total: 564)\n",
      "(12) ReductionCell1: 53  layers (total: 617)\n",
      "(13) FirstCell   : 53  layers (total: 670)\n",
      "(14) NormalCell  : 49  layers (total: 719)\n",
      "(15) NormalCell  : 49  layers (total: 768)\n",
      "(16) NormalCell  : 49  layers (total: 817)\n",
      "(17) ReLU        : 1   layers (total: 818)\n",
      "(18) AvgPool2d   : 1   layers (total: 819)\n",
      "(19) Dropout     : 1   layers (total: 820)\n",
      "(20) Linear      : 1   layers (total: 821)\n",
      "(21) Identity    : 1   layers (total: 822)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: nasnetamobile(False)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We need to set the `cut` to `noop` as we have already replaced the last few layers (as I'll explain later, we can't cut directly). For learning rates, layer (8) is a good point to split. So we set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[nasnetamobile] =  { 'cut': noop, \n",
    "                               'split': lambda m: (list(m[0][0].children())[8], m[1]) }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create the learner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), nasnetamobile, pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Sequential', 'CellStem0', 'CellStem1', 'FirstCell', 'NormalCell', 'NormalCell', 'NormalCell', 'ReductionCell0']\n",
      "Group 2: ['FirstCell', 'NormalCell', 'NormalCell', 'NormalCell', 'ReductionCell1', 'FirstCell', 'NormalCell', 'NormalCell', 'NormalCell', 'ReLU', 'AvgPool2d', 'Dropout', 'Linear']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0].children(), *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The groups for different learning rates are indeed split at layer (8). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "__NOTE__: why do I need to replace the `logits`, rather than just set `cut=-1` as usual? If you set `cut=-1`, then in `create_cnn` the model will be converted into: \n",
    "\n",
    "`nn.Sequential(*list(model.children())[:cut]`\n",
    "\n",
    "And this new model doesn't work. It's because `nn.Sequential` expects the modules to have exactly one argument in `forward`. However in this case, a few layers, such as `CellStem1`, take two arguments in `forward`. So we don't want to convert the model, we must use it as a whole. \n",
    "\n",
    "Now if we use the model as a whole, how can we cut the last few layers? My solution is to change it into an identity function, and it seems to work fine.\n",
    "\n",
    "Now even we have used the model as a whole, we can still set different learning rates. In the `split`, `m[0][0]` is the `NASNetAMobile` object (it's wrapped by two `nn.Sequential`, both with 1 layer only), it's not iterable so we need to make a list of its children. This time we're fine, because the `split` is only used to set different learning rates, we won't call `forward` in here. \n",
    "\n",
    "A final note: this model doesn't seem to work for small images (smaller than 64x64). Apparently the original model is designed for ImageNet, although we replaced with adaptive pooling layer, some parts in the model body still don't work with small images. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## PNASNet-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "PNASNet-5 is proposed in \"[Progressive Neural Architecture Search](https://arxiv.org/abs/1712.00559)\". Wrap the Cadene implementation into torchvision model API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def identity(x): return x\n",
    "\n",
    "def pnasnet5large(pretrained=False):    \n",
    "    pretrained = 'imagenet' if pretrained else None\n",
    "    model = pretrainedmodels.pnasnet5large(pretrained=pretrained, num_classes=1000) \n",
    "    model.logits = identity\n",
    "    return nn.Sequential(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Note the `logits` part is replaced by identity, see the \"NASNetAMobile\" section for more explanations. Same as above, if you use lambda function for the identity, the model won't be picklable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 2   layers (total: 2)\n",
      "(1) CellStem0   : 59  layers (total: 61)\n",
      "(2) Cell        : 64  layers (total: 125)\n",
      "(3) Cell        : 61  layers (total: 186)\n",
      "(4) Cell        : 57  layers (total: 243)\n",
      "(5) Cell        : 57  layers (total: 300)\n",
      "(6) Cell        : 57  layers (total: 357)\n",
      "(7) Cell        : 68  layers (total: 425)\n",
      "(8) Cell        : 61  layers (total: 486)\n",
      "(9) Cell        : 57  layers (total: 543)\n",
      "(10) Cell        : 57  layers (total: 600)\n",
      "(11) Cell        : 60  layers (total: 660)\n",
      "(12) Cell        : 61  layers (total: 721)\n",
      "(13) Cell        : 57  layers (total: 778)\n",
      "(14) Cell        : 57  layers (total: 835)\n",
      "(15) ReLU        : 1   layers (total: 836)\n",
      "(16) AvgPool2d   : 1   layers (total: 837)\n",
      "(17) Dropout     : 1   layers (total: 838)\n",
      "(18) Linear      : 1   layers (total: 839)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: pnasnet5large(False)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "We need to set the `cut` to `noop` as we have already replaced the last few layers (we can't cut directly, see the \"NASNetAMobile\" section for more explanations). For learning rates, layer (8) is a good point to split. So we set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "model_meta[pnasnet5large] =  { 'cut': noop, \n",
    "                               'split': lambda m: (list(m[0][0].children())[8], m[1]) }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create the learner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), pnasnet5large, pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Sequential', 'CellStem0', 'Cell', 'Cell', 'Cell', 'Cell', 'Cell', 'Cell']\n",
      "Group 2: ['Cell', 'Cell', 'Cell', 'Cell', 'Cell', 'Cell', 'Cell', 'ReLU', 'AvgPool2d', 'Dropout', 'Linear']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0].children(), *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The groups for different learning rates are indeed split at layer (8)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Again, we did something special here, such as replacing the `logits` rather than just setting `cut=-1`. See the \"NASNetAMobile\" section for more explanations (for the same reason, this model doesn't work for small images (size smaller than 64x64)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## VGG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "VGG is proposed in \"[Very Deep Convolutional Networks for Large-Scale Image Recognition](https://arxiv.org/pdf/1409.1556.pdf)\". In the torchvision implementation, the VGG-16 model has the following structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "hidden": true,
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 44  layers (total: 44)\n",
      "(1) Sequential  : 7   layers (total: 51)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(vgg16_bn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The first `Sequential` has:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) Conv2d      : 1   layers (total: 4)\n",
      "(4) BatchNorm2d : 1   layers (total: 5)\n",
      "(5) ReLU        : 1   layers (total: 6)\n",
      "(6) MaxPool2d   : 1   layers (total: 7)\n",
      "(7) Conv2d      : 1   layers (total: 8)\n",
      "(8) BatchNorm2d : 1   layers (total: 9)\n",
      "(9) ReLU        : 1   layers (total: 10)\n",
      "(10) Conv2d      : 1   layers (total: 11)\n",
      "(11) BatchNorm2d : 1   layers (total: 12)\n",
      "(12) ReLU        : 1   layers (total: 13)\n",
      "(13) MaxPool2d   : 1   layers (total: 14)\n",
      "(14) Conv2d      : 1   layers (total: 15)\n",
      "(15) BatchNorm2d : 1   layers (total: 16)\n",
      "(16) ReLU        : 1   layers (total: 17)\n",
      "(17) Conv2d      : 1   layers (total: 18)\n",
      "(18) BatchNorm2d : 1   layers (total: 19)\n",
      "(19) ReLU        : 1   layers (total: 20)\n",
      "(20) Conv2d      : 1   layers (total: 21)\n",
      "(21) BatchNorm2d : 1   layers (total: 22)\n",
      "(22) ReLU        : 1   layers (total: 23)\n",
      "(23) MaxPool2d   : 1   layers (total: 24)\n",
      "(24) Conv2d      : 1   layers (total: 25)\n",
      "(25) BatchNorm2d : 1   layers (total: 26)\n",
      "(26) ReLU        : 1   layers (total: 27)\n",
      "(27) Conv2d      : 1   layers (total: 28)\n",
      "(28) BatchNorm2d : 1   layers (total: 29)\n",
      "(29) ReLU        : 1   layers (total: 30)\n",
      "(30) Conv2d      : 1   layers (total: 31)\n",
      "(31) BatchNorm2d : 1   layers (total: 32)\n",
      "(32) ReLU        : 1   layers (total: 33)\n",
      "(33) MaxPool2d   : 1   layers (total: 34)\n",
      "(34) Conv2d      : 1   layers (total: 35)\n",
      "(35) BatchNorm2d : 1   layers (total: 36)\n",
      "(36) ReLU        : 1   layers (total: 37)\n",
      "(37) Conv2d      : 1   layers (total: 38)\n",
      "(38) BatchNorm2d : 1   layers (total: 39)\n",
      "(39) ReLU        : 1   layers (total: 40)\n",
      "(40) Conv2d      : 1   layers (total: 41)\n",
      "(41) BatchNorm2d : 1   layers (total: 42)\n",
      "(42) ReLU        : 1   layers (total: 43)\n",
      "(43) MaxPool2d   : 1   layers (total: 44)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(vgg16_bn(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The second `Sequential` has:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Linear      : 1   layers (total: 1)\n",
      "(1) ReLU        : 1   layers (total: 2)\n",
      "(2) Dropout     : 1   layers (total: 3)\n",
      "(3) Linear      : 1   layers (total: 4)\n",
      "(4) ReLU        : 1   layers (total: 5)\n",
      "(5) Dropout     : 1   layers (total: 6)\n",
      "(6) Linear      : 1   layers (total: 7)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: list(vgg16_bn(False).children())[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So we'll cut the second `Sequential` out, and split the first `Sequential` at (22):\n",
    "\n",
    "```\n",
    "{'cut':-1, 'split':_vgg_split}\n",
    "```\n",
    "\n",
    "where:\n",
    "```\n",
    "def _vgg_split(m:nn.Module): return (m[0][0][22], m[1])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), vgg16_bn, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d']\n",
      "Group 2: ['ReLU', 'MaxPool2d', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'Conv2d', 'BatchNorm2d', 'ReLU', 'MaxPool2d']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "VGG-19 has a similar structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 53  layers (total: 53)\n",
      "(1) Sequential  : 7   layers (total: 60)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(vgg19_bn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) Conv2d      : 1   layers (total: 4)\n",
      "(4) BatchNorm2d : 1   layers (total: 5)\n",
      "(5) ReLU        : 1   layers (total: 6)\n",
      "(6) MaxPool2d   : 1   layers (total: 7)\n",
      "(7) Conv2d      : 1   layers (total: 8)\n",
      "(8) BatchNorm2d : 1   layers (total: 9)\n",
      "(9) ReLU        : 1   layers (total: 10)\n",
      "(10) Conv2d      : 1   layers (total: 11)\n",
      "(11) BatchNorm2d : 1   layers (total: 12)\n",
      "(12) ReLU        : 1   layers (total: 13)\n",
      "(13) MaxPool2d   : 1   layers (total: 14)\n",
      "(14) Conv2d      : 1   layers (total: 15)\n",
      "(15) BatchNorm2d : 1   layers (total: 16)\n",
      "(16) ReLU        : 1   layers (total: 17)\n",
      "(17) Conv2d      : 1   layers (total: 18)\n",
      "(18) BatchNorm2d : 1   layers (total: 19)\n",
      "(19) ReLU        : 1   layers (total: 20)\n",
      "(20) Conv2d      : 1   layers (total: 21)\n",
      "(21) BatchNorm2d : 1   layers (total: 22)\n",
      "(22) ReLU        : 1   layers (total: 23)\n",
      "(23) Conv2d      : 1   layers (total: 24)\n",
      "(24) BatchNorm2d : 1   layers (total: 25)\n",
      "(25) ReLU        : 1   layers (total: 26)\n",
      "(26) MaxPool2d   : 1   layers (total: 27)\n",
      "(27) Conv2d      : 1   layers (total: 28)\n",
      "(28) BatchNorm2d : 1   layers (total: 29)\n",
      "(29) ReLU        : 1   layers (total: 30)\n",
      "(30) Conv2d      : 1   layers (total: 31)\n",
      "(31) BatchNorm2d : 1   layers (total: 32)\n",
      "(32) ReLU        : 1   layers (total: 33)\n",
      "(33) Conv2d      : 1   layers (total: 34)\n",
      "(34) BatchNorm2d : 1   layers (total: 35)\n",
      "(35) ReLU        : 1   layers (total: 36)\n",
      "(36) Conv2d      : 1   layers (total: 37)\n",
      "(37) BatchNorm2d : 1   layers (total: 38)\n",
      "(38) ReLU        : 1   layers (total: 39)\n",
      "(39) MaxPool2d   : 1   layers (total: 40)\n",
      "(40) Conv2d      : 1   layers (total: 41)\n",
      "(41) BatchNorm2d : 1   layers (total: 42)\n",
      "(42) ReLU        : 1   layers (total: 43)\n",
      "(43) Conv2d      : 1   layers (total: 44)\n",
      "(44) BatchNorm2d : 1   layers (total: 45)\n",
      "(45) ReLU        : 1   layers (total: 46)\n",
      "(46) Conv2d      : 1   layers (total: 47)\n",
      "(47) BatchNorm2d : 1   layers (total: 48)\n",
      "(48) ReLU        : 1   layers (total: 49)\n",
      "(49) Conv2d      : 1   layers (total: 50)\n",
      "(50) BatchNorm2d : 1   layers (total: 51)\n",
      "(51) ReLU        : 1   layers (total: 52)\n",
      "(52) MaxPool2d   : 1   layers (total: 53)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(vgg19_bn(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The first `Sequential` has more layers than VGG-16, but we can still spit it at (22), it doesn't matter that much. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Linear      : 1   layers (total: 1)\n",
      "(1) ReLU        : 1   layers (total: 2)\n",
      "(2) Dropout     : 1   layers (total: 3)\n",
      "(3) Linear      : 1   layers (total: 4)\n",
      "(4) ReLU        : 1   layers (total: 5)\n",
      "(5) Dropout     : 1   layers (total: 6)\n",
      "(6) Linear      : 1   layers (total: 7)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: list(vgg19_bn(False).children())[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## AlexNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "AlexNet is proposed in \"[One weird trick for parallelizing convolutional neural networks](https://arxiv.org/abs/1404.5997)\". In the torchvision implementation, the model has the following structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Sequential  : 13  layers (total: 13)\n",
      "(1) Sequential  : 7   layers (total: 20)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(alexnet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The first `Sequential` has:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) ReLU        : 1   layers (total: 2)\n",
      "(2) MaxPool2d   : 1   layers (total: 3)\n",
      "(3) Conv2d      : 1   layers (total: 4)\n",
      "(4) ReLU        : 1   layers (total: 5)\n",
      "(5) MaxPool2d   : 1   layers (total: 6)\n",
      "(6) Conv2d      : 1   layers (total: 7)\n",
      "(7) ReLU        : 1   layers (total: 8)\n",
      "(8) Conv2d      : 1   layers (total: 9)\n",
      "(9) ReLU        : 1   layers (total: 10)\n",
      "(10) Conv2d      : 1   layers (total: 11)\n",
      "(11) ReLU        : 1   layers (total: 12)\n",
      "(12) MaxPool2d   : 1   layers (total: 13)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: next(alexnet(False).children()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "The second `Sequential` has:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Dropout     : 1   layers (total: 1)\n",
      "(1) Linear      : 1   layers (total: 2)\n",
      "(2) ReLU        : 1   layers (total: 3)\n",
      "(3) Dropout     : 1   layers (total: 4)\n",
      "(4) Linear      : 1   layers (total: 5)\n",
      "(5) ReLU        : 1   layers (total: 6)\n",
      "(6) Linear      : 1   layers (total: 7)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: list(alexnet(False).children())[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "So we'll cut the second `Sequential` out, and split the first `Sequential` at (6):\n",
    "\n",
    "```\n",
    "{'cut':-1, 'split':_alexnet_split}\n",
    "```\n",
    "\n",
    "where:\n",
    "```\n",
    "def _alexnet_split(m:nn.Module): return (m[0][0][6], m[1])\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "Now we can create a learner: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "hidden": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), alexnet, pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "hidden": true
   },
   "source": [
    "To check the cut and split work as we expected, we extract the groups: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2d', 'ReLU', 'MaxPool2d', 'Conv2d', 'ReLU', 'MaxPool2d']\n",
      "Group 2: ['Conv2d', 'ReLU', 'Conv2d', 'ReLU', 'Conv2d', 'ReLU', 'MaxPool2d']\n",
      "Group 3: ['AdaptiveConcatPool2d', 'Flatten', 'BatchNorm1d', 'Dropout', 'Linear', 'ReLU', 'BatchNorm1d', 'Dropout', 'Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0][0], *learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EfficientNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "EfficientNet is proposed in \"[EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://arxiv.org/abs/1905.11946)\". This [PyTorch implementation](https://github.com/lukemelas/EfficientNet-PyTorch) is used in here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from efficientnet_pytorch import EfficientNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take a look at the basic structures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = EfficientNet.from_pretrained('efficientnet-b0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) Conv2dSamePadding: 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "(2) ModuleList  : 126 layers (total: 128)\n",
      "(3) Conv2dSamePadding: 1   layers (total: 129)\n",
      "(4) BatchNorm2d : 1   layers (total: 130)\n",
      "(5) Linear      : 1   layers (total: 131)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) MBConvBlock : 6   layers (total: 6)\n",
      "(1) MBConvBlock : 8   layers (total: 14)\n",
      "(2) MBConvBlock : 8   layers (total: 22)\n",
      "(3) MBConvBlock : 8   layers (total: 30)\n",
      "(4) MBConvBlock : 8   layers (total: 38)\n",
      "(5) MBConvBlock : 8   layers (total: 46)\n",
      "(6) MBConvBlock : 8   layers (total: 54)\n",
      "(7) MBConvBlock : 8   layers (total: 62)\n",
      "(8) MBConvBlock : 8   layers (total: 70)\n",
      "(9) MBConvBlock : 8   layers (total: 78)\n",
      "(10) MBConvBlock : 8   layers (total: 86)\n",
      "(11) MBConvBlock : 8   layers (total: 94)\n",
      "(12) MBConvBlock : 8   layers (total: 102)\n",
      "(13) MBConvBlock : 8   layers (total: 110)\n",
      "(14) MBConvBlock : 8   layers (total: 118)\n",
      "(15) MBConvBlock : 8   layers (total: 126)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: list(m.children())[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrap it in a function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"pretrained\" is hardcoded to adapt to the PyTorch model function\n",
    "def efficient_net_b0(pretrained=True):\n",
    "    model = EfficientNet.from_pretrained('efficientnet-b0')\n",
    "    return nn.Sequential(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that it's not very easy to initialize a non-pretrained instance (and we'll want to use the pretrained anyway), so the \"pretrained\" parameter is hardcoded here. The `efficient_net_b0` does nothing, it just makes it easier to feed to fastai's `cnn_learner`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to set the `cut` to `noop`, see the \"NASNetAMobile\" section for more explanations. For learning rates, I divided the 18 `MBConvBlock` blocks into two parts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_meta[efficient_net_b0] =  { 'cut': noop, \n",
    "                               'split': lambda m: (list(m[0][0].children())[2][7], m[1]) }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the original model, the final output size is 1000:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained weights for efficientnet-b0\n"
     ]
    }
   ],
   "source": [
    "output_size = list(efficient_net_b0()[0].children())[-1].out_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And an adaptive layer is already added, so we don't want to do it again. In here, I'll just add a custom head to adapt the output size to our problem (in here, it has two classes):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = FakeData()\n",
    "custom_head = nn.Linear(output_size, data.c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained weights for efficientnet-b0\n"
     ]
    }
   ],
   "source": [
    "learn = cnn_learner(data, efficient_net_b0, custom_head = custom_head)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the learning rate groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['Conv2dSamePadding', 'BatchNorm2d', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock']\n",
      "Group 2: ['MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'MBConvBlock', 'Conv2dSamePadding', 'BatchNorm2d', 'Linear']\n",
      "Group 3: ['Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*list(learn.model[0][0].children())[:2], \n",
    "                         *list(learn.model[0][0].children())[2],\n",
    "                         *list(learn.model[0][0].children())[3:],\n",
    "                         learn.model[1]), \n",
    "           learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It works as expected. Note we can add more layers to group 3, just put it into the `custom_head`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test with a random image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-2.5571,  0.0241]])\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    learn.model.eval()\n",
    "    print(learn.model(torch.randn(1,3,96,96)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## _To Be Continued..._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}