11_ConvNeXt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-10-14T01:26:04.580511Z",
     "start_time": "2024-10-14T01:26:04.546438Z"
    }
   },
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import random_split, DataLoader\n",
    "from torchaudio.models.wav2vec2.components import LayerNorm\n",
    "from torchinfo import summary\n",
    "from torchvision import datasets, transforms, models\n",
    "\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(DEVICE)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "execution_count": 32
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Implementatation",
   "id": "6728b1ca1cd1e4dc"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T01:26:04.621586Z",
     "start_time": "2024-10-14T01:26:04.598254Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torchvision.ops.misc import Permute\n",
    "\n",
    "class ChannelLayerNoorm(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super().__init__()\n",
    "        self.permute1 = Permute([0, 2, 3, 1])\n",
    "        self.layer_norm = nn.LayerNorm(channels)\n",
    "        self.permute2 = Permute([0, 3, 1, 2])\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.permute1(x)\n",
    "        x = self.layer_norm(x)\n",
    "        x = self.permute2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class ConvNeXtBlock(nn.Module):\n",
    "    def __init__(self, input_channels, inner_channels):\n",
    "        super().__init__()\n",
    "        self.residual_block = nn.Sequential(\n",
    "            nn.Conv2d(input_channels, inner_channels, kernel_size=7, stride=1, padding=3, groups=inner_channels),\n",
    "            nn.GELU(),\n",
    "            Permute([0, 2, 3, 1]),\n",
    "            nn.LayerNorm(input_channels),\n",
    "            nn.Linear(inner_channels, inner_channels*4),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(inner_channels*4, inner_channels),\n",
    "            Permute([0, 3, 1, 2])\n",
    "        )\n",
    "        self.layer_scale = nn.Parameter(torch.full((1, inner_channels, 1, 1), 1e-6))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        shortcut = x\n",
    "        x = self.residual_block(x) * self.layer_scale\n",
    "        return x + shortcut\n",
    "\n",
    "\n",
    "class ConvNeXt(nn.Module):\n",
    "    def __init__(self, n_channels, n_classes):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Sequential(\n",
    "            nn.Conv2d(n_channels, 96, kernel_size=4, stride=4),\n",
    "            ChannelLayerNoorm(96),\n",
    "            nn.GELU(),\n",
    "        )\n",
    "        self.layers1 = self.make_layers(96, 96, 3, False)\n",
    "        self.layers2 = self.make_layers(96, 192, 3, True)\n",
    "        self.layers3 = self.make_layers(192, 384, 9, True)\n",
    "        self.layers4 = self.make_layers(384, 768, 3, True)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.classifier = nn.Sequential(\n",
    "            ChannelLayerNoorm(768),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(768, n_classes)\n",
    "        )\n",
    "    \n",
    "    def make_layers(self, input_channels, inner_channels, n_blocks, downsample):\n",
    "        layers = [nn.Sequential(\n",
    "            ChannelLayerNoorm(input_channels),\n",
    "            nn.Conv2d(input_channels, inner_channels, kernel_size=2, stride=2),\n",
    "        )] if downsample else []\n",
    "        for layer in range(n_blocks):\n",
    "            layers.append(ConvNeXtBlock(inner_channels, inner_channels))\n",
    "        \n",
    "        return nn.Sequential(*layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.layers1(x)\n",
    "        x = self.layers2(x)\n",
    "        x = self.layers3(x)\n",
    "        x = self.layers4(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = self.classifier(x)\n",
    "        return x"
   ],
   "id": "8f6ffb7b9ada31d4",
   "outputs": [],
   "execution_count": 33
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T01:26:05.089826Z",
     "start_time": "2024-10-14T01:26:04.641800Z"
    }
   },
   "cell_type": "code",
   "source": [
    "convnext_model = ConvNeXt(3, 1000)\n",
    "summary(convnext_model, input_size=(1, 3, 224, 224), col_names=['output_size', 'num_params', 'mult_adds'], depth=4, device='cpu')"
   ],
   "id": "81c19b8e00414744",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "===================================================================================================================\n",
       "Layer (type:depth-idx)                   Output Shape              Param #                   Mult-Adds\n",
       "===================================================================================================================\n",
       "ConvNeXt                                 [1, 1000]                 --                        --\n",
       "├─Sequential: 1-1                        [1, 96, 56, 56]           --                        --\n",
       "│    └─Conv2d: 2-1                       [1, 96, 56, 56]           4,704                     14,751,744\n",
       "│    └─ChannelLayerNoorm: 2-2            [1, 96, 56, 56]           --                        --\n",
       "│    │    └─Permute: 3-1                 [1, 56, 56, 96]           --                        --\n",
       "│    │    └─LayerNorm: 3-2               [1, 56, 56, 96]           192                       192\n",
       "│    │    └─Permute: 3-3                 [1, 96, 56, 56]           --                        --\n",
       "│    └─GELU: 2-3                         [1, 96, 56, 56]           --                        --\n",
       "├─Sequential: 1-2                        [1, 96, 56, 56]           --                        --\n",
       "│    └─ConvNeXtBlock: 2-4                [1, 96, 56, 56]           96                        --\n",
       "│    │    └─Sequential: 3-4              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Conv2d: 4-1             [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    └─GELU: 4-2               [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Permute: 4-3            [1, 56, 56, 96]           --                        --\n",
       "│    │    │    └─LayerNorm: 4-4          [1, 56, 56, 96]           192                       192\n",
       "│    │    │    └─Linear: 4-5             [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    └─GELU: 4-6               [1, 56, 56, 384]          --                        --\n",
       "│    │    │    └─Linear: 4-7             [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    └─Permute: 4-8            [1, 96, 56, 56]           --                        --\n",
       "│    └─ConvNeXtBlock: 2-5                [1, 96, 56, 56]           96                        --\n",
       "│    │    └─Sequential: 3-5              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Conv2d: 4-9             [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    └─GELU: 4-10              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Permute: 4-11           [1, 56, 56, 96]           --                        --\n",
       "│    │    │    └─LayerNorm: 4-12         [1, 56, 56, 96]           192                       192\n",
       "│    │    │    └─Linear: 4-13            [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    └─GELU: 4-14              [1, 56, 56, 384]          --                        --\n",
       "│    │    │    └─Linear: 4-15            [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    └─Permute: 4-16           [1, 96, 56, 56]           --                        --\n",
       "│    └─ConvNeXtBlock: 2-6                [1, 96, 56, 56]           96                        --\n",
       "│    │    └─Sequential: 3-6              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Conv2d: 4-17            [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    └─GELU: 4-18              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Permute: 4-19           [1, 56, 56, 96]           --                        --\n",
       "│    │    │    └─LayerNorm: 4-20         [1, 56, 56, 96]           192                       192\n",
       "│    │    │    └─Linear: 4-21            [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    └─GELU: 4-22              [1, 56, 56, 384]          --                        --\n",
       "│    │    │    └─Linear: 4-23            [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    └─Permute: 4-24           [1, 96, 56, 56]           --                        --\n",
       "├─Sequential: 1-3                        [1, 192, 28, 28]          --                        --\n",
       "│    └─Sequential: 2-7                   [1, 192, 28, 28]          --                        --\n",
       "│    │    └─ChannelLayerNoorm: 3-7       [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─Permute: 4-25           [1, 56, 56, 96]           --                        --\n",
       "│    │    │    └─LayerNorm: 4-26         [1, 56, 56, 96]           192                       192\n",
       "│    │    │    └─Permute: 4-27           [1, 96, 56, 56]           --                        --\n",
       "│    │    └─Conv2d: 3-8                  [1, 192, 28, 28]          73,920                    57,953,280\n",
       "│    └─ConvNeXtBlock: 2-8                [1, 192, 28, 28]          192                       --\n",
       "│    │    └─Sequential: 3-9              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Conv2d: 4-28            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    └─GELU: 4-29              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Permute: 4-30           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-31         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    └─Linear: 4-32            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    └─GELU: 4-33              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    └─Linear: 4-34            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    └─Permute: 4-35           [1, 192, 28, 28]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-9                [1, 192, 28, 28]          192                       --\n",
       "│    │    └─Sequential: 3-10             [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Conv2d: 4-36            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    └─GELU: 4-37              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Permute: 4-38           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-39         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    └─Linear: 4-40            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    └─GELU: 4-41              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    └─Linear: 4-42            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    └─Permute: 4-43           [1, 192, 28, 28]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-10               [1, 192, 28, 28]          192                       --\n",
       "│    │    └─Sequential: 3-11             [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Conv2d: 4-44            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    └─GELU: 4-45              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Permute: 4-46           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-47         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    └─Linear: 4-48            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    └─GELU: 4-49              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    └─Linear: 4-50            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    └─Permute: 4-51           [1, 192, 28, 28]          --                        --\n",
       "├─Sequential: 1-4                        [1, 384, 14, 14]          --                        --\n",
       "│    └─Sequential: 2-11                  [1, 384, 14, 14]          --                        --\n",
       "│    │    └─ChannelLayerNoorm: 3-12      [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─Permute: 4-52           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-53         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    └─Permute: 4-54           [1, 192, 28, 28]          --                        --\n",
       "│    │    └─Conv2d: 3-13                 [1, 384, 14, 14]          295,296                   57,878,016\n",
       "│    └─ConvNeXtBlock: 2-12               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-14             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-55            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-56              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-57           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-58         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-59            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-60              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-61            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-62           [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-13               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-15             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-63            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-64              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-65           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-66         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-67            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-68              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-69            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-70           [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-14               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-16             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-71            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-72              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-73           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-74         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-75            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-76              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-77            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-78           [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-15               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-17             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-79            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-80              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-81           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-82         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-83            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-84              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-85            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-86           [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-16               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-18             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-87            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-88              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-89           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-90         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-91            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-92              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-93            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-94           [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-17               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-19             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-95            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-96              [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-97           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-98         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-99            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-100             [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-101           [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-102          [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-18               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-20             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-103           [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-104             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-105          [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-106        [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-107           [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-108             [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-109           [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-110          [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-19               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-21             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-111           [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-112             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-113          [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-114        [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-115           [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-116             [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-117           [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-118          [1, 384, 14, 14]          --                        --\n",
       "│    └─ConvNeXtBlock: 2-20               [1, 384, 14, 14]          384                       --\n",
       "│    │    └─Sequential: 3-22             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Conv2d: 4-119           [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    └─GELU: 4-120             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-121          [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-122        [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Linear: 4-123           [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    └─GELU: 4-124             [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    └─Linear: 4-125           [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    └─Permute: 4-126          [1, 384, 14, 14]          --                        --\n",
       "├─Sequential: 1-5                        [1, 768, 7, 7]            --                        --\n",
       "│    └─Sequential: 2-21                  [1, 768, 7, 7]            --                        --\n",
       "│    │    └─ChannelLayerNoorm: 3-23      [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─Permute: 4-127          [1, 14, 14, 384]          --                        --\n",
       "│    │    │    └─LayerNorm: 4-128        [1, 14, 14, 384]          768                       768\n",
       "│    │    │    └─Permute: 4-129          [1, 384, 14, 14]          --                        --\n",
       "│    │    └─Conv2d: 3-24                 [1, 768, 7, 7]            1,180,416                 57,840,384\n",
       "│    └─ConvNeXtBlock: 2-22               [1, 768, 7, 7]            768                       --\n",
       "│    │    └─Sequential: 3-25             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Conv2d: 4-130           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    └─GELU: 4-131             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Permute: 4-132          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    └─LayerNorm: 4-133        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    └─Linear: 4-134           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    └─GELU: 4-135             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    └─Linear: 4-136           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    └─Permute: 4-137          [1, 768, 7, 7]            --                        --\n",
       "│    └─ConvNeXtBlock: 2-23               [1, 768, 7, 7]            768                       --\n",
       "│    │    └─Sequential: 3-26             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Conv2d: 4-138           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    └─GELU: 4-139             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Permute: 4-140          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    └─LayerNorm: 4-141        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    └─Linear: 4-142           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    └─GELU: 4-143             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    └─Linear: 4-144           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    └─Permute: 4-145          [1, 768, 7, 7]            --                        --\n",
       "│    └─ConvNeXtBlock: 2-24               [1, 768, 7, 7]            768                       --\n",
       "│    │    └─Sequential: 3-27             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Conv2d: 4-146           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    └─GELU: 4-147             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─Permute: 4-148          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    └─LayerNorm: 4-149        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    └─Linear: 4-150           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    └─GELU: 4-151             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    └─Linear: 4-152           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    └─Permute: 4-153          [1, 768, 7, 7]            --                        --\n",
       "├─AdaptiveAvgPool2d: 1-6                 [1, 768, 1, 1]            --                        --\n",
       "├─Sequential: 1-7                        [1, 1000]                 --                        --\n",
       "│    └─ChannelLayerNoorm: 2-25           [1, 768, 1, 1]            --                        --\n",
       "│    │    └─Permute: 3-28                [1, 1, 1, 768]            --                        --\n",
       "│    │    └─LayerNorm: 3-29              [1, 1, 1, 768]            1,536                     1,536\n",
       "│    │    └─Permute: 3-30                [1, 768, 1, 1]            --                        --\n",
       "│    └─Flatten: 2-26                     [1, 768]                  --                        --\n",
       "│    └─Linear: 2-27                      [1, 1000]                 769,000                   769,000\n",
       "===================================================================================================================\n",
       "Total params: 28,589,128\n",
       "Trainable params: 28,589,128\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 322.37\n",
       "===================================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 131.27\n",
       "Params size (MB): 114.33\n",
       "Estimated Total Size (MB): 246.21\n",
       "==================================================================================================================="
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 34
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T01:26:05.542393Z",
     "start_time": "2024-10-14T01:26:05.101727Z"
    }
   },
   "cell_type": "code",
   "source": [
    "convnext_torch_model = models.convnext_tiny()\n",
    "summary(convnext_torch_model, input_size=(1, 3, 224, 224), col_names=['output_size', 'num_params', 'mult_adds'], depth=5, device='cpu')"
   ],
   "id": "db1cb73d8c7629f1",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "========================================================================================================================\n",
       "Layer (type:depth-idx)                        Output Shape              Param #                   Mult-Adds\n",
       "========================================================================================================================\n",
       "ConvNeXt                                      [1, 1000]                 --                        --\n",
       "├─Sequential: 1-1                             [1, 768, 7, 7]            --                        --\n",
       "│    └─Conv2dNormActivation: 2-1              [1, 96, 56, 56]           --                        --\n",
       "│    │    └─Conv2d: 3-1                       [1, 96, 56, 56]           4,704                     14,751,744\n",
       "│    │    └─LayerNorm2d: 3-2                  [1, 96, 56, 56]           192                       192\n",
       "│    └─Sequential: 2-2                        [1, 96, 56, 56]           --                        --\n",
       "│    │    └─CNBlock: 3-3                      [1, 96, 56, 56]           96                        --\n",
       "│    │    │    └─Sequential: 4-1              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    │    └─Conv2d: 5-1             [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    │    └─Permute: 5-2            [1, 56, 56, 96]           --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-3          [1, 56, 56, 96]           192                       192\n",
       "│    │    │    │    └─Linear: 5-4             [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    │    └─GELU: 5-5               [1, 56, 56, 384]          --                        --\n",
       "│    │    │    │    └─Linear: 5-6             [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    │    └─Permute: 5-7            [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─StochasticDepth: 4-2         [1, 96, 56, 56]           --                        --\n",
       "│    │    └─CNBlock: 3-4                      [1, 96, 56, 56]           96                        --\n",
       "│    │    │    └─Sequential: 4-3              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    │    └─Conv2d: 5-8             [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    │    └─Permute: 5-9            [1, 56, 56, 96]           --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-10         [1, 56, 56, 96]           192                       192\n",
       "│    │    │    │    └─Linear: 5-11            [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    │    └─GELU: 5-12              [1, 56, 56, 384]          --                        --\n",
       "│    │    │    │    └─Linear: 5-13            [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    │    └─Permute: 5-14           [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─StochasticDepth: 4-4         [1, 96, 56, 56]           --                        --\n",
       "│    │    └─CNBlock: 3-5                      [1, 96, 56, 56]           96                        --\n",
       "│    │    │    └─Sequential: 4-5              [1, 96, 56, 56]           --                        --\n",
       "│    │    │    │    └─Conv2d: 5-15            [1, 96, 56, 56]           4,800                     15,052,800\n",
       "│    │    │    │    └─Permute: 5-16           [1, 56, 56, 96]           --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-17         [1, 56, 56, 96]           192                       192\n",
       "│    │    │    │    └─Linear: 5-18            [1, 56, 56, 384]          37,248                    37,248\n",
       "│    │    │    │    └─GELU: 5-19              [1, 56, 56, 384]          --                        --\n",
       "│    │    │    │    └─Linear: 5-20            [1, 56, 56, 96]           36,960                    36,960\n",
       "│    │    │    │    └─Permute: 5-21           [1, 96, 56, 56]           --                        --\n",
       "│    │    │    └─StochasticDepth: 4-6         [1, 96, 56, 56]           --                        --\n",
       "│    └─Sequential: 2-3                        [1, 192, 28, 28]          --                        --\n",
       "│    │    └─LayerNorm2d: 3-6                  [1, 96, 56, 56]           192                       192\n",
       "│    │    └─Conv2d: 3-7                       [1, 192, 28, 28]          73,920                    57,953,280\n",
       "│    └─Sequential: 2-4                        [1, 192, 28, 28]          --                        --\n",
       "│    │    └─CNBlock: 3-8                      [1, 192, 28, 28]          192                       --\n",
       "│    │    │    └─Sequential: 4-7              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-22            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    │    └─Permute: 5-23           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-24         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    │    └─Linear: 5-25            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    │    └─GELU: 5-26              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    │    └─Linear: 5-27            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    │    └─Permute: 5-28           [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-8         [1, 192, 28, 28]          --                        --\n",
       "│    │    └─CNBlock: 3-9                      [1, 192, 28, 28]          192                       --\n",
       "│    │    │    └─Sequential: 4-9              [1, 192, 28, 28]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-29            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    │    └─Permute: 5-30           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-31         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    │    └─Linear: 5-32            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    │    └─GELU: 5-33              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    │    └─Linear: 5-34            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    │    └─Permute: 5-35           [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-10        [1, 192, 28, 28]          --                        --\n",
       "│    │    └─CNBlock: 3-10                     [1, 192, 28, 28]          192                       --\n",
       "│    │    │    └─Sequential: 4-11             [1, 192, 28, 28]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-36            [1, 192, 28, 28]          9,600                     7,526,400\n",
       "│    │    │    │    └─Permute: 5-37           [1, 28, 28, 192]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-38         [1, 28, 28, 192]          384                       384\n",
       "│    │    │    │    └─Linear: 5-39            [1, 28, 28, 768]          148,224                   148,224\n",
       "│    │    │    │    └─GELU: 5-40              [1, 28, 28, 768]          --                        --\n",
       "│    │    │    │    └─Linear: 5-41            [1, 28, 28, 192]          147,648                   147,648\n",
       "│    │    │    │    └─Permute: 5-42           [1, 192, 28, 28]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-12        [1, 192, 28, 28]          --                        --\n",
       "│    └─Sequential: 2-5                        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─LayerNorm2d: 3-11                 [1, 192, 28, 28]          384                       384\n",
       "│    │    └─Conv2d: 3-12                      [1, 384, 14, 14]          295,296                   57,878,016\n",
       "│    └─Sequential: 2-6                        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-13                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-13             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-43            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-44           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-45         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-46            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-47              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-48            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-49           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-14        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-14                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-15             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-50            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-51           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-52         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-53            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-54              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-55            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-56           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-16        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-15                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-17             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-57            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-58           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-59         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-60            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-61              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-62            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-63           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-18        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-16                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-19             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-64            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-65           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-66         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-67            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-68              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-69            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-70           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-20        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-17                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-21             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-71            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-72           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-73         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-74            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-75              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-76            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-77           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-22        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-18                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-23             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-78            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-79           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-80         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-81            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-82              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-83            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-84           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-24        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-19                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-25             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-85            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-86           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-87         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-88            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-89              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-90            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-91           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-26        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-20                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-27             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-92            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-93           [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-94         [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-95            [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-96              [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-97            [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-98           [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-28        [1, 384, 14, 14]          --                        --\n",
       "│    │    └─CNBlock: 3-21                     [1, 384, 14, 14]          384                       --\n",
       "│    │    │    └─Sequential: 4-29             [1, 384, 14, 14]          --                        --\n",
       "│    │    │    │    └─Conv2d: 5-99            [1, 384, 14, 14]          19,200                    3,763,200\n",
       "│    │    │    │    └─Permute: 5-100          [1, 14, 14, 384]          --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-101        [1, 14, 14, 384]          768                       768\n",
       "│    │    │    │    └─Linear: 5-102           [1, 14, 14, 1536]         591,360                   591,360\n",
       "│    │    │    │    └─GELU: 5-103             [1, 14, 14, 1536]         --                        --\n",
       "│    │    │    │    └─Linear: 5-104           [1, 14, 14, 384]          590,208                   590,208\n",
       "│    │    │    │    └─Permute: 5-105          [1, 384, 14, 14]          --                        --\n",
       "│    │    │    └─StochasticDepth: 4-30        [1, 384, 14, 14]          --                        --\n",
       "│    └─Sequential: 2-7                        [1, 768, 7, 7]            --                        --\n",
       "│    │    └─LayerNorm2d: 3-22                 [1, 384, 14, 14]          768                       768\n",
       "│    │    └─Conv2d: 3-23                      [1, 768, 7, 7]            1,180,416                 57,840,384\n",
       "│    └─Sequential: 2-8                        [1, 768, 7, 7]            --                        --\n",
       "│    │    └─CNBlock: 3-24                     [1, 768, 7, 7]            768                       --\n",
       "│    │    │    └─Sequential: 4-31             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    │    └─Conv2d: 5-106           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    │    └─Permute: 5-107          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-108        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    │    └─Linear: 5-109           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    │    └─GELU: 5-110             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    │    └─Linear: 5-111           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    │    └─Permute: 5-112          [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─StochasticDepth: 4-32        [1, 768, 7, 7]            --                        --\n",
       "│    │    └─CNBlock: 3-25                     [1, 768, 7, 7]            768                       --\n",
       "│    │    │    └─Sequential: 4-33             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    │    └─Conv2d: 5-113           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    │    └─Permute: 5-114          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-115        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    │    └─Linear: 5-116           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    │    └─GELU: 5-117             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    │    └─Linear: 5-118           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    │    └─Permute: 5-119          [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─StochasticDepth: 4-34        [1, 768, 7, 7]            --                        --\n",
       "│    │    └─CNBlock: 3-26                     [1, 768, 7, 7]            768                       --\n",
       "│    │    │    └─Sequential: 4-35             [1, 768, 7, 7]            --                        --\n",
       "│    │    │    │    └─Conv2d: 5-120           [1, 768, 7, 7]            38,400                    1,881,600\n",
       "│    │    │    │    └─Permute: 5-121          [1, 7, 7, 768]            --                        --\n",
       "│    │    │    │    └─LayerNorm: 5-122        [1, 7, 7, 768]            1,536                     1,536\n",
       "│    │    │    │    └─Linear: 5-123           [1, 7, 7, 3072]           2,362,368                 2,362,368\n",
       "│    │    │    │    └─GELU: 5-124             [1, 7, 7, 3072]           --                        --\n",
       "│    │    │    │    └─Linear: 5-125           [1, 7, 7, 768]            2,360,064                 2,360,064\n",
       "│    │    │    │    └─Permute: 5-126          [1, 768, 7, 7]            --                        --\n",
       "│    │    │    └─StochasticDepth: 4-36        [1, 768, 7, 7]            --                        --\n",
       "├─AdaptiveAvgPool2d: 1-2                      [1, 768, 1, 1]            --                        --\n",
       "├─Sequential: 1-3                             [1, 1000]                 --                        --\n",
       "│    └─LayerNorm2d: 2-9                       [1, 768, 1, 1]            1,536                     1,536\n",
       "│    └─Flatten: 2-10                          [1, 768]                  --                        --\n",
       "│    └─Linear: 2-11                           [1, 1000]                 769,000                   769,000\n",
       "========================================================================================================================\n",
       "Total params: 28,589,128\n",
       "Trainable params: 28,589,128\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 322.37\n",
       "========================================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 131.27\n",
       "Params size (MB): 114.33\n",
       "Estimated Total Size (MB): 246.21\n",
       "========================================================================================================================"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 35
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Training",
   "id": "e2f0ce17f284d70c"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T01:26:07.138318Z",
     "start_time": "2024-10-14T01:26:05.682939Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from pathlib import Path\n",
    "\n",
    "TRAIN_RATIO = 0.8\n",
    "data_dir = Path('./data/')\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "train_ds = datasets.CIFAR100(data_dir, train=True, download=True, transform=transform)\n",
    "train_ds, val_ds = random_split(train_ds, (TRAIN_RATIO, 1 - TRAIN_RATIO))\n",
    "val_ds.transform = transform\n",
    "test_ds = datasets.CIFAR100(data_dir, train=False, download=True, transform=transform)"
   ],
   "id": "42d21d784414347d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "execution_count": 36
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T04:11:13.650836Z",
     "start_time": "2024-10-14T01:26:07.149308Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import wandb\n",
    "from src.engine import *\n",
    "\n",
    "config = dict(batch_size=64, lr=5e-4, epochs=20, dataset='CIFAR100')\n",
    "with wandb.init(project='pytorch-study', name='ConvNeXt50', config=config) as run:\n",
    "    w_config = run.config\n",
    "    train_dl = DataLoader(train_ds, batch_size=w_config.batch_size, shuffle=True)\n",
    "    val_dl = DataLoader(val_ds, batch_size=w_config.batch_size, shuffle=True)\n",
    "    \n",
    "    n_classes = len(train_ds.dataset.classes)\n",
    "    convnext_mdoel = ConvNeXt(3, n_classes).to(DEVICE)\n",
    "        \n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(convnext_mdoel.parameters(), lr=w_config.lr)\n",
    "    \n",
    "    loss_history, acc_history = train(convnext_mdoel, train_dl, val_dl, criterion, optimizer, w_config.epochs, DEVICE, run) "
   ],
   "id": "7feb6758495dd8a0",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch=20: 100%|██████████| 20/20 [2:44:53<00:00, 494.67s/it, train_loss=0.112, train_acc=96.36%, val_loss=4.290, val_acc=44.60%]  \n"
     ]
    }
   ],
   "execution_count": 37
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}