code/utils/model-creation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** <amirhosseinheydari78@gmail.com> - 📍 **Origin:** [mr-pylin/pytorch-workshop](https://github.com/mr-pylin/pytorch-workshop)\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Table of contents**<a id='toc0_'></a>    \n",
    "- [Dependencies](#toc1_)    \n",
    "- [Model Creation](#toc2_)    \n",
    "  - [Built-in Models](#toc2_1_)    \n",
    "    - [Torchvision Models](#toc2_1_1_)    \n",
    "    - [Torchaudio Models](#toc2_1_2_)    \n",
    "  - [Hugging Face](#toc2_2_)    \n",
    "  - [Custom Models](#toc2_3_)    \n",
    "    - [Generate Artificial Data](#toc2_3_1_)    \n",
    "    - [Sequential Model](#toc2_3_2_)    \n",
    "      - [Example 1: Using `nn.Sequential`](#toc2_3_2_1_)    \n",
    "      - [Example 2: Using `nn.ModuleList`](#toc2_3_2_2_)    \n",
    "      - [Example 3: Mix of `nn.Sequential` and `nn.ModuleList`](#toc2_3_2_3_)    \n",
    "    - [Non-Sequential (Functional) Model](#toc2_3_3_)    \n",
    "      - [`torch.Tensor` vs. `torch.nn.Parameter`](#toc2_3_3_1_)    \n",
    "        - [`torch.Tensor`](#toc2_3_3_1_1_)    \n",
    "        - [`torch.nn.Parameter`](#toc2_3_3_1_2_)    \n",
    "      - [Example 1: Using `nn.Linear`, `nn.Conv2d`, ...](#toc2_3_3_2_)    \n",
    "      - [Example 2: Mix of Sequential and Non-sequential methods](#toc2_3_3_3_)    \n",
    "      - [Example 3: Separate Class for each Module](#toc2_3_3_4_)    \n",
    "\n",
    "<!-- vscode-jupyter-toc-config\n",
    "\tnumbering=false\n",
    "\tanchor=true\n",
    "\tflat=false\n",
    "\tminLevel=1\n",
    "\tmaxLevel=6\n",
    "\t/vscode-jupyter-toc-config -->\n",
    "<!-- THIS CELL WILL BE REPLACED ON TOC UPDATE. DO NOT WRITE YOUR TEXT IN THIS CELL -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc1_'></a>[Dependencies](#toc0_)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torchaudio.models import hubert_base, wav2vec2_base\n",
    "from torchinfo import summary\n",
    "from torchvision import models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a seed for deterministic results\n",
    "seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id='toc2_'></a>[Model Creation](#toc0_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_1_'></a>[Built-in Models](#toc0_)\n",
    "\n",
    "- PyTorch provides a variety of **pre-trained models** for different tasks, including **image classification**, **object detection**, **segmentation**, and **audio processing**.\n",
    "- These models are available directly in libraries like `torchvision`, `torchaudio`, and `torchtext`, making it easier to leverage state-of-the-art architectures.\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- Torchvision Models: [pytorch.org/vision/stable/models.html](https://pytorch.org/vision/stable/models.html)\n",
    "- Torchaudio Models: [pytorch.org/audio/stable/models.html](https://pytorch.org/audio/stable/models.html)\n",
    "- Torchtext Models: [pytorch.org/text/stable/models.html](https://pytorch.org/text/stable/models.html)\n",
    "- Check Manual Implementations with details in [**../models/**](../models/) directory.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_1_1_'></a>[Torchvision Models](#toc0_)\n",
    "\n",
    "- This is a **subset** of available pre-trained models in `torchvision`.\n",
    "\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Task</th>\n",
    "      <th>Model</th>\n",
    "      <th>Type</th>\n",
    "      <th>Input Format</th>\n",
    "      <th>Import Path</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td rowspan=\"6\">Image Classification</td>\n",
    "      <td>ResNet-18, ResNet-50</td>\n",
    "      <td>Residual Network</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>VGG-16, VGG-19</td>\n",
    "      <td>CNN</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>DenseNet-121, DenseNet-161</td>\n",
    "      <td>Dense Network</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>MobileNetV2, MobileNetV3</td>\n",
    "      <td>Lightweight CNN</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Inception v3</td>\n",
    "      <td>Inception Network</td>\n",
    "      <td>RGB images (299x299)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>EfficientNet</td>\n",
    "      <td>Efficient Network</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td rowspan=\"4\">Object Detection & Segmentation</td>\n",
    "      <td>Faster R-CNN (ResNet-50)</td>\n",
    "      <td>Region Proposal Network</td>\n",
    "      <td>RGB images (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.detection</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Mask R-CNN (ResNet-50)</td>\n",
    "      <td>Instance Segmentation</td>\n",
    "      <td>RGB images (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.detection</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>RetinaNet</td>\n",
    "      <td>Single-stage Object Detection</td>\n",
    "      <td>RGB images (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.detection</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Keypoint R-CNN</td>\n",
    "      <td>Keypoint Detection</td>\n",
    "      <td>RGB images (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.detection</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td rowspan=\"2\">Semantic Segmentation</td>\n",
    "      <td>DeepLabV3</td>\n",
    "      <td>Atrous Convolution Network</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.segmentation</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>FCN (Fully Convolutional Network)</td>\n",
    "      <td>Fully Convolutional Network</td>\n",
    "      <td>RGB images (224x224)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.segmentation</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td rowspan=\"2\">Video Classification</td>\n",
    "      <td>ResNet3D</td>\n",
    "      <td>3D Convolution Network</td>\n",
    "      <td>Video (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.video</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Swin Transformer 3D</td>\n",
    "      <td>Transformer-based 3D Video Classification</td>\n",
    "      <td>Video (varied sizes)</td>\n",
    "      <td style=\"font-family: monospace;\">torchvision.models.video</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resnet50 = models.resnet50(weights=None)\n",
    "\n",
    "# log\n",
    "print(resnet50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fasterrcnn_resnet50 = models.detection.fasterrcnn_resnet50_fpn(weights=None)\n",
    "\n",
    "# log\n",
    "print(fasterrcnn_resnet50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deeplabv3 = models.segmentation.deeplabv3_resnet50(weights=None)\n",
    "\n",
    "# log\n",
    "print(deeplabv3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_1_2_'></a>[Torchaudio Models](#toc0_)\n",
    "\n",
    "- This is a **subset** of available pre-trained models in `torchaudio`.\n",
    "\n",
    "<table style=\"margin:0 auto;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th>Task</th>\n",
    "      <th>Model</th>\n",
    "      <th>Type</th>\n",
    "      <th>Input Format</th>\n",
    "      <th>Import Path</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td rowspan=\"4\">Speech Recognition</td>\n",
    "      <td>Wav2Vec2</td>\n",
    "      <td>Self-Supervised Speech Model</td>\n",
    "      <td>Waveform (1D Tensor)</td>\n",
    "      <td style=\"font-family: monospace;\">torchaudio.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Hubert</td>\n",
    "      <td>Self-Supervised Speech Model</td>\n",
    "      <td>Waveform (1D Tensor)</td>\n",
    "      <td style=\"font-family: monospace;\">torchaudio.models</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>DeepSpeech</td>\n",
    "      <td>End-to-End Speech Recognition</td>\n",
    "      <td>Waveform (1D Tensor)</td>\n",
    "      <td style=\"font-family: monospace;\">torchaudio.models.deepspeech</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>Conformer</td>\n",
    "      <td>Convolution-augmented Transformer</td>\n",
    "      <td>Waveform (1D Tensor)</td>\n",
    "      <td style=\"font-family: monospace;\">torchaudio.models.conformer</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wav2vec2 = wav2vec2_base()\n",
    "\n",
    "# log\n",
    "print(wav2vec2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hubert = hubert_base()\n",
    "\n",
    "# log\n",
    "print(hubert)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_2_'></a>[Hugging Face](#toc0_)\n",
    "\n",
    "- Hugging Face offers a vast collection of models and pretrained weights.\n",
    "- It is renowned for state-of-the-art models in NLP, computer vision, and more.\n",
    "- Models include transformers, BERT, GPT, and many others (pretrained on large datasets).\n",
    "- They can be fine-tuned for specific tasks using regular pytorch code.\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- Documentations: [huggingface.co/docs](https://huggingface.co/docs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## <a id='toc2_3_'></a>[Custom Models](#toc0_)\n",
    "\n",
    "- PyTorch allows you to define **custom** models by extending the `torch.nn.Module` class.\n",
    "- To create a custom model, subclass `torch.nn.Module` and implement the `__init__` and `forward` methods, where `__init__` initializes the layers, and `forward` defines the computation that takes place when the model is called.\n",
    "\n",
    "📝 **Docs**:\n",
    "\n",
    "- Building Models with PyTorch: [pytorch.org/tutorials/beginner/introyt/modelsyt_tutorial.html](https://pytorch.org/tutorials/beginner/introyt/modelsyt_tutorial.html)\n",
    "- Build the Neural Network: [pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html](https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html)\n",
    "- Neural Networks: [pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial](https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial)\n",
    "- `torch.nn`: [pytorch.org/docs/stable/nn.html](https://pytorch.org/docs/stable/nn.html)\n",
    "- `nn.Module`: [pytorch.org/docs/stable/generated/torch.nn.Module.html](https://pytorch.org/docs/stable/generated/torch.nn.Module.html)\n",
    "- `nn.Sequential`: [pytorch.org/docs/stable/generated/torch.nn.Sequential.html](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)\n",
    "- `nn.ModuleList`: [pytorch.org/docs/stable/generated/torch.nn.ModuleList.html](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_3_1_'></a>[Generate Artificial Data](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, num_features, num_classes = 4, 6, 2\n",
    "\n",
    "x = torch.randn(size=(batch_size, num_features))\n",
    "y = torch.randint(low=0, high=num_classes, size=(batch_size,))\n",
    "\n",
    "# log\n",
    "print(f\"x:\\n{x}\\n\")\n",
    "print(f\"y:\\n{y}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_3_2_'></a>[Sequential Model](#toc0_)\n",
    "\n",
    "- **Overview**:\n",
    "  - Ideal for **simpler** models where layers are stacked in a **linear** sequence.\n",
    "  - The `torch.nn.Sequential` class allows you to stack layers in a sequence, passing the output of one layer directly to the next.\n",
    "  - Suitable for straightforward models like fully-connected neural networks or basic CNNs.\n",
    "  - `nn.ModuleList` provides more flexibility, allowing dynamic layer configurations and custom forward passes.\n",
    "\n",
    "- **Key Points**:\n",
    "  - **nn.Sequential**:\n",
    "    - Layers are defined in the order they are passed to `Sequential`.\n",
    "    - No need to manually define the `forward` method; PyTorch handles it for you.\n",
    "  - **nn.ModuleList**:\n",
    "    - Layers are stored in a list-like structure.\n",
    "    - You have full control over the forward pass, allowing for more complex architectures."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_2_1_'></a>[Example 1: Using `nn.Sequential`](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "sequential_model_1 = nn.Sequential(nn.Linear(num_features, 20), nn.ReLU(), nn.Linear(20, num_classes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(sequential_model_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(sequential_model_1, input_size=(batch_size, num_features), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "sequential_model_1(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_2_2_'></a>[Example 2: Using `nn.ModuleList`](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "sequential_model_2 = nn.ModuleList([nn.Linear(num_features, 20), nn.ReLU(), nn.Linear(20, num_classes)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(sequential_model_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "def forward(layers: nn.ModuleList, x: torch.Tensor) -> torch.Tensor:\n",
    "    for layer in layers:\n",
    "        x = layer(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "forward(sequential_model_2, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_2_3_'></a>[Example 3: Mix of `nn.Sequential` and `nn.ModuleList`](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sequential_model_3 = nn.ModuleList(\n",
    "    [\n",
    "        nn.Linear(num_features, 10),\n",
    "        nn.Sequential(\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(10, 20),\n",
    "            nn.ReLU(),\n",
    "        ),\n",
    "        nn.Linear(20, num_classes),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(sequential_model_3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "def forward(layers: nn.ModuleList, x: torch.Tensor) -> torch.Tensor:\n",
    "    for layer in layers:\n",
    "        x = layer(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "forward(sequential_model_3, x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### <a id='toc2_3_3_'></a>[Non-Sequential (Functional) Model](#toc0_)\n",
    "\n",
    "- **Overview**:\n",
    "\n",
    "  - Allows for **complex** architectures with **non-linear** layer connections (e.g., skip connections in ResNet).\n",
    "  - Models are created by subclassing `torch.nn.Module`.\n",
    "  - Enables the definition of any neural network architecture, from simple feedforward networks to complex architectures like GANs or transformers.\n",
    "\n",
    "- **Key Points**:\n",
    "\n",
    "  - Use `torch.nn.Module` as the parent class and implement the `forward` method.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_3_1_'></a>[`torch.Tensor` vs. `torch.nn.Parameter`](#toc0_)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### <a id='toc2_3_3_1_1_'></a>[`torch.Tensor`](#toc0_)\n",
    "\n",
    "- **Definition**: A general-purpose tensor used to store data in PyTorch.\n",
    "- **Gradient Tracking**: Gradients are only tracked if `requires_grad=True`.\n",
    "- **Optimization**: It is not automatically registered as a parameter in a model when assigned as an attribute.\n",
    "- **Use Case**: Storing data, intermediate computations, or tensors that do not need to be optimized during training.\n",
    "- **Integration with Optimizer**: Must be explicitly added to the optimizer (if `requires_grad=True`).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel1(nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super().__init__()\n",
    "\n",
    "        # manually define weights as nn.Parameter\n",
    "        self.weight1 = torch.randn(num_features, 20)\n",
    "        self.bias1 = torch.randn(20)\n",
    "        self.weight2 = torch.randn(20, num_classes)\n",
    "        self.bias2 = torch.randn(num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.matmul(x, self.weight1) + self.bias1  # x * weight1 + bias1\n",
    "        x = torch.relu(x)\n",
    "        x = torch.matmul(x, self.weight2) + self.bias2\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialization\n",
    "functional_model_1 = CustomModel1(num_features, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(functional_model_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(functional_model_1, input_size=(batch_size, num_features), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "functional_model_1(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### <a id='toc2_3_3_1_2_'></a>[`torch.nn.Parameter`](#toc0_)\n",
    "\n",
    "- **Definition**: A subclass of `torch.Tensor` specifically designed to represent learnable parameters in `torch.nn.Module`.\n",
    "- **Gradient Tracking**: Always tracks gradients (<code>requires_grad=True</code> by default).\n",
    "- **Optimization**: It is automatically registered as a parameter of the model if assigned as an attribute to a subclass of `torch.nn.Module`.\n",
    "- **Use Case**: Learnable weights or biases of a model.\n",
    "- **Integration with Optimizer**: Automatically included in `model.parameters()` when assigned as an attribute to an `nn.Module`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel2(nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super().__init__()\n",
    "\n",
    "        # manually define weights as nn.Parameter\n",
    "        self.weight1 = nn.Parameter(torch.randn(num_features, 20))\n",
    "        self.bias1 = nn.Parameter(torch.randn(20))\n",
    "        self.weight2 = nn.Parameter(torch.randn(20, num_classes))\n",
    "        self.bias2 = nn.Parameter(torch.randn(num_classes))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.matmul(x, self.weight1) + self.bias1  # x * weight1 + bias1\n",
    "        x = torch.relu(x)\n",
    "        x = torch.matmul(x, self.weight2) + self.bias2\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialization\n",
    "functional_model_2 = CustomModel2(num_features, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(functional_model_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(functional_model_2, input_size=(batch_size, num_features), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "functional_model_2(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_3_2_'></a>[Example 1: Using `nn.Linear`, `nn.Conv2d`, ...](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel3(nn.Module):\n",
    "    def __init__(self, num_features, num_classes):\n",
    "        super().__init__()\n",
    "\n",
    "        # define linear transformation layers using `nn.Linear`\n",
    "        self.fc1 = nn.Linear(num_features, 20)\n",
    "        self.fc2 = nn.Linear(20, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialization\n",
    "functional_model_3 = CustomModel3(num_features, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(functional_model_3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(functional_model_3, input_size=(batch_size, num_features), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# feed-forward\n",
    "functional_model_3(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_3_3_'></a>[Example 2: Mix of Sequential and Non-sequential methods](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel4(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super().__init__()\n",
    "\n",
    "        # feature extractor\n",
    "        self.feature_extractor = nn.Sequential(\n",
    "            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        )\n",
    "\n",
    "        # classifier\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Flatten(start_dim=1), nn.Linear(32 * 8 * 8, 128), nn.ReLU(), nn.Linear(128, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.feature_extractor(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialization\n",
    "functional_model_4 = CustomModel4(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(functional_model_4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(functional_model_4, input_size=(1, 3, 32, 32), device=\"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### <a id='toc2_3_3_4_'></a>[Example 3: Separate Class for each Module](#toc0_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureExtractor(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)\n",
    "\n",
    "\n",
    "class Classifier(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super().__init__()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Flatten(start_dim=1), nn.Linear(32 * 8 * 8, 128), nn.ReLU(), nn.Linear(128, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.fc(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel5(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super().__init__()\n",
    "        self.feature_extractor = FeatureExtractor()\n",
    "        self.classifier = Classifier(num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.feature_extractor(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "# initialization\n",
    "functional_model_5 = CustomModel5(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# log\n",
    "print(functional_model_5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# summary\n",
    "summary(functional_model_5, input_size=(1, 3, 32, 32), device=\"cpu\")"
   ]
  }
 ],
 "metadata": {
  "author_email": "AmirhosseinHeydari78@gmail.com",
  "author_github": "https://github.com/mr-pylin",
  "author_name": "Amirhossein Heydari",
  "kernelspec": {
   "display_name": "pytorch-workshop-U_zYfVTd-py3.12",
   "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.8"
  },
  "origin_repo": "https://github.com/mr-pylin/pytorch-workshop"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}