{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Rec 0: Block Processing\n", "\n", "## The concept of a `Layer`\n", "`Layers` are the building blocks of neural networks. They:\n", "- take input tensors\n", "- apply some computation\n", "- return output tensors\n", "- have learnable parameters\n", "\n", "## The Concept of a `Block`:\n", "A neural network block:\n", "- Can represent a single layer, a set of layers, or the entire model.\n", "- These blocks can be combined recursively into larger, more complex structures.\n", "- This modular approach allows for compact code while enabling the implementation of complex networks.\n", "\n", "Motivation:\n", "- For large models with many layers, it's easier to manage and understand if we break them down into smaller, more manageable blocks.\n", "- This approach allows us to focus on the core functionality of each block, rather than the entire model.\n", "- It also makes the code more modular and easier to modify.\n", "\n", "# ![Neural Network Blocks](https://classic.d2l.ai/_images/blocks.svg \"Multiple layers are combined into blocks, forming repeating patterns of larger models.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Imports" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torchinfo import summary" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# `nn.Module`: The main building block!\n", "- Base class for all neural network modules. Your models MUST also subclass this class. \n", "- Modules can also contain other Modules, allowing to nest them in a tree structure. \n", "- You can assign the submodules as regular attributes. Conv2d, Linear, BatchNorm2d all subclass nn.Module. \n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Example: Simple CNN Classifier\n", "The following is a simple CNN classifier with two convolutional layers and two fully connected layers. Dont worry about the details of the layers, you will learm about them through the course. \n", "\n", "Problem? \n", "- If we want to add a layer we have to again write lots of code in the __init__ and in the forward function. \n", "- 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.\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Conv2d: 1-1 [1, 32, 28, 28] 320\n", "├─BatchNorm2d: 1-2 [1, 32, 28, 28] 64\n", "├─Conv2d: 1-3 [1, 64, 28, 28] 18,496\n", "├─BatchNorm2d: 1-4 [1, 64, 28, 28] 128\n", "├─Linear: 1-5 [1, 1024] 51,381,248\n", "├─Linear: 1-6 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 51,410,506\n", "Trainable params: 51,410,506\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.14\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.21\n", "Params size (MB): 205.64\n", "Estimated Total Size (MB): 206.86\n", "==========================================================================================" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleCNNClassifier(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) # 3x3 convolution\n", " x = self.bn1(x) # batch normalization\n", " x = F.relu(x) # ReLU activation\n", " \n", " x = self.conv2(x) # 3x3 convolution\n", " x = self.bn2(x) # batch normalization\n", " x = F.relu(x) # ReLU activation\n", "\n", " x = x.view(x.size(0), -1) # flatten the tensor \n", " \n", " x = self.fc1(x) # fully connected layer\n", " x = F.sigmoid(x) # sigmoid activation\n", " x = self.fc2(x) # fully connected layer\n", " \n", " return x\n", " \n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## `nn.Sequential`: Stack and merge layers!\n", "- A sequential container. Modules will be added to it in the order they are passed in the constructor. \n", "- The data is passed through all the modules in the same order as they are added to the `Sequential` container. \n", "- You can use `nn.Sequential` to create a simple neural networks. \n", "- You can also use `nn.Sequential` to create a more complex network by adding multiple instances of it. \n", "\n" ] }, { "cell_type": "code", "execution_count": 26, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 32, 28, 28] --\n", "│ └─Conv2d: 2-1 [1, 32, 28, 28] 320\n", "│ └─BatchNorm2d: 2-2 [1, 32, 28, 28] 64\n", "│ └─ReLU: 2-3 [1, 32, 28, 28] --\n", "├─Sequential: 1-2 [1, 64, 28, 28] --\n", "│ └─Conv2d: 2-4 [1, 64, 28, 28] 18,496\n", "│ └─BatchNorm2d: 2-5 [1, 64, 28, 28] 128\n", "│ └─ReLU: 2-6 [1, 64, 28, 28] --\n", "├─Sequential: 1-3 [1, 10] --\n", "│ └─Linear: 2-7 [1, 1024] 51,381,248\n", "│ └─Sigmoid: 2-8 [1, 1024] --\n", "│ └─Linear: 2-9 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 51,410,506\n", "Trainable params: 51,410,506\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.14\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.21\n", "Params size (MB): 205.64\n", "Estimated Total Size (MB): 206.86\n", "==========================================================================================" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleCNNClassifier(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", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " return x\n", " \n", "\n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Much better! But still not perfect. Did you notice that self.conv_block1 and self.conv_block2 look exactly the same? Can we do better?\n", "- Yes! We could create a function that reteurns a nn.Sequential to even simplify the code! Then we can just call this function in our Module\n" ] }, { "cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 32, 28, 28] --\n", "│ └─Conv2d: 2-1 [1, 32, 28, 28] 320\n", "│ └─BatchNorm2d: 2-2 [1, 32, 28, 28] 64\n", "│ └─ReLU: 2-3 [1, 32, 28, 28] --\n", "├─Sequential: 1-2 [1, 64, 28, 28] --\n", "│ └─Conv2d: 2-4 [1, 64, 28, 28] 18,496\n", "│ └─BatchNorm2d: 2-5 [1, 64, 28, 28] 128\n", "│ └─ReLU: 2-6 [1, 64, 28, 28] --\n", "├─Sequential: 1-3 [1, 10] --\n", "│ └─Linear: 2-7 [1, 1024] 51,381,248\n", "│ └─Sigmoid: 2-8 [1, 1024] --\n", "│ └─Linear: 2-9 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 51,410,506\n", "Trainable params: 51,410,506\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.14\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.21\n", "Params size (MB): 205.64\n", "Estimated Total Size (MB): 206.86\n", "==========================================================================================" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "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", " )\n", "\n", "class SimpleCNNClassifier(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", " self.decoder = nn.Sequential(\n", " nn.Linear(64 * 28 * 28, 1024),\n", " nn.Sigmoid(),\n", " nn.Linear(1024, n_classes)\n", " )\n", "\n", " def forward(self, x):\n", " x = self.conv_block1(x)\n", " x = self.conv_block2(x)\n", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " return x\n", " \n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Even cleaner! Still conv_block1 and conv_block2 are almost the same! We can do away with the conv_block function and merge them using nn.Sequential" ] }, { "cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 64, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 32, 28, 28] --\n", "│ │ └─Conv2d: 3-1 [1, 32, 28, 28] 320\n", "│ │ └─BatchNorm2d: 3-2 [1, 32, 28, 28] 64\n", "│ │ └─ReLU: 3-3 [1, 32, 28, 28] --\n", "│ └─Sequential: 2-2 [1, 64, 28, 28] --\n", "│ │ └─Conv2d: 3-4 [1, 64, 28, 28] 18,496\n", "│ │ └─BatchNorm2d: 3-5 [1, 64, 28, 28] 128\n", "│ │ └─ReLU: 3-6 [1, 64, 28, 28] --\n", "├─Sequential: 1-2 [1, 10] --\n", "│ └─Linear: 2-3 [1, 1024] 51,381,248\n", "│ └─Sigmoid: 2-4 [1, 1024] --\n", "│ └─Linear: 2-5 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 51,410,506\n", "Trainable params: 51,410,506\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.14\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.21\n", "Params size (MB): 205.64\n", "Estimated Total Size (MB): 206.86\n", "==========================================================================================" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleCNNClassifier(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", " 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", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " return x\n", "\n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Dynamic Sequential: Create multiple layers at once!\n", "What if we want more layers in our `self.encoder`? Hardcoding is not convenient.\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", "```\n", "\n", "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`. \n", "\n" ] }, { "cell_type": "code", "execution_count": 29, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 64, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 32, 28, 28] --\n", "│ │ └─Conv2d: 3-1 [1, 32, 28, 28] 320\n", "│ │ └─BatchNorm2d: 3-2 [1, 32, 28, 28] 64\n", "│ │ └─ReLU: 3-3 [1, 32, 28, 28] --\n", "│ └─Sequential: 2-2 [1, 64, 28, 28] --\n", "│ │ └─Conv2d: 3-4 [1, 64, 28, 28] 18,496\n", "│ │ └─BatchNorm2d: 3-5 [1, 64, 28, 28] 128\n", "│ │ └─ReLU: 3-6 [1, 64, 28, 28] --\n", "├─Sequential: 1-2 [1, 10] --\n", "│ └─Linear: 2-3 [1, 1024] 51,381,248\n", "│ └─Sigmoid: 2-4 [1, 1024] --\n", "│ └─Linear: 2-5 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 51,410,506\n", "Trainable params: 51,410,506\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.14\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.21\n", "Params size (MB): 205.64\n", "Estimated Total Size (MB): 206.86\n", "==========================================================================================" ] }, "execution_count": 29, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleCNNClassifier(nn.Module):\n", " def __init__(self, in_c, n_classes):\n", " super().__init__()\n", " self.enc_sizes = [in_c, 32, 64]\n", " \n", " # zip stops when the shortest list is exhausted\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", " # Sequential cannot take a list of modules, so we need to unpack it\n", " self.encoder = nn.Sequential(*conv_blocks)\n", "\n", " self.decoder = nn.Sequential(\n", " nn.Linear(self.enc_sizes[-1] * 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", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " return x\n", "\n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "code", "execution_count": 68, "metadata": {}, "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": [ "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": 33, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 128, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 32, 28, 28] --\n", "│ │ └─Conv2d: 3-1 [1, 32, 28, 28] 320\n", "│ │ └─BatchNorm2d: 3-2 [1, 32, 28, 28] 64\n", "│ │ └─ReLU: 3-3 [1, 32, 28, 28] --\n", "│ └─Sequential: 2-2 [1, 64, 28, 28] --\n", "│ │ └─Conv2d: 3-4 [1, 64, 28, 28] 18,496\n", "│ │ └─BatchNorm2d: 3-5 [1, 64, 28, 28] 128\n", "│ │ └─ReLU: 3-6 [1, 64, 28, 28] --\n", "│ └─Sequential: 2-3 [1, 128, 28, 28] --\n", "│ │ └─Conv2d: 3-7 [1, 128, 28, 28] 73,856\n", "│ │ └─BatchNorm2d: 3-8 [1, 128, 28, 28] 256\n", "│ │ └─ReLU: 3-9 [1, 128, 28, 28] --\n", "├─Sequential: 1-2 [1, 10] --\n", "│ └─Linear: 2-4 [1, 1024] 102,761,472\n", "│ └─Sigmoid: 2-5 [1, 1024] --\n", "│ └─Linear: 2-6 [1, 10] 10,250\n", "==========================================================================================\n", "Total params: 102,864,842\n", "Trainable params: 102,864,842\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 175.43\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 2.82\n", "Params size (MB): 411.46\n", "Estimated Total Size (MB): 414.28\n", "==========================================================================================" ] }, "execution_count": 33, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleCNNClassifier(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(self.enc_sizes[-1] * 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", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " return x\n", "\n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, enc_sizes=[32, 64, 128], n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can do the same for the decoder part!" ] }, { "cell_type": "code", "execution_count": 34, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─Sequential: 1-1 [1, 128, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 32, 28, 28] --\n", "│ │ └─Conv2d: 3-1 [1, 32, 28, 28] 320\n", "│ │ └─BatchNorm2d: 3-2 [1, 32, 28, 28] 64\n", "│ │ └─ReLU: 3-3 [1, 32, 28, 28] --\n", "│ └─Sequential: 2-2 [1, 64, 28, 28] --\n", "│ │ └─Conv2d: 3-4 [1, 64, 28, 28] 18,496\n", "│ │ └─BatchNorm2d: 3-5 [1, 64, 28, 28] 128\n", "│ │ └─ReLU: 3-6 [1, 64, 28, 28] --\n", "│ └─Sequential: 2-3 [1, 128, 28, 28] --\n", "│ │ └─Conv2d: 3-7 [1, 128, 28, 28] 73,856\n", "│ │ └─BatchNorm2d: 3-8 [1, 128, 28, 28] 256\n", "│ │ └─ReLU: 3-9 [1, 128, 28, 28] --\n", "├─Sequential: 1-2 [1, 512] --\n", "│ └─Sequential: 2-4 [1, 1024] --\n", "│ │ └─Linear: 3-10 [1, 1024] 102,761,472\n", "│ │ └─Sigmoid: 3-11 [1, 1024] --\n", "│ └─Sequential: 2-5 [1, 512] --\n", "│ │ └─Linear: 3-12 [1, 512] 524,800\n", "│ │ └─Sigmoid: 3-13 [1, 512] --\n", "├─Linear: 1-3 [1, 10] 5,130\n", "==========================================================================================\n", "Total params: 103,384,522\n", "Trainable params: 103,384,522\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 175.95\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 2.82\n", "Params size (MB): 413.54\n", "Estimated Total Size (MB): 416.36\n", "==========================================================================================" ] }, "execution_count": 34, "metadata": {}, "output_type": "execute_result" } ], "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 SimpleCNNClassifier(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 = [enc_sizes[-1] * 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", " 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", " self.last = nn.Linear(self.dec_sizes[-1], n_classes)\n", "\n", " \n", " def forward(self, x):\n", " x = self.encoder(x)\n", " x = x.view(x.size(0), -1) # flat\n", " x = self.decoder(x)\n", " x = self.last(x)\n", " return x\n", "\n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, enc_sizes=[32, 64, 128], dec_sizes=[1024, 512], n_classes=10), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can make things even more modular by splitting the encoder and decoder into separate classes!" ] }, { "cell_type": "code", "execution_count": 38, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "SimpleCNNClassifier [1, 10] --\n", "├─SimpleEncoder: 1-1 [1, 128, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 128, 28, 28] --\n", "│ │ └─Sequential: 3-1 [1, 32, 28, 28] 384\n", "│ │ └─Sequential: 3-2 [1, 64, 28, 28] 18,624\n", "│ │ └─Sequential: 3-3 [1, 128, 28, 28] 74,112\n", "├─SimpleDecoder: 1-2 [1, 10] --\n", "│ └─Sequential: 2-2 [1, 512] --\n", "│ │ └─Sequential: 3-4 [1, 1024] 102,761,472\n", "│ │ └─Sequential: 3-5 [1, 512] 524,800\n", "│ └─Linear: 2-3 [1, 10] 5,130\n", "==========================================================================================\n", "Total params: 103,384,522\n", "Trainable params: 103,384,522\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 175.95\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 2.82\n", "Params size (MB): 413.54\n", "Estimated Total Size (MB): 416.36\n", "==========================================================================================" ] }, "execution_count": 38, "metadata": {}, "output_type": "execute_result" } ], "source": [ "class SimpleEncoder(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 SimpleDecoder(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", " x = self.dec_blocks(x)\n", " x = self.last(x)\n", " return x\n", " \n", " \n", "class SimpleCNNClassifier(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 = SimpleEncoder(self.enc_sizes)\n", " self.decoder = SimpleDecoder(self.dec_sizes, n_classes)\n", " \n", " def forward(self, x):\n", " x = self.encoder(x)\n", " x = x.flatten(1) # flat\n", " x = self.decoder(x)\n", " return x\n", " \n", "# Summary\n", "summary(SimpleCNNClassifier(in_c=1, enc_sizes=[32, 64, 128], dec_sizes=[1024, 512], n_classes=10), input_size=(1, 1, 28, 28))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "By diving our module into submodules it is easier to share the code, debug it and test it." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# `nn.ModuleList`: When we need to iterate!\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,\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": 44, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "torch.Size([4, 16])\n", "torch.Size([4, 32])\n" ] } ], "source": [ "class SimpleModule(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\n", "\n", "# Summary\n", "summary(SimpleModule([1, 16, 32]), input_size=(1, 1))\n", "\n", "# Example: Use for storing intermediate outputs/debugging \n", "model = SimpleModule([1, 16, 32])\n", "model(torch.rand((4,1)))\n", "\n", "for trace in model.trace:\n", " print(trace.shape)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# `nn.ModuleDict`: When we need to choose!\n", "ModuleDict allows you to store Module as a dictionary. It can be useful when you need to choose a module based on some condition. \n", "What if we want to switch to LeakyReLU 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": 50, "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": [ "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", " )\n", "\n", "# Summary\n", "print(conv_block(1, 32, kernel_size=3, padding=1, activation='lrelu'))\n", "print(conv_block(1, 32, kernel_size=3, padding=1, activation='relu'))\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Putting it all together!" ] }, { "cell_type": "code", "execution_count": 67, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "==========================================================================================\n", "Layer (type:depth-idx) Output Shape Param #\n", "==========================================================================================\n", "MyCNNClassifier [1, 10] --\n", "├─MyEncoder: 1-1 [1, 64, 28, 28] --\n", "│ └─Sequential: 2-1 [1, 64, 28, 28] --\n", "│ │ └─Sequential: 3-1 [1, 32, 28, 28] 384\n", "│ │ └─Sequential: 3-2 [1, 64, 28, 28] 18,624\n", "├─MyDecoder: 1-2 [1, 10] --\n", "│ └─Sequential: 2-2 [1, 512] --\n", "│ │ └─Sequential: 3-3 [1, 1024] 51,381,248\n", "│ │ └─Sequential: 3-4 [1, 512] 524,800\n", "│ └─Linear: 2-3 [1, 10] 5,130\n", "==========================================================================================\n", "Total params: 51,930,186\n", "Trainable params: 51,930,186\n", "Non-trainable params: 0\n", "Total mult-adds (Units.MEGABYTES): 66.66\n", "==========================================================================================\n", "Input size (MB): 0.00\n", "Forward/backward pass size (MB): 1.22\n", "Params size (MB): 207.72\n", "Estimated Total Size (MB): 208.94\n", "==========================================================================================" ] }, "execution_count": 67, "metadata": {}, "output_type": "execute_result" } ], "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", " x = self.dec_blocks(x) \n", " x = self.last(x)\n", " return x\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 = [self.enc_sizes[-1] * 28 * 28, *dec_sizes]\n", "\n", " self.encoder = MyEncoder(self.enc_sizes, activation=activation)\n", " self.decoder = MyDecoder(self.dec_sizes, n_classes)\n", " \n", " def forward(self, x):\n", " x = self.encoder(x)\n", " x = x.flatten(1) # flat\n", " x = self.decoder(x)\n", " return x\n", "\n", "# Summary\n", "summary(MyCNNClassifier(in_c=1, enc_sizes=[32, 64], dec_sizes=[1024, 512], n_classes=10, activation='lrelu'), input_size=(1, 1, 28, 28))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Conclusion\n", "- In summary, we have seen how to create a modular neural network using `nn.Module`, `nn.Sequential`, `nn.ModuleList`, and `nn.ModuleDict`.\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", "For more information, you can check out the [official documentation](https://pytorch.org/docs/stable/nn.html#containers)\n" ] } ], "metadata": { "kernelspec": { "display_name": "base", "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.12.4" } }, "nbformat": 4, "nbformat_minor": 2 }