CNN_archs/cnn_archs_more.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# More CNN Architectures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The CNN models in [this notebook (cnn_archs)](cnn_archs.ipynb) are all from [torchvision](https://pytorch.org/docs/stable/torchvision/index.html) and [Cadene](https://github.com/Cadene/pretrained-models.pytorch). Here's another [great repo (pytorchcv)](https://github.com/osmr/imgclsmob/tree/master/pytorch) providing even more comprehensive list of models. Moreover, the implementation is _much_ easier to use in fastai, comparing with Cadene. E.g., the model body are encapsulated in `features`, which is a `Sequential`. With this, for fastai's `cut`, we can just take the `features` part; for fastai's `split`, since it's a `Sequential`, we can just slice it!\n",
    "\n",
    "Here're some examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorchcv.model_provider import get_model as ptcv_get_model\n",
    "from torchvision.models import *\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": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.53.dev0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ptcv_get_model(\"resnet18\", pretrained=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model body is in the `features`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) ResInitBlock: 4   layers (total: 4)\n",
      "(1) Sequential  : 12  layers (total: 16)\n",
      "(2) Sequential  : 14  layers (total: 30)\n",
      "(3) Sequential  : 14  layers (total: 44)\n",
      "(4) Sequential  : 14  layers (total: 58)\n",
      "(5) AvgPool2d   : 1   layers (total: 59)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m.features)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model head is in the `output`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=512, out_features=1000, bias=True)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When `features` are `Sequential`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When the `features` are `Sequential`, it's very straightforward for fastai's `split`: just slice it! Here're some examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### inceptionv3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ptcv_get_model(\"inceptionv3\", pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) InceptInitBlock: 17  layers (total: 17)\n",
      "(1) Sequential  : 66  layers (total: 83)\n",
      "(2) Sequential  : 137 layers (total: 220)\n",
      "(3) Sequential  : 75  layers (total: 295)\n",
      "(4) AvgPool2d   : 1   layers (total: 296)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m.features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------(0)---------\n",
      "(0) InceptConv  : 3   layers (total: 3)\n",
      "(1) InceptConv  : 3   layers (total: 6)\n",
      "(2) InceptConv  : 3   layers (total: 9)\n",
      "(3) MaxPool2d   : 1   layers (total: 10)\n",
      "(4) InceptConv  : 3   layers (total: 13)\n",
      "(5) InceptConv  : 3   layers (total: 16)\n",
      "(6) MaxPool2d   : 1   layers (total: 17)\n",
      "---------(1)---------\n",
      "(0) InceptionAUnit: 22  layers (total: 22)\n",
      "(1) InceptionAUnit: 22  layers (total: 44)\n",
      "(2) InceptionAUnit: 22  layers (total: 66)\n",
      "---------(2)---------\n",
      "(0) ReductionAUnit: 13  layers (total: 13)\n",
      "(1) InceptionBUnit: 31  layers (total: 44)\n",
      "(2) InceptionBUnit: 31  layers (total: 75)\n",
      "(3) InceptionBUnit: 31  layers (total: 106)\n",
      "(4) InceptionBUnit: 31  layers (total: 137)\n",
      "---------(3)---------\n",
      "(0) ReductionBUnit: 19  layers (total: 19)\n",
      "(1) InceptionCUnit: 28  layers (total: 47)\n",
      "(2) InceptionCUnit: 28  layers (total: 75)\n"
     ]
    }
   ],
   "source": [
    "for i in range(4):\n",
    "    print(f'---------({i})---------')\n",
    "    arch_summary(lambda _: m.features[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inceptionv3(pretrained=False):\n",
    "    return ptcv_get_model(\"inceptionv3\", pretrained=False).features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have already cut the head out, so don't need to anything for `cut`. Now create learner, split at layer (3): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), inceptionv3, pretrained=False,\n",
    "                    cut=noop, split_on=lambda m: (m[0][3], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the cut and split work as we expected, we extract the groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['InceptInitBlock', 'Sequential', 'Sequential']\n",
      "Group 2: ['Sequential', 'AvgPool2d']\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": {},
   "source": [
    "The customization works as expected. Isn't it much easier, comparing with what we did in [this notebook](cnn_archs.ipynb)?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### EfficientNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ptcv_get_model(\"efficientnet_b0\", pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.nn.modules.container.Sequential"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.features.__class__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) EffiInitBlock: 3   layers (total: 3)\n",
      "(1) Sequential  : 10  layers (total: 13)\n",
      "(2) Sequential  : 26  layers (total: 39)\n",
      "(3) Sequential  : 26  layers (total: 65)\n",
      "(4) Sequential  : 78  layers (total: 143)\n",
      "(5) Sequential  : 65  layers (total: 208)\n",
      "(6) ConvBlock   : 3   layers (total: 211)\n",
      "(7) AdaptiveAvgPool2d: 1   layers (total: 212)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m.features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------(0)---------\n",
      "(0) ConvBlock   : 3   layers (total: 3)\n",
      "---------(1)---------\n",
      "(0) EffiDwsConvUnit: 10  layers (total: 10)\n",
      "---------(2)---------\n",
      "(0) EffiInvResUnit: 13  layers (total: 13)\n",
      "(1) EffiInvResUnit: 13  layers (total: 26)\n",
      "---------(3)---------\n",
      "(0) EffiInvResUnit: 13  layers (total: 13)\n",
      "(1) EffiInvResUnit: 13  layers (total: 26)\n",
      "---------(4)---------\n",
      "(0) EffiInvResUnit: 13  layers (total: 13)\n",
      "(1) EffiInvResUnit: 13  layers (total: 26)\n",
      "(2) EffiInvResUnit: 13  layers (total: 39)\n",
      "(3) EffiInvResUnit: 13  layers (total: 52)\n",
      "(4) EffiInvResUnit: 13  layers (total: 65)\n",
      "(5) EffiInvResUnit: 13  layers (total: 78)\n",
      "---------(5)---------\n",
      "(0) EffiInvResUnit: 13  layers (total: 13)\n",
      "(1) EffiInvResUnit: 13  layers (total: 26)\n",
      "(2) EffiInvResUnit: 13  layers (total: 39)\n",
      "(3) EffiInvResUnit: 13  layers (total: 52)\n",
      "(4) EffiInvResUnit: 13  layers (total: 65)\n"
     ]
    }
   ],
   "source": [
    "for i in range(6):\n",
    "    print(f'---------({i})---------')\n",
    "    arch_summary(lambda _: m.features[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def efficientnet_b0(pretrained=False):\n",
    "    return ptcv_get_model(\"efficientnet_b0\", pretrained=False).features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create learner, split at layer (4): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = cnn_learner(FakeData(), efficientnet_b0, pretrained=False,\n",
    "                    cut=noop, split_on=lambda m: (m[0][4], m[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the cut and split work as we expected, we extract the groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['EffiInitBlock', 'Sequential', 'Sequential', 'Sequential']\n",
      "Group 2: ['Sequential', 'Sequential', 'ConvBlock', 'AdaptiveAvgPool2d']\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": {},
   "source": [
    "The customization works as expected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## When `features` are not `Sequential`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actually, in the author's implementation, the `features` are always `Sequential`. But in the following case, it inherits `Sequential` and customized the `forward` method. Although we can still split by slicing, but because of some hardcoded part in fastai, we'll have some error. Let's take a look."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NASNet-A-Mobile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ptcv_get_model(\"nasnet_4a1056\", pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pytorchcv.models.common.DualPathSequential"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.features.__class__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) NASNetInitBlock: 2   layers (total: 2)\n",
      "(1) Stem1Unit   : 47  layers (total: 49)\n",
      "(2) Stem2Unit   : 62  layers (total: 111)\n",
      "(3) DualPathSequential: 200 layers (total: 311)\n",
      "(4) DualPathSequential: 258 layers (total: 569)\n",
      "(5) DualPathSequential: 258 layers (total: 827)\n",
      "(6) ReLU        : 1   layers (total: 828)\n",
      "(7) AvgPool2d   : 1   layers (total: 829)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m.features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------(0)---------\n",
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "---------(1)---------\n",
      "(0) NasConv     : 3   layers (total: 3)\n",
      "(1) DwsBranch   : 8   layers (total: 11)\n",
      "(2) DwsBranch   : 8   layers (total: 19)\n",
      "(3) NasMaxPoolBlock: 1   layers (total: 20)\n",
      "(4) DwsBranch   : 8   layers (total: 28)\n",
      "(5) AvgPool2d   : 1   layers (total: 29)\n",
      "(6) DwsBranch   : 8   layers (total: 37)\n",
      "(7) AvgPool2d   : 1   layers (total: 38)\n",
      "(8) DwsBranch   : 8   layers (total: 46)\n",
      "(9) NasMaxPoolBlock: 1   layers (total: 47)\n",
      "---------(2)---------\n",
      "(0) NasConv     : 3   layers (total: 3)\n",
      "(1) NasPathBlock: 7   layers (total: 10)\n",
      "(2) DwsBranch   : 9   layers (total: 19)\n",
      "(3) DwsBranch   : 9   layers (total: 28)\n",
      "(4) NasMaxPoolBlock: 2   layers (total: 30)\n",
      "(5) DwsBranch   : 9   layers (total: 39)\n",
      "(6) NasAvgPoolBlock: 2   layers (total: 41)\n",
      "(7) DwsBranch   : 9   layers (total: 50)\n",
      "(8) AvgPool2d   : 1   layers (total: 51)\n",
      "(9) DwsBranch   : 9   layers (total: 60)\n",
      "(10) NasMaxPoolBlock: 2   layers (total: 62)\n",
      "---------(3)---------\n",
      "(0) FirstUnit   : 53  layers (total: 53)\n",
      "(1) NormalUnit  : 49  layers (total: 102)\n",
      "(2) NormalUnit  : 49  layers (total: 151)\n",
      "(3) NormalUnit  : 49  layers (total: 200)\n",
      "---------(4)---------\n",
      "(0) Reduction1Unit: 58  layers (total: 58)\n",
      "(1) FirstUnit   : 53  layers (total: 111)\n",
      "(2) NormalUnit  : 49  layers (total: 160)\n",
      "(3) NormalUnit  : 49  layers (total: 209)\n",
      "(4) NormalUnit  : 49  layers (total: 258)\n",
      "---------(5)---------\n",
      "(0) Reduction2Unit: 58  layers (total: 58)\n",
      "(1) FirstUnit   : 53  layers (total: 111)\n",
      "(2) NormalUnit  : 49  layers (total: 160)\n",
      "(3) NormalUnit  : 49  layers (total: 209)\n",
      "(4) NormalUnit  : 49  layers (total: 258)\n"
     ]
    }
   ],
   "source": [
    "for i in range(6):\n",
    "    print(f'---------({i})---------')\n",
    "    arch_summary(lambda _: m.features[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (dropout): Dropout(p=0.5)\n",
       "  (fc): Linear(in_features=1056, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_size = m.output[1].out_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nasnetamobile(pretrained=False):\n",
    "    return ptcv_get_model(\"nasnet_4a1056\", pretrained=False).features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this case, we can still take the `features` part, but we can't use it directly in fastai. \n",
    "\n",
    "Why? Because in fastai's `create_cnn_model`, if `custom_head` is not given, then it'll try to convert the body into a `Sequential` to determine the number of features (which is really not necessary, should have an option to let user set the number `nf`). \n",
    "\n",
    "The body in here was `DualPathSequential`, which has two parameters in the `forward`, so it won't work after converting. We must prevent fastai from reaching this line. Looks like right now the only way is to define a `custom_head` so this line will be bypassed. In the following example, a simple head is used for demonstration. You can use more complicated head, like the one fastai adds by default."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create learner, split at layer (4): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = FakeData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(data, nasnetamobile, pretrained=False,\n",
    "                    cut=noop, split_on=lambda m: (m[0][4], m[1]),\n",
    "                    custom_head=nn.Linear(output_size, data.c))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the cut and split work as we expected, we extract the groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['NASNetInitBlock', 'Stem1Unit', 'Stem2Unit', 'DualPathSequential']\n",
      "Group 2: ['DualPathSequential', 'DualPathSequential', 'ReLU', 'AvgPool2d']\n",
      "Group 3: ['Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The customization works as expected. You can add more complicated head in here.\n",
    "\n",
    "Another way is to modify fastai's `create_cnn_model` so that it takes a `nf` from input, rather than making the body `Sequential` just for calculating the number of features. This can be an improvement for fastai."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PNASNet-5-Large"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One more example, which is just the same as `NASNet-A-Mobile`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ptcv_get_model(\"pnasnet5large\", pretrained=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pytorchcv.models.common.DualPathSequential"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.features.__class__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0) NASNetInitBlock: 2   layers (total: 2)\n",
      "(1) Stem1Unit   : 59  layers (total: 61)\n",
      "(2) DualPathSequential: 296 layers (total: 357)\n",
      "(3) DualPathSequential: 243 layers (total: 600)\n",
      "(4) DualPathSequential: 235 layers (total: 835)\n",
      "(5) ReLU        : 1   layers (total: 836)\n",
      "(6) AvgPool2d   : 1   layers (total: 837)\n"
     ]
    }
   ],
   "source": [
    "arch_summary(lambda _: m.features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------(0)---------\n",
      "(0) Conv2d      : 1   layers (total: 1)\n",
      "(1) BatchNorm2d : 1   layers (total: 2)\n",
      "---------(1)---------\n",
      "(0) NasConv     : 3   layers (total: 3)\n",
      "(1) DwsBranch   : 8   layers (total: 11)\n",
      "(2) PnasMaxPathBlock: 3   layers (total: 14)\n",
      "(3) DwsBranch   : 8   layers (total: 22)\n",
      "(4) PnasMaxPoolBlock: 1   layers (total: 23)\n",
      "(5) DwsBranch   : 8   layers (total: 31)\n",
      "(6) DwsBranch   : 8   layers (total: 39)\n",
      "(7) DwsBranch   : 8   layers (total: 47)\n",
      "(8) PnasMaxPoolBlock: 1   layers (total: 48)\n",
      "(9) DwsBranch   : 8   layers (total: 56)\n",
      "(10) NasConv     : 3   layers (total: 59)\n",
      "---------(2)---------\n",
      "(0) PnasUnit    : 64  layers (total: 64)\n",
      "(1) PnasUnit    : 61  layers (total: 125)\n",
      "(2) PnasUnit    : 57  layers (total: 182)\n",
      "(3) PnasUnit    : 57  layers (total: 239)\n",
      "(4) PnasUnit    : 57  layers (total: 296)\n",
      "---------(3)---------\n",
      "(0) PnasUnit    : 68  layers (total: 68)\n",
      "(1) PnasUnit    : 61  layers (total: 129)\n",
      "(2) PnasUnit    : 57  layers (total: 186)\n",
      "(3) PnasUnit    : 57  layers (total: 243)\n",
      "---------(4)---------\n",
      "(0) PnasUnit    : 60  layers (total: 60)\n",
      "(1) PnasUnit    : 61  layers (total: 121)\n",
      "(2) PnasUnit    : 57  layers (total: 178)\n",
      "(3) PnasUnit    : 57  layers (total: 235)\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    print(f'---------({i})---------')\n",
    "    arch_summary(lambda _: m.features[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (dropout): Dropout(p=0.5)\n",
       "  (fc): Linear(in_features=4320, out_features=1000, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_size = m.output[1].out_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pnasnet5large(pretrained=False):\n",
    "    return ptcv_get_model(\"pnasnet5large\", pretrained=False).features"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create learner, split at layer (3): "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = FakeData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "learn = cnn_learner(data, pnasnet5large, pretrained=False,\n",
    "                    cut=noop, split_on=lambda m: (m[0][3], m[1]),\n",
    "                    custom_head=nn.Linear(output_size, data.c))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To check the cut and split work as we expected, we extract the groups:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Group 1: ['NASNetInitBlock', 'Stem1Unit', 'DualPathSequential']\n",
      "Group 2: ['DualPathSequential', 'DualPathSequential', 'ReLU', 'AvgPool2d']\n",
      "Group 3: ['Linear']\n"
     ]
    }
   ],
   "source": [
    "get_groups(nn.Sequential(*learn.model[0], learn.model[1]), learn.layer_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The customization works as expected."
   ]
  },
  {
   "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
}