codes/12_convolutional-neural-networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📝 **Author:** Amirhossein Heydari - 📧 **Email:** amirhosseinheydari78@gmail.com - 📍 **Linktree:** [linktr.ee/mr_pylin](https://linktr.ee/mr_pylin)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn\n",
    "from torchinfo import summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a seed for deterministic results\n",
    "random_state = 42\n",
    "torch.manual_seed(random_state)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check if cuda is available\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convolution vs. Correlation\n",
    "   - Convolution and correlation are both operations used in signal processing and image analysis\n",
    "\n",
    "**[Convolution](https://en.wikipedia.org/wiki/Convolution)**:\n",
    "   - Convolution measures how one function (the kernel) modifies the other function (the signal or image).\n",
    "   - In the context of image processing, it's used to apply a filter or kernel to an image.\n",
    "   - Mathematical Formulation (discrete signals):\n",
    "   $$(f * g)[i] = \\sum_{j} f[j] \\cdot g[i - j]$$\n",
    "\n",
    "**[Correlation](https://en.wikipedia.org/wiki/Correlation)**:\n",
    "   - Correlation measures the similarity between two signals as one is shifted over the other.\n",
    "   - In image processing, it's used to detect patterns by sliding a filter over an image.\n",
    "   - Mathematical Formulation (discrete signals):\n",
    "   $$(f \\star g)[i] = \\sum_{j} f[j] \\cdot g[i + j]$$\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "    <img src=\"../resources/images/SVGs/correlation-and-convolution.svg\" alt=\"correlation-and-convolution.svg\" style=\"width: 100%;\">\n",
    "    <figcaption>Correlation vs. Convolution</figcaption>\n",
    "</figure>\n",
    "\n",
    "**Basic Concepts**:\n",
    "   - [Padding](https://medium.com/analytics-vidhya/convolution-padding-stride-and-pooling-in-cnn-13dc1f3ada26#:~:text=1%29%20%E2%88%97%20%28%F0%9D%91%9A%20%E2%88%92%20%F0%9D%91%9B%20%2B%201%29.-,Padding,-There%20are%20two)\n",
    "      - It refers to adding extra values (usually zeros) around the input tensor (signal or image) before applying the convolution operation\n",
    "      - Padding is used to control the size of the output and to allow the kernel to process the edges of the input\n",
    "      - `padding='same'`\n",
    "         - To ensure that the output of the convolution operation has the same spatial dimensions (width and height for 2D convolutions, length for 1D convolutions) as the input\n",
    "         $$p = \\left\\lceil \\frac{k - 1}{2} \\right\\rceil$$\n",
    "      - `padding='valid'`\n",
    "         - Means no padding is applied to the input\n",
    "         $$\\text{Output Size} = \\left\\lfloor \\frac{\\text{Input Size} - k}{s} + 1 \\right\\rfloor$$\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "    <img src=\"../resources/images/SVGs/convolution-padding.svg\" alt=\"convolution-padding.svg\" style=\"width: 75%;\">\n",
    "    <figcaption>Padding for Convolution</figcaption>\n",
    "</figure>\n",
    "\n",
    "   - [Stride](https://medium.com/analytics-vidhya/convolution-padding-stride-and-pooling-in-cnn-13dc1f3ada26#:~:text=in%20this%20case.-,Stride,-left%20image%3A%20stride)\n",
    "      - It defines how much the kernel moves over the input tensor during the convolution\n",
    "      - A stride of `1` means the kernel moves one step at a time, fully overlapping with each adjacent position\n",
    "      - A stride of `2` means the kernel skips one element at a time, leading to downsampling (reducing the size of the output)\n",
    "\n",
    "   - [Dilation](https://towardsdatascience.com/review-dilated-convolution-semantic-segmentation-9d5a5bd768f5)\n",
    "      - It introduces gaps between the elements of the kernel, effectively \"spreading out\" the kernel\n",
    "      - This allows the kernel to cover a larger area of the input without increasing the number of parameters (kernel size)\n",
    "      - Dilation is useful for capturing long-range dependencies in the input.\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "    <img src=\"../resources/images/SVGs/convolution-dilation.svg\" alt=\"convolution-dilation.svg\" style=\"width: 75%;\">\n",
    "    <figcaption>Dilation for Convolution</figcaption>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convolution in PyTorch\n",
    "   - Convolution operations (e.g. `nn.Conv1d`, `nn.Conv2d`) in PyTorch (and most deep learning frameworks) technically performs **correlation, not convolution!**\n",
    "   - Although the operation is named e.g. `Conv2d`, the correlation operation is preferred in practice for a few reasons\n",
    "      1. **Simplicity**:\n",
    "         - Correlation is easier to implement and understand since it doesn't require flipping the kernel\n",
    "      1. **Equivalence in Learning**:\n",
    "         - In the context of CNNs, the kernel weights are learned during training\n",
    "         - Since the kernels are learned, whether you use convolution or cross-correlation doesn't matter\n",
    "         - The network can learn equivalent filters regardless of whether the kernel is flipped or not\n",
    "\n",
    "**Docs**:\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.Conv1d.html](https://pytorch.org/docs/stable/generated/torch.nn.Conv1d.html)\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.Conv2d.html](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html)\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.Conv3d.html](https://pytorch.org/docs/stable/generated/torch.nn.Conv3d.html)\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.functional.conv1d.html](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv1d.html)\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.functional.conv2d.html](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv2d.html)\n",
    "   - [pytorch.org/docs/stable/generated/torch.nn.functional.conv3d.html](https://pytorch.org/docs/stable/generated/torch.nn.functional.conv3d.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1D Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a 1D signal and a kernel\n",
    "signal_1d = torch.arange(1, 10).reshape(1, 1, -1)      # shape: [1, 1, 10] -> (batch_size, num_channels, signal_length)\n",
    "kernel_1d = torch.tensor([2, 1, 2]).reshape(1, 1, -1)  # shape: [1, 1,  3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv_1d_1 : tensor([[[ 5, 10, 15, 20, 25, 30, 35, 40, 25]]])\n",
      "conv_1d_2 : tensor([[[10, 15, 20, 25, 30, 35, 40]]])\n",
      "conv_1d_3 : tensor([[[ 2, 10, 20, 30, 40, 18]]])\n"
     ]
    }
   ],
   "source": [
    "# convolution using torch.nn.functional.conv1d\n",
    "conv_1d_1 = F.conv1d(signal_1d, kernel_1d, padding='same')       # applies convolution with \"same\" padding, output size is the same as input size\n",
    "conv_1d_2 = F.conv1d(signal_1d, kernel_1d, padding='valid')      # applies convolution with \"valid\" padding, no padding is added, so the output size is reduced\n",
    "conv_1d_3 = F.conv1d(signal_1d, kernel_1d, padding=2, stride=2)  # applies convolution with a padding of 2 and a stride of 2, which results in downsampling the output\n",
    "\n",
    "# log\n",
    "print(f\"conv_1d_1 : {conv_1d_1}\")\n",
    "print(f\"conv_1d_2 : {conv_1d_2}\")\n",
    "print(f\"conv_1d_3 : {conv_1d_3}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig, axs = plt.subplots(nrows=1, ncols=4, figsize=(16, 4), layout='compressed')\n",
    "\n",
    "axs[0].plot(signal_1d.squeeze(), marker='o', label='Original Signal')\n",
    "axs[0].plot(kernel_1d.squeeze(), marker='o', color='purple', label='Kernel')\n",
    "axs[0].set_title(\"Original Signal\")\n",
    "axs[0].legend()\n",
    "axs[1].plot(conv_1d_1.squeeze(), marker='o', color='orange')\n",
    "axs[1].set_title(\"Convolution with \\\"Same\\\" Padding\")\n",
    "axs[2].plot(conv_1d_2.squeeze(), marker='o', color='green')\n",
    "axs[2].set_title(\"Convolution with \\\"Valid\\\" Padding\")\n",
    "axs[3].plot(conv_1d_3.squeeze(), marker='o', color='red')\n",
    "axs[3].set_title(\"Convolution with Custom Padding and Stride\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2D Correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a 2D signal (image) and a kernel\n",
    "signal_2d = torch.arange(1, 26, dtype=torch.float32).reshape(1, 1, 5, 5)                                 # shape: [1, 1, 5, 5] -> (batch_size, num_channels, signal_length)\n",
    "kernel_2d = torch.tensor([[1, 0, -1], [1, 0, -1], [1, 0, -1]], dtype=torch.float32).reshape(1, 1, 3, 3)  # shape: [1, 1, 3, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convolution using torch.nn.functional.conv2d\n",
    "conv_2d_1 = F.conv2d(signal_2d, kernel_2d, padding='same')       # applies convolution with \"same\" padding, output size is the same as input size\n",
    "conv_2d_2 = F.conv2d(signal_2d, kernel_2d, padding='valid')      # applies convolution with \"valid\" padding, no padding is added, so the output size is reduced\n",
    "conv_2d_3 = F.conv2d(signal_2d, kernel_2d, padding=1, stride=2)  # applies convolution with a padding of 1 and a stride of 2, which results in downsampling the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "fig, axs = plt.subplots(nrows=1, ncols=5, figsize=(20, 4), layout='compressed')\n",
    "\n",
    "axs[0].imshow(signal_2d.squeeze(), cmap='gray')\n",
    "axs[0].set(title=\"Original Signal (Image)\", xticks=range(signal_2d.shape[3]), yticks=range(signal_2d.shape[2]))\n",
    "axs[1].imshow(kernel_2d.squeeze(), cmap='gray')\n",
    "axs[1].set(title='Kernel', xticks=range(kernel_2d.shape[3]), yticks=range(kernel_2d.shape[2]))\n",
    "axs[2].imshow(conv_2d_1.squeeze(), cmap='gray')\n",
    "axs[2].set(title=\"Convolution with \\\"Same\\\" Padding\", xticks=range(conv_2d_1.shape[3]), yticks=range(conv_2d_1.shape[2]))\n",
    "axs[3].imshow(conv_2d_2.squeeze(), cmap='gray')\n",
    "axs[3].set(title=\"Convolution with \\\"Valid\\\" Padding\", xticks=range(conv_2d_2.shape[3]), yticks=range(conv_2d_2.shape[2]))\n",
    "axs[4].imshow(conv_2d_3.squeeze(), cmap='gray')\n",
    "axs[4].set(title=\"Convolution with Custom Padding and Stride\", xticks=range(conv_2d_3.shape[3]), yticks=range(conv_2d_3.shape[2]))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convolutional Neural Networks\n",
    "   - CNNs are a class of deep learning models specifically designed for processing structured grid-like data, such as images, videos, and even certain types of sequential data\n",
    "\n",
    "**Key Components of CNNs**\n",
    "   1. Feature Extraction\n",
    "      - Convolutional Layers\n",
    "         - This is the core building block of a CNN\n",
    "         - It involves sliding a filter (kernel) over the input data to produce a feature map\n",
    "      - Pooling Layers\n",
    "         - Pooling layers reduce the spatial dimensions of the feature maps, which helps in making the model invariant to small translations and reducing computational load\n",
    "         - Types:\n",
    "            - Max Pooling: Takes the maximum value from each patch of the feature map.\n",
    "            - Average Pooling: Takes the average value from each patch.\n",
    "   2. Classification\n",
    "      - After feature extraction, the resulting features are flattened and passed into a series of fully connected layers, forming a [Multi-Layer Perceptron (MLP)](https://github.com/mr-pylin/pytorch-workshop/blob/main/codes/08_multi-layer-perceptron.ipynb)\n",
    "      - This section performs the final classification or regression task based on the features extracted by the previous layers\n",
    "\n",
    "<figure style=\"text-align: center;\">\n",
    "    <img src=\"../resources/images/SVGs/convolutional-neural-network.svg\" alt=\"convolutional-neural-network.svg\" style=\"width: 100%;\">\n",
    "    <figcaption>Convolutional Neural Networks Model</figcaption>\n",
    "</figure>\n",
    "\n",
    "<table style=\"margin-left:auto;margin-right:auto;text-align:center;\">\n",
    "  <thead>\n",
    "    <tr>\n",
    "      <th colspan=\"4\" style=\"text-align:center;\">Feature Extraction</th>\n",
    "      <th colspan=\"4\" style=\"text-align:center;\">Classification</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <th colspan=\"2\">Convolution_1 parameters</th>\n",
    "      <th colspan=\"2\">Convolution_2 parameters</th>\n",
    "      <th colspan=\"2\">hidden<sub>1</sub> parameters</th>\n",
    "      <th colspan=\"2\">logits parameters</th>\n",
    "    </tr>\n",
    "  </thead>\n",
    "  <tbody>\n",
    "    <tr>\n",
    "      <td>Weights</td>\n",
    "      <td>Biases</td>\n",
    "      <td>Weights</td>\n",
    "      <td>Biases</td>\n",
    "      <td>Weights</td>\n",
    "      <td>Biases</td>\n",
    "      <td>Weights</td>\n",
    "      <td>Biases</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "      <td>(1 x 3 × 3) × A</td>\n",
    "      <td>A</td>\n",
    "      <td>(A x 3 × 3) × B</td>\n",
    "      <td>B</td>\n",
    "      <td>C × D</td>\n",
    "      <td>D</td>\n",
    "      <td>D × E</td>\n",
    "      <td>E</td>\n",
    "    </tr>\n",
    "  </tbody>\n",
    "  <tfoot>\n",
    "    <tr>\n",
    "      <td colspan=\"2\">(1 × 3 × 3 + 1) × A</td>\n",
    "      <td colspan=\"2\">(A × 3 × 3 + 1) × B</td>\n",
    "      <td colspan=\"2\">(C + 1) × D</td>\n",
    "      <td colspan=\"2\">(D + 1) × E</td>\n",
    "    </tr>\n",
    "  </tfoot>\n",
    "</table>\n",
    "\n",
    "\n",
    "**Training a CNN**:\n",
    "   - Forward Pass: Calculate the output using the current weights and biases.\n",
    "   - Loss Function: Commonly used loss functions for CNNs include Cross-Entropy Loss for classification tasks and Mean Squared Error for regression tasks.\n",
    "   - Backward Pass (Backpropagation): Calculate the gradient of the loss function with respect to each weight and bias.\n",
    "   - Weight Update: Update the weights and biases using an optimization algorithm like Gradient Descent or Adam.\n",
    "   - Regularization: Techniques like Dropout and Batch Normalization are used to prevent overfitting and stabilize training.\n",
    "\n",
    "**Applications of CNNs**:\n",
    "   - Image Classification: Identifying the class label of an input image.\n",
    "   - Object Detection: Locating objects within an image and identifying their class.\n",
    "   - Segmentation: Classifying each pixel in an image into a category.\n",
    "   - Face Recognition: Identifying or verifying a person based on an image of their face.\n",
    "\n",
    "**[Popular CNN Architectures](https://github.com/mr-pylin/pytorch-workshop/tree/main/codes/models/0_CNN)**\n",
    "   - LeNet-5: One of the earliest CNNs, designed for handwritten digit recognition.\n",
    "   - AlexNet: A deeper CNN that won the ImageNet competition in 2012, popularizing CNNs for large-scale image classification.\n",
    "   - VGGNet: Known for its simplicity and use of very small (3x3) filters, VGGNet showed that depth (more layers) can lead to better performance.\n",
    "   - ResNet (Residual Networks): Introduces skip connections to combat the vanishing gradient problem, enabling much deeper networks.\n",
    "   - ...\n",
    "\n",
    "**Notes**:\n",
    "   - `torch.nn.Conv2d`\n",
    "      - loss function : \n",
    "         - multi-class classification : `torch.nn.CrossEntropyLoss` = `torch.nn.LogSoftmax` + `torch.nn.NLLLoss`\n",
    "         - [pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)\n",
    "         - [pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html](https://pytorch.org/docs/stable/generated/torch.nn.NLLLoss.html)\n",
    "      - activation function for the last layer:\n",
    "         - when using `torch.nn.CrossEntropyLoss` as a loss function, the output layer doesn't need an activation function\n",
    "         - `torch.nn.CrossEntropyLoss` calculates `torch.nn.LogSoftmax` and `torch.nn.NLLLoss` internally.\n",
    "         - [pytorch.org/docs/stable/generated/torch.nn.Softmax.html](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html)\n",
    "         - [pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html](https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html)\n",
    "      - Weights\n",
    "         - Initialized based on a scheme similar to Kaiming/He initialization\n",
    "         - Uniform Distribution [default]: $W \\sim \\mathcal{U}\\left(-\\sqrt{\\frac{6}{n_{\\text{in}}}}, \\sqrt{\\frac{6}{n_{\\text{in}}}}\\right)$\n",
    "         - Normal Distribution: $W \\sim \\mathcal{N}\\left(0, \\frac{2}{n_{\\text{in}}}\\right)$\n",
    "      - Biases:\n",
    "         - Initialized to zero\n",
    "      - [pytorch.org/docs/stable/nn.init.html](https://pytorch.org/docs/stable/nn.init.html)\n",
    "      - Paper: [Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification - He, K. et al. (2015).](https://www.cv-foundation.org/openaccess/content_iccv_2015/papers/He_Delving_Deep_into_ICCV_2015_paper.pdf)\n",
    "\n",
    "**Playground**:\n",
    "   - [poloclub.github.io/cnn-explainer](https://poloclub.github.io/cnn-explainer/)\n",
    "   - [convnetplayground.fastforwardlabs.com](https://convnetplayground.fastforwardlabs.com/)\n",
    "   - [alexlenail.me/NN-SVG](https://alexlenail.me/NN-SVG/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolutional Neural Networks Using PyTorch\n",
    "   - Refer to this [notebook](https://github.com/mr-pylin/pytorch-workshop/blob/main/codes/projects/01_convolutional-neural-networks.ipynb) for a comprehensive example on the CNN concept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CIFAR10Model(nn.Module):\n",
    "    def __init__(self, in_channels, output_dim):\n",
    "        super(CIFAR10Model, self).__init__()\n",
    "        self.feature_extractor = nn.Sequential(\n",
    "\n",
    "            # 3x32x32\n",
    "            nn.Conv2d(in_channels, out_channels=32, kernel_size=3),\n",
    "            nn.BatchNorm2d(32),  # StandardScaler along channel axis\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            # 32x15x15\n",
    "\n",
    "            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(kernel_size=2),\n",
    "            # 64x6x6\n",
    "\n",
    "            nn.AdaptiveAvgPool2d(output_size=(1, 1))\n",
    "            # 64x1x1\n",
    "        )\n",
    "\n",
    "        self.flatten = nn.Flatten(start_dim=1)\n",
    "\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(64, output_dim),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.feature_extractor(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example input\n",
    "batch_size = 3\n",
    "x = torch.randn(batch_size, 3, 32, 32)\n",
    "y = torch.tensor([1, 0, 1], dtype=torch.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CIFAR10Model(\n",
       "  (feature_extractor): Sequential(\n",
       "    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (5): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (6): ReLU()\n",
       "    (7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (8): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  )\n",
       "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  (classifier): Sequential(\n",
       "    (0): Linear(in_features=64, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# initialize the CNN\n",
    "in_channels = x.shape[1]  # \n",
    "output_dim = 10           # \n",
    "\n",
    "model = CIFAR10Model(in_channels, output_dim)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==========================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #\n",
       "==========================================================================================\n",
       "CIFAR10Model                             [3, 10]                   --\n",
       "├─Sequential: 1-1                        [3, 64, 1, 1]             --\n",
       "│    └─Conv2d: 2-1                       [3, 32, 30, 30]           896\n",
       "│    └─BatchNorm2d: 2-2                  [3, 32, 30, 30]           64\n",
       "│    └─ReLU: 2-3                         [3, 32, 30, 30]           --\n",
       "│    └─MaxPool2d: 2-4                    [3, 32, 15, 15]           --\n",
       "│    └─Conv2d: 2-5                       [3, 64, 13, 13]           18,496\n",
       "│    └─BatchNorm2d: 2-6                  [3, 64, 13, 13]           128\n",
       "│    └─ReLU: 2-7                         [3, 64, 13, 13]           --\n",
       "│    └─MaxPool2d: 2-8                    [3, 64, 6, 6]             --\n",
       "│    └─AdaptiveAvgPool2d: 2-9            [3, 64, 1, 1]             --\n",
       "├─Flatten: 1-2                           [3, 64]                   --\n",
       "├─Sequential: 1-3                        [3, 10]                   --\n",
       "│    └─Linear: 2-10                      [3, 10]                   650\n",
       "==========================================================================================\n",
       "Total params: 20,234\n",
       "Trainable params: 20,234\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 11.80\n",
       "==========================================================================================\n",
       "Input size (MB): 0.04\n",
       "Forward/backward pass size (MB): 1.90\n",
       "Params size (MB): 0.08\n",
       "Estimated Total Size (MB): 2.02\n",
       "=========================================================================================="
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "summary(model, input_size=x.size(), device='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch   1/10  ->  Loss: 2.8423\n",
      "epoch   2/10  ->  Loss: 1.0186\n",
      "epoch   3/10  ->  Loss: 0.3171\n",
      "epoch   4/10  ->  Loss: 0.1235\n",
      "epoch   5/10  ->  Loss: 0.0634\n",
      "epoch   6/10  ->  Loss: 0.0399\n",
      "epoch   7/10  ->  Loss: 0.0284\n",
      "epoch   8/10  ->  Loss: 0.0215\n",
      "epoch   9/10  ->  Loss: 0.0168\n",
      "epoch  10/10  ->  Loss: 0.0133\n"
     ]
    }
   ],
   "source": [
    "# define a loss function\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# define an optimizer (e.g., SGD)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "\n",
    "# training loop\n",
    "num_epochs = 10  # Number of epochs\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "\n",
    "    # forward pass\n",
    "    output = model(x)\n",
    "    \n",
    "    # compute the loss\n",
    "    loss = criterion(output, y)\n",
    "    \n",
    "    # perform backward propagation automatically\n",
    "    loss.backward()\n",
    "    \n",
    "    # update the weights & zero the gradients\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # log\n",
    "    print(f'epoch {epoch+1:3}/{num_epochs}  ->  Loss: {loss.item():.4f}')"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}