notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pytorch: how and when to use Module, Sequential, ModuleList and ModuleDict\n",
    "### Effective way to share, reuse and break down the complexity of your models\n",
    "\n",
    "Updated at Pytorch 1.5\n",
    "\n",
    "You can find the code [here](https://github.com/FrancescoSaverioZuppichini/Pytorch-how-and-when-to-use-Module-Sequential-ModuleList-and-ModuleDict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Pytorch](https://pytorch.org/) is an open source deep learning frameworks that provide a smart way to create ML models. Even if the documentation is well made, I still find that most people still are able to write bad and not organized PyTorch code.\n",
    "\n",
    "Today, we are going to see how to use the three main building blocks of PyTorch: `Module, Sequential and ModuleList`. We are going to start with an example and iteratively we will make it better.\n",
    "\n",
    "All these four classes are contained into `torch.nn`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "# nn.Module\n",
    "# nn.Sequential\n",
    "# nn.Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module: the main building block\n",
    "\n",
    "The [Module](https://pytorch.org/docs/stable/nn.html?highlight=module) is the main building block, it defines the base class for all neural network and you MUST subclass it. \n",
    "\n",
    "Let's create a classic CNN classifier as example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, n_classes):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(in_c, 32, kernel_size=3, stride=1, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(32)\n",
    "        \n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(64)\n",
    "\n",
    "        self.fc1 = nn.Linear(64 * 28 * 28, 1024)\n",
    "        self.fc2 = nn.Linear(1024, n_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.fc1(x)\n",
    "        x = F.sigmoid(x)\n",
    "        x = self.fc2(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (fc1): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "  (fc2): Linear(in_features=1024, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a very simple classifier with an encoding part that uses two layers with 3x3 convs + batchnorm + relu and a decoding part with two linear layers. If you are not new to PyTorch you may have seen this type of coding before, but there are two problems.\n",
    "\n",
    "If we want to add a layer we have to again write lots of code in the `__init__` and in the `forward` function. Also, if we have some common block that we want to use in another model, e.g. the 3x3 conv + batchnorm + relu, we have to write it again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sequential: stack and merge layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Sequential](https://pytorch.org/docs/stable/nn.html?highlight=sequential#torch.nn.Sequential) is a container of Modules that can be stacked together and run at the same time.\n",
    "\n",
    "You can notice that we have to store into `self` everything. We can use `Sequential` to improve our code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, n_classes):\n",
    "        super().__init__()\n",
    "        self.conv_block1 = nn.Sequential(\n",
    "            nn.Conv2d(in_c, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        self.conv_block2 = nn.Sequential(\n",
    "            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(64 * 28 * 28, 1024),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(1024, n_classes)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv_block1(x)\n",
    "        x = self.conv_block2(x)\n",
    "\n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (conv_block1): Sequential(\n",
      "    (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "  )\n",
      "  (conv_block2): Sequential(\n",
      "    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "    (1): Sigmoid()\n",
      "    (2): Linear(in_features=1024, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much Better uhu?\n",
    "\n",
    "Did you notice that `conv_block1` and `conv_block2` looks almost the same? We could create a function that reteurns a `nn.Sequential` to even simplify the code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_block(in_f, out_f, *args, **kwargs):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_f, out_f, *args, **kwargs),\n",
    "        nn.BatchNorm2d(out_f),\n",
    "        nn.ReLU()\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can just call this function in our Module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, n_classes):\n",
    "        super().__init__()\n",
    "        self.conv_block1 = conv_block(in_c, 32, kernel_size=3, padding=1)\n",
    "        \n",
    "        self.conv_block2 = conv_block(32, 64, kernel_size=3, padding=1)\n",
    "\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(64 * 28 * 28, 1024),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(1024, n_classes)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv_block1(x)\n",
    "        x = self.conv_block2(x)\n",
    "\n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (conv_block1): Sequential(\n",
      "    (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "  )\n",
      "  (conv_block2): Sequential(\n",
      "    (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (2): ReLU()\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "    (1): Sigmoid()\n",
      "    (2): Linear(in_features=1024, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even cleaner! Still `conv_block1` and `conv_block2` are almost the same! We can merge them using `nn.Sequential`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, n_classes):\n",
    "        super().__init__()\n",
    "        self.encoder = nn.Sequential(\n",
    "            conv_block(in_c, 32, kernel_size=3, padding=1),\n",
    "            conv_block(32, 64, kernel_size=3, padding=1)\n",
    "        )\n",
    "\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(64 * 28 * 28, 1024),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(1024, n_classes)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "    (1): Sigmoid()\n",
      "    (2): Linear(in_features=1024, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`self.encoder` now holds booth `conv_block`. We have decoupled logic for our model and make it easier to read and reuse. Our `conv_block` function can be imported and used in another model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamic Sequential: create multiple layers at once\n",
    "\n",
    "What if we can to add a new layers in `self.encoder`, hardcoded them is not convinient:\n",
    "\n",
    "```python\n",
    "self.encoder = nn.Sequential(\n",
    "            conv_block(in_c, 32, kernel_size=3, padding=1),\n",
    "            conv_block(32, 64, kernel_size=3, padding=1),\n",
    "            conv_block(64, 128, kernel_size=3, padding=1),\n",
    "            conv_block(128, 256, kernel_size=3, padding=1),\n",
    "\n",
    "        )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Would it be nice if we can define the sizes as an array and automatically create all the layers without writing each one of them? Fortunately we can create an array and pass it to `Sequential`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, n_classes):\n",
    "        super().__init__()\n",
    "        self.enc_sizes = [in_c, 32, 64]\n",
    "        \n",
    "        conv_blocks = [conv_block(in_f, out_f, kernel_size=3, padding=1) \n",
    "                       for in_f, out_f in zip(self.enc_sizes, self.enc_sizes[1:])]\n",
    "        \n",
    "        self.encoder = nn.Sequential(*conv_blocks)\n",
    "\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(64 * 28 * 28, 1024),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(1024, n_classes)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "    (1): Sigmoid()\n",
      "    (2): Linear(in_features=1024, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's break it down. We created an array `self.enc_sizes` that holds the sizes of our encoder. Then we create an array `conv_blocks` by iterating the sizes. Since we have to give booth a in size and an outsize for each layer we `zip`ed the size'array with itself by shifting it by one. \n",
    "\n",
    "Just to be clear, take a look at the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 32\n",
      "32 64\n"
     ]
    }
   ],
   "source": [
    "sizes = [1, 32, 64]\n",
    "\n",
    "for in_f,out_f in zip(sizes, sizes[1:]):\n",
    "    print(in_f,out_f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, since `Sequential` does not accept a list, we decompose it by using the `*` operator.\n",
    "\n",
    "Tada! Now if we just want to add a size, we can easily add a new number to the list. It is a common practice to make the size a parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, enc_sizes, n_classes):\n",
    "        super().__init__()\n",
    "        self.enc_sizes = [in_c, *enc_sizes]\n",
    "        \n",
    "        conv_blocks = [conv_block(in_f, out_f, kernel_size=3, padding=1) \n",
    "                       for in_f, out_f in zip(self.enc_sizes, self.enc_sizes[1:])]\n",
    "        \n",
    "        self.encoder = nn.Sequential(*conv_blocks)\n",
    "\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(64 * 28 * 28, 1024),\n",
    "            nn.Sigmoid(),\n",
    "            nn.Linear(1024, n_classes)\n",
    "        )\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "    (2): Sequential(\n",
      "      (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "    (1): Sigmoid()\n",
      "    (2): Linear(in_features=1024, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, [32,64, 128], 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do the same for the decoder part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dec_block(in_f, out_f):\n",
    "    return nn.Sequential(\n",
    "        nn.Linear(in_f, out_f),\n",
    "        nn.Sigmoid()\n",
    "    )\n",
    "\n",
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, enc_sizes, dec_sizes,  n_classes):\n",
    "        super().__init__()\n",
    "        self.enc_sizes = [in_c, *enc_sizes]\n",
    "        self.dec_sizes = [64 * 28 * 28, *dec_sizes]\n",
    "\n",
    "        conv_blocks = [conv_block(in_f, out_f, kernel_size=3, padding=1) \n",
    "                       for in_f, out_f in zip(self.enc_sizes, self.enc_sizes[1:])]\n",
    "        \n",
    "        self.encoder = nn.Sequential(*conv_blocks)\n",
    "\n",
    "        \n",
    "        dec_blocks = [dec_block(in_f, out_f) \n",
    "                       for in_f, out_f in zip(self.dec_sizes, self.dec_sizes[1:])]\n",
    "        \n",
    "        self.decoder = nn.Sequential(*dec_blocks)\n",
    "        \n",
    "        self.last = nn.Linear(self.dec_sizes[-1], n_classes)\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.view(x.size(0), -1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (2): ReLU()\n",
      "    )\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Sequential(\n",
      "      (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "      (1): Sigmoid()\n",
      "    )\n",
      "    (1): Sequential(\n",
      "      (0): Linear(in_features=1024, out_features=512, bias=True)\n",
      "      (1): Sigmoid()\n",
      "    )\n",
      "  )\n",
      "  (last): Linear(in_features=512, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, [32,64], [1024, 512], 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We followed the same pattern, we create a new block for the decoding part, linear + sigmoid, and we pass an array with the sizes. We had to add a `self.last` since we do not want to activate the output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can even break down our model in two! Encoder + Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyEncoder(nn.Module):\n",
    "    def __init__(self, enc_sizes):\n",
    "        super().__init__()\n",
    "        self.conv_blocks = nn.Sequential(*[conv_block(in_f, out_f, kernel_size=3, padding=1) \n",
    "                       for in_f, out_f in zip(enc_sizes, enc_sizes[1:])])\n",
    "\n",
    "        def forward(self, x):\n",
    "            return self.conv_blocks(x)\n",
    "        \n",
    "class MyDecoder(nn.Module):\n",
    "    def __init__(self, dec_sizes, n_classes):\n",
    "        super().__init__()\n",
    "        self.dec_blocks = nn.Sequential(*[dec_block(in_f, out_f) \n",
    "                       for in_f, out_f in zip(dec_sizes, dec_sizes[1:])])\n",
    "        self.last = nn.Linear(dec_sizes[-1], n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.dec_blocks()\n",
    "    \n",
    "    \n",
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, enc_sizes, dec_sizes,  n_classes):\n",
    "        super().__init__()\n",
    "        self.enc_sizes = [in_c, *enc_sizes]\n",
    "        self.dec_sizes = [self.enc_sizes[-1] * 28 * 28, *dec_sizes]\n",
    "\n",
    "        self.encoder = MyEncoder(self.enc_sizes)\n",
    "        \n",
    "        self.decoder = MyDecoder(self.dec_sizes, n_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.flatten(1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): MyEncoder(\n",
      "    (conv_blocks): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (2): ReLU()\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (2): ReLU()\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (decoder): MyDecoder(\n",
      "    (dec_blocks): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Linear(in_features=50176, out_features=1024, bias=True)\n",
      "        (1): Sigmoid()\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (1): Sigmoid()\n",
      "      )\n",
      "    )\n",
      "    (last): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, [32,64], [1024, 512], 10)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Be aware that `MyEncoder` and `MyDecoder` could also be functions that returns a `nn.Sequential`. I prefer to use the first pattern for models and the second for building blocks.\n",
    "\n",
    "By diving our module into submodules it is easier to **share** the code, **debug** it and **test** it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModuleList : when we need to iterate\n",
    "\n",
    "`ModuleList` allows you to store `Module` as a list. It can be useful when you need to iterate through layer and store/use some information, like in U-net.\n",
    "\n",
    "The main difference between `Sequential` is that `ModuleList` have not a `forward` method so the inner layers are not connected. Assuming we need each output of each layer in the decoder, we can store it by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModule(nn.Module):\n",
    "    def __init__(self, sizes):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([nn.Linear(in_f, out_f) for in_f, out_f in zip(sizes, sizes[1:])])\n",
    "        self.trace = []\n",
    "        \n",
    "    def forward(self,x):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "            self.trace.append(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([4, 16])\n",
      "torch.Size([4, 32])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[None, None]"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MyModule([1, 16, 32])\n",
    "import torch\n",
    "\n",
    "model(torch.rand((4,1)))\n",
    "\n",
    "[print(trace.shape) for trace in model.trace]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ModuleDict: when we need to choose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What if we want to switch to `LearkyRelu` in our `conv_block`? We can use `ModuleDict` to create a dictionary of `Module` and dynamically switch `Module` when we want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_block(in_f, out_f, activation='relu', *args, **kwargs):\n",
    "    \n",
    "    activations = nn.ModuleDict([\n",
    "                ['lrelu', nn.LeakyReLU()],\n",
    "                ['relu', nn.ReLU()]\n",
    "    ])\n",
    "    \n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_f, out_f, *args, **kwargs),\n",
    "        nn.BatchNorm2d(out_f),\n",
    "        activations[activation]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequential(\n",
      "  (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (2): LeakyReLU(negative_slope=0.01)\n",
      ")\n",
      "Sequential(\n",
      "  (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (2): ReLU()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(conv_block(1, 32,'lrelu', kernel_size=3, padding=1))\n",
    "print(conv_block(1, 32,'relu', kernel_size=3, padding=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's wrap it up everything!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_block(in_f, out_f, activation='relu', *args, **kwargs):\n",
    "    activations = nn.ModuleDict([\n",
    "                ['lrelu', nn.LeakyReLU()],\n",
    "                ['relu', nn.ReLU()]\n",
    "    ])\n",
    "    \n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_f, out_f, *args, **kwargs),\n",
    "        nn.BatchNorm2d(out_f),\n",
    "        activations[activation]\n",
    "    )\n",
    "\n",
    "def dec_block(in_f, out_f):\n",
    "    return nn.Sequential(\n",
    "        nn.Linear(in_f, out_f),\n",
    "        nn.Sigmoid()\n",
    "    )\n",
    "\n",
    "class MyEncoder(nn.Module):\n",
    "    def __init__(self, enc_sizes, *args, **kwargs):\n",
    "        super().__init__()\n",
    "        self.conv_blocks = nn.Sequential(*[conv_block(in_f, out_f, kernel_size=3, padding=1, *args, **kwargs) \n",
    "                       for in_f, out_f in zip(enc_sizes, enc_sizes[1:])])\n",
    "        \n",
    "        def forward(self, x):\n",
    "            return self.conv_blocks(x)\n",
    "        \n",
    "class MyDecoder(nn.Module):\n",
    "    def __init__(self, dec_sizes, n_classes):\n",
    "        super().__init__()\n",
    "        self.dec_blocks = nn.Sequential(*[dec_block(in_f, out_f) \n",
    "                       for in_f, out_f in zip(dec_sizes, dec_sizes[1:])])\n",
    "        self.last = nn.Linear(dec_sizes[-1], n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.dec_blocks()\n",
    "    \n",
    "    \n",
    "class MyCNNClassifier(nn.Module):\n",
    "    def __init__(self, in_c, enc_sizes, dec_sizes,  n_classes, activation='relu'):\n",
    "        super().__init__()\n",
    "        self.enc_sizes = [in_c, *enc_sizes]\n",
    "        self.dec_sizes = [32 * 28 * 28, *dec_sizes]\n",
    "\n",
    "        self.encoder = MyEncoder(self.enc_sizes, activation=activation)\n",
    "        \n",
    "        self.decoder = MyDecoder(dec_sizes, n_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "        \n",
    "        x = x.flatten(1) # flat\n",
    "        \n",
    "        x = self.decoder(x)\n",
    "        \n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MyCNNClassifier(\n",
      "  (encoder): MyEncoder(\n",
      "    (conv_blocks): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (2): LeakyReLU(negative_slope=0.01)\n",
      "      )\n",
      "      (1): Sequential(\n",
      "        (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        (2): LeakyReLU(negative_slope=0.01)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (decoder): MyDecoder(\n",
      "    (dec_blocks): Sequential(\n",
      "      (0): Sequential(\n",
      "        (0): Linear(in_features=1024, out_features=512, bias=True)\n",
      "        (1): Sigmoid()\n",
      "      )\n",
      "    )\n",
      "    (last): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "model = MyCNNClassifier(1, [32,64], [1024, 512], 10, activation='lrelu')\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "So, in summary.\n",
    "\n",
    "- Use `Module` when you have a big block compose of multiple smaller blocks\n",
    "- Use `Sequential` when you want to create a small block from layers\n",
    "- Use `ModuleList` when you need to iterate through some layers or building blocks and do something\n",
    "- Use `ModuleDict` when you need to parametise some blocks of your model, for example an activation function\n",
    "\n",
    "That's all folks!\n",
    "\n",
    "Thank you for reading"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}