8_MobileNetV2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:23.192118Z",
     "start_time": "2024-10-13T05:24:19.694871Z"
    }
   },
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import random_split, DataLoader\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": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Implementatation",
   "id": "6728b1ca1cd1e4dc"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:23.209186Z",
     "start_time": "2024-10-13T05:24:23.202238Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class GroupedConvolution(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, n_groups, stride):\n",
    "        super().__init__()\n",
    "        self.group_channels = in_channels // n_groups\n",
    "        self.grouped_conv_paths = nn.ModuleList([\n",
    "            nn.Conv2d(self.group_channels, out_channels // n_groups, kernel_size=3, stride=stride, padding=1,\n",
    "                      bias=False) for _ in range(n_groups)\n",
    "        ])\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = torch.concat([\n",
    "            path(x[:, self.group_channels * i:self.group_channels * (i + 1), ...]) for i, path in\n",
    "            enumerate(self.grouped_conv_paths)], dim=1)\n",
    "        return x\n",
    "\n",
    "\n",
    "class DepthwiseSeparableConvolution(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride):\n",
    "        super().__init__()\n",
    "        self.depthwise_conv = nn.Sequential(\n",
    "            # GroupedConvolution(in_channels, in_channels, in_channels, stride),\n",
    "            nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, bias=False,\n",
    "                      groups=in_channels),\n",
    "            nn.BatchNorm2d(in_channels),\n",
    "            nn.ReLU6(),\n",
    "        )\n",
    "        self.seeparable_conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.depthwise_conv(x)\n",
    "        x = self.seeparable_conv(x)\n",
    "        return x"
   ],
   "id": "63d7939dd7eefc63",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:23.685930Z",
     "start_time": "2024-10-13T05:24:23.440877Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class InvertedResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, expansion_factor, stride=1):\n",
    "        super().__init__()\n",
    "        self.expansion_factor = expansion_factor\n",
    "        self.use_skip_connection = (stride == 1 and in_channels == out_channels)\n",
    "        self.bottleneck = nn.Sequential(\n",
    "            nn.Sequential(\n",
    "                nn.Conv2d(in_channels, in_channels * expansion_factor, kernel_size=1, stride=1, bias=False),\n",
    "                nn.BatchNorm2d(in_channels * expansion_factor),\n",
    "                nn.ReLU6(),\n",
    "            ) if self.expansion_factor != 1 else nn.Identity(),\n",
    "            DepthwiseSeparableConvolution(in_channels * expansion_factor, out_channels, stride)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        bottleneck = self.bottleneck(x)\n",
    "        bottleneck = bottleneck + x if self.use_skip_connection else bottleneck\n",
    "        return bottleneck\n",
    "\n",
    "\n",
    "x_batch = torch.randn(32, 64, 52, 52)\n",
    "print(InvertedResidualBlock(64, 32, 6, 2)(x_batch).shape)"
   ],
   "id": "c26d419e1df292aa",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 32, 26, 26])\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:25.618571Z",
     "start_time": "2024-10-13T05:24:23.712032Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MobileNetV2(nn.Module):\n",
    "    def __init__(self, n_channels, n_classes):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Sequential(\n",
    "            nn.Conv2d(n_channels, 32, kernel_size=3, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU6()\n",
    "        )\n",
    "        self.conv2 = InvertedResidualBlock(32, 16, 1, 1)\n",
    "        self.conv3 = self.make_layers(16, 24, 6, 2, 2)\n",
    "        self.conv4 = self.make_layers(24, 32, 6, 3, 2)\n",
    "        self.conv5 = self.make_layers(32, 64, 6, 4, 2)\n",
    "        self.conv6 = self.make_layers(64, 96, 6, 3, 1)\n",
    "        self.conv7 = self.make_layers(96, 160, 6, 3, 2)\n",
    "        self.conv8 = self.make_layers(160, 320, 6, 1, 1)\n",
    "        self.conv9 = nn.Sequential(\n",
    "            nn.Conv2d(320, 1280, kernel_size=1, bias=False),\n",
    "            nn.BatchNorm2d(1280),\n",
    "            nn.ReLU6()\n",
    "        )\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Dropout(0.2),\n",
    "            nn.Linear(1280, n_classes),\n",
    "        )\n",
    "    \n",
    "    def make_layers(self, in_channels, out_channels, expansion_factor, n_blocks, stride):\n",
    "        layers = [InvertedResidualBlock(in_channels, out_channels, expansion_factor, stride)]\n",
    "        for layer in range(n_blocks - 1):\n",
    "            layers.append(InvertedResidualBlock(out_channels, out_channels, expansion_factor))\n",
    "        \n",
    "        return nn.Sequential(*layers)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.conv4(x)\n",
    "        x = self.conv5(x)\n",
    "        x = self.conv6(x)\n",
    "        x = self.conv7(x)\n",
    "        x = self.conv8(x)\n",
    "        x = self.conv9(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "x_batch = torch.randn(32, 3, 224, 224)\n",
    "print(MobileNetV2(3, 1000)(x_batch).shape)"
   ],
   "id": "8f6ffb7b9ada31d4",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 1000])\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:25.951694Z",
     "start_time": "2024-10-13T05:24:25.637572Z"
    }
   },
   "cell_type": "code",
   "source": [
    "mobilenet_v2_model = MobileNetV2(3, 1000)\n",
    "summary(mobilenet_v2_model, input_size=(1, 3, 224, 224), col_names=['num_params', 'mult_adds'], device='cpu', depth=2)"
   ],
   "id": "81c19b8e00414744",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "==============================================================================================================\n",
       "Layer (type:depth-idx)                                       Param #                   Mult-Adds\n",
       "==============================================================================================================\n",
       "MobileNetV2                                                  --                        --\n",
       "├─Sequential: 1-1                                            --                        --\n",
       "│    └─Conv2d: 2-1                                           864                       10,838,016\n",
       "│    └─BatchNorm2d: 2-2                                      64                        64\n",
       "│    └─ReLU6: 2-3                                            --                        --\n",
       "├─InvertedResidualBlock: 1-2                                 --                        --\n",
       "│    └─Sequential: 2-4                                       896                       10,035,296\n",
       "├─Sequential: 1-3                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-5                            5,136                     29,202,864\n",
       "│    └─InvertedResidualBlock: 2-6                            8,832                     25,740,912\n",
       "├─Sequential: 1-4                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-7                            10,000                    15,467,392\n",
       "│    └─InvertedResidualBlock: 2-8                            14,848                    10,989,376\n",
       "│    └─InvertedResidualBlock: 2-9                            14,848                    10,989,376\n",
       "├─Sequential: 1-5                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-10                           21,056                    7,564,928\n",
       "│    └─InvertedResidualBlock: 2-11                           54,272                    10,312,832\n",
       "│    └─InvertedResidualBlock: 2-12                           54,272                    10,312,832\n",
       "│    └─InvertedResidualBlock: 2-13                           54,272                    10,312,832\n",
       "├─Sequential: 1-6                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-14                           66,624                    12,721,344\n",
       "│    └─InvertedResidualBlock: 2-15                           118,272                   22,694,592\n",
       "│    └─InvertedResidualBlock: 2-16                           118,272                   22,694,592\n",
       "├─Sequential: 1-7                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-17                           155,264                   15,610,496\n",
       "│    └─InvertedResidualBlock: 2-18                           320,000                   15,480,320\n",
       "│    └─InvertedResidualBlock: 2-19                           320,000                   15,480,320\n",
       "├─Sequential: 1-8                                            --                        --\n",
       "│    └─InvertedResidualBlock: 2-20                           473,920                   23,007,040\n",
       "├─Sequential: 1-9                                            --                        --\n",
       "│    └─Conv2d: 2-21                                          409,600                   20,070,400\n",
       "│    └─BatchNorm2d: 2-22                                     2,560                     2,560\n",
       "│    └─ReLU6: 2-23                                           --                        --\n",
       "├─AdaptiveAvgPool2d: 1-10                                    --                        --\n",
       "├─Sequential: 1-11                                           --                        --\n",
       "│    └─Flatten: 2-24                                         --                        --\n",
       "│    └─Dropout: 2-25                                         --                        --\n",
       "│    └─Linear: 2-26                                          1,281,000                 1,281,000\n",
       "==============================================================================================================\n",
       "Total params: 3,504,872\n",
       "Trainable params: 3,504,872\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 300.81\n",
       "==============================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 106.86\n",
       "Params size (MB): 14.02\n",
       "Estimated Total Size (MB): 121.48\n",
       "=============================================================================================================="
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:26.234848Z",
     "start_time": "2024-10-13T05:24:25.970989Z"
    }
   },
   "cell_type": "code",
   "source": [
    "mobilenet_v2_torch_model = models.mobilenet_v2()\n",
    "summary(mobilenet_v2_torch_model, input_size=(1, 3, 224, 224),\n",
    "        col_names=['num_params', 'mult_adds'], device='cpu', depth=2)"
   ],
   "id": "d725afd43334e3c6",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "====================================================================================================\n",
       "Layer (type:depth-idx)                             Param #                   Mult-Adds\n",
       "====================================================================================================\n",
       "MobileNetV2                                        --                        --\n",
       "├─Sequential: 1-1                                  --                        --\n",
       "│    └─Conv2dNormActivation: 2-1                   928                       10,838,080\n",
       "│    └─InvertedResidual: 2-2                       896                       10,035,296\n",
       "│    └─InvertedResidual: 2-3                       5,136                     29,202,864\n",
       "│    └─InvertedResidual: 2-4                       8,832                     25,740,912\n",
       "│    └─InvertedResidual: 2-5                       10,000                    15,467,392\n",
       "│    └─InvertedResidual: 2-6                       14,848                    10,989,376\n",
       "│    └─InvertedResidual: 2-7                       14,848                    10,989,376\n",
       "│    └─InvertedResidual: 2-8                       21,056                    7,564,928\n",
       "│    └─InvertedResidual: 2-9                       54,272                    10,312,832\n",
       "│    └─InvertedResidual: 2-10                      54,272                    10,312,832\n",
       "│    └─InvertedResidual: 2-11                      54,272                    10,312,832\n",
       "│    └─InvertedResidual: 2-12                      66,624                    12,721,344\n",
       "│    └─InvertedResidual: 2-13                      118,272                   22,694,592\n",
       "│    └─InvertedResidual: 2-14                      118,272                   22,694,592\n",
       "│    └─InvertedResidual: 2-15                      155,264                   15,610,496\n",
       "│    └─InvertedResidual: 2-16                      320,000                   15,480,320\n",
       "│    └─InvertedResidual: 2-17                      320,000                   15,480,320\n",
       "│    └─InvertedResidual: 2-18                      473,920                   23,007,040\n",
       "│    └─Conv2dNormActivation: 2-19                  412,160                   20,072,960\n",
       "├─Sequential: 1-2                                  --                        --\n",
       "│    └─Dropout: 2-20                               --                        --\n",
       "│    └─Linear: 2-21                                1,281,000                 1,281,000\n",
       "====================================================================================================\n",
       "Total params: 3,504,872\n",
       "Trainable params: 3,504,872\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 300.81\n",
       "====================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 106.86\n",
       "Params size (MB): 14.02\n",
       "Estimated Total Size (MB): 121.48\n",
       "===================================================================================================="
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Training",
   "id": "e2f0ce17f284d70c"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-13T05:24:28.061181Z",
     "start_time": "2024-10-13T05:24:26.393345Z"
    }
   },
   "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": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-10-13T05:24:28.084217Z"
    }
   },
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch=18:  85%|████████▌ | 17/20 [39:06<07:27, 149.04s/it, train_loss=0.383, train_acc=87.64%, val_loss=2.523, val_acc=48.65%]"
     ]
    }
   ],
   "execution_count": null,
   "source": [
    "import wandb\n",
    "from src.engine import *\n",
    "\n",
    "config = dict(batch_size=64, lr=3e-4, epochs=20, dataset='CIFAR100')\n",
    "with wandb.init(project='pytorch-study', name='MobileNetV2', 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",
    "    mobilenet_v2_model = MobileNetV2(3, n_classes).to(DEVICE)\n",
    "    \n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(mobilenet_v2_model.parameters(), lr=w_config.lr)\n",
    "    \n",
    "    loss_history, acc_history = train(mobilenet_v2_model, train_dl, val_dl, criterion, optimizer, w_config.epochs,\n",
    "                                      DEVICE, run) "
   ],
   "id": "7feb6758495dd8a0"
  }
 ],
 "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
}