VisionTransformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.init as init\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "from medmnist import DermaMNIST"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Vision Transformer Model"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Attention, self).__init__()\n",
    "\n",
    "    def forward(self, Query, Key, Value, D):\n",
    "        scaling_factor = torch.sqrt(torch.tensor(D, dtype=torch.float32))\n",
    "\n",
    "        QueryKey = torch.matmul(Query, torch.transpose(Key, -1, -2))\n",
    "        attention_scores = torch.div(QueryKey, scaling_factor)\n",
    "        softmax = F.softmax(attention_scores, dim = -1)\n",
    "        return torch.matmul(softmax, Value)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, num_heads, D):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        dimensions = D / num_heads\n",
    "        assert dimensions % 1 == 0, f\"D must be divisible by heads, got heads: {num_heads}, D: {D}\"\n",
    "        self.num_heads = num_heads\n",
    "        self.dimensions = int(dimensions)\n",
    "        self.D = D\n",
    "\n",
    "        self.Q_Linear = nn.Linear(D, D)\n",
    "        self.K_Linear = nn.Linear(D, D)\n",
    "        self.V_Linear = nn.Linear(D, D)\n",
    "\n",
    "        self.attention = Attention()\n",
    "\n",
    "        self.final_layer = nn.Linear(D, D)\n",
    "\n",
    "    def forward(self, X):\n",
    "        batch_size = X.shape[:-2]\n",
    "\n",
    "\n",
    "        Query = self.Q_Linear(X)\n",
    "        Key = self.K_Linear(X)\n",
    "        Value = self.V_Linear(X)\n",
    "\n",
    "        Query = Query.view((*batch_size, -1, self.num_heads , self.dimensions)).transpose(-3,-2)\n",
    "        Key = Key.view((*batch_size, -1, self.num_heads , self.dimensions)).transpose(-3,-2)\n",
    "        Value = Value.view((*batch_size, -1, self.num_heads , self.dimensions)).transpose(-3,-2)\n",
    "\n",
    "        heads = self.attention(Query, Key, Value, self.dimensions)\n",
    "        concatenated_heads = heads.transpose(-3, -2).contiguous().view(*batch_size, -1, D)\n",
    "\n",
    "        return self.final_layer(concatenated_heads)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, num_heads, D, mlp_width = 248):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "\n",
    "        # The operation follow in this order\n",
    "        self.norm_1 = nn.LayerNorm(D)\n",
    "        self.multi_head_attention = MultiHeadAttention(num_heads, D)\n",
    "        # Residual Connection will come in here\n",
    "        self.norm_2 = nn.LayerNorm(D)\n",
    "        self.MLP = nn.Sequential(\n",
    "            nn.Linear(D, mlp_width),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(mlp_width, D)\n",
    "        )\n",
    "        # Another Residual Connection will come here\n",
    "\n",
    "    def forward(self, X):\n",
    "        first_output = self.norm_1(X)\n",
    "        first_output = self.multi_head_attention(first_output)\n",
    "        first_output = X + first_output # Residual connection\n",
    "        second_output = self.norm_2(first_output)\n",
    "        second_output = self.MLP(second_output)\n",
    "        final_output = first_output + second_output # Second residual connection\n",
    "        return final_output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "outputs": [],
   "source": [
    "class LinearProjectionOfFlattenedPatches(nn.Module):\n",
    "    def __init__(self, in_channels, heigth, width, patch_size,  D):\n",
    "        super(LinearProjectionOfFlattenedPatches, self).__init__()\n",
    "        self.img_size = (heigth, width)\n",
    "        self.patch_size = patch_size  # P\n",
    "        self.in_channels = in_channels  # C\n",
    "\n",
    "        self.N = (self.img_size[0] * self.img_size[1]) / (patch_size**2)\n",
    "        assert self.N % 1 == 0, f\"num_patches must be divisible by patch size: {patch_size}, size: {self.img_size}\"\n",
    "        self.N = int(self.N)\n",
    "        patch_dimensions = patch_size * patch_size * in_channels\n",
    "\n",
    "        # Linear projection of patches into latent vector of dimension D\n",
    "        self.projection = nn.Linear(patch_dimensions, D)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[:-3]\n",
    "        permutate_dim = (0, 2, 1, 3) if batch_size else (1, 0, 2)\n",
    "\n",
    "        # Divide images into patches and flatten\n",
    "        x = x.unfold(2, self.patch_size, self.patch_size).unfold(3, self.patch_size, self.patch_size)\n",
    "        x = x.contiguous().view(*batch_size, self.in_channels, self.N, -1)\n",
    "        x = torch.permute(x, permutate_dim).contiguous().view(*batch_size, self.N, -1)\n",
    "\n",
    "        # Apply linear projection to each patch\n",
    "        x = self.projection(x)\n",
    "\n",
    "        return x"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "outputs": [],
   "source": [
    "class VisionTransformer(nn.Module):\n",
    "    def __init__(self,\n",
    "                 in_channels,\n",
    "                 height,\n",
    "                 width,\n",
    "                 patch_size,\n",
    "                 num_heads,\n",
    "                 D,\n",
    "                 L,\n",
    "                 num_classes,\n",
    "                 mlp_width = 248,\n",
    "                 head_width = 248):\n",
    "        super(VisionTransformer, self).__init__()\n",
    "        self.img_size = (height, width)\n",
    "\n",
    "        # Calculating the correct dimensions:\n",
    "        N = (self.img_size[0] * self.img_size[1]) / (patch_size**2)\n",
    "        assert N % 1 == 0, f\"num_patches must be divisible by patch size: {patch_size}, size: {self.img_size}\"\n",
    "        N = int(N)\n",
    "\n",
    "        # Initialising Extra Learnable [class] Embedding\n",
    "        self.class_embedding = nn.Parameter(torch.randn(1, D))\n",
    "\n",
    "        # Initialising Learnable Position Embedding\n",
    "        self.position_embedding = nn.Parameter(torch.randn(N+1, D))\n",
    "\n",
    "        # ----- Building the Model ----------\n",
    "        # The operations follow in this order\n",
    "        self.linear_projection_of_flattened_patches = LinearProjectionOfFlattenedPatches(\n",
    "            in_channels, height, width, patch_size, D)\n",
    "\n",
    "        transformer_encoder_list = nn.ModuleList()\n",
    "        for _ in range(L): # Adding the TransformerEncoder block L times\n",
    "            transformer_encoder_list.append(TransformerEncoder(num_heads, D, mlp_width))\n",
    "\n",
    "        self.transformer_encoder = nn.Sequential(*transformer_encoder_list) # Unpacking L blocks into one sequential block\n",
    "\n",
    "        # Only the class token will be passed to the head, which will be concatenated at position 0 of our image_patches\n",
    "        self.head = nn.Sequential(\n",
    "            nn.Linear(D, head_width),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(head_width, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, X):\n",
    "        batch_size = X.shape[:-3]\n",
    "        image_patches = self.linear_projection_of_flattened_patches(X)\n",
    "        # Concatenating with the [class] token\n",
    "        class_embedding = self.class_embedding.unsqueeze(0).repeat(*batch_size, 1, 1) if batch_size else self.class_embedding\n",
    "        image_patches = torch.cat((class_embedding, image_patches), -2)\n",
    "        image_patches = image_patches + self.position_embedding # Adding positional embedding\n",
    "        image_patches = self.transformer_encoder(image_patches)\n",
    "        # Extract the class token as we're only using a class token for the classification\n",
    "        class_token = image_patches[:,0,:] if batch_size else image_patches[0, :]\n",
    "        output = self.head(class_token)\n",
    "        return output"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# CNN Model"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "outputs": [],
   "source": [
    "class CNN(nn.Module):\n",
    "    def __init__(self, in_channels, num_classes):\n",
    "        super(CNN, self).__init__()\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=1, padding=1), # 28x28\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2, 0), #14x14\n",
    "            nn.Conv2d(64, 32, 3, 1, 1),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.ReLU(),\n",
    "            nn.MaxPool2d(2, 2, 0), #7x7\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(32*7*7, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, X):\n",
    "        return self.model(X)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loading the Data"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "from medmnist import DermaMNIST\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "train_transformations = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomRotation(10),\n",
    "    transforms.ColorJitter(brightness=0.2),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(0.5, 0.5),\n",
    "])\n",
    "\n",
    "\n",
    "val_transformations = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(.5, .5),\n",
    "])\n",
    "\n",
    "\n",
    "train_download = DermaMNIST(split='train', transform=train_transformations, download=True)\n",
    "val_download = DermaMNIST(split='val', transform=val_transformations, download=True)\n",
    "test_download = DermaMNIST(split='test', transform=val_transformations, download=True)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n",
     "is_executing": true
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Model Init"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "outputs": [],
   "source": [
    "channels, height, width = train_download[0][0].shape\n",
    "patch_size = 4\n",
    "num_heads = 8\n",
    "latent_vector = 64\n",
    "transformer_blocks = 2\n",
    "num_classes = 7\n",
    "mlp_width = 248\n",
    "head_width = 248\n",
    "\n",
    "vit = VisionTransformer(\n",
    "    in_channels=channels,\n",
    "    height=height,\n",
    "    width=width,\n",
    "    patch_size=patch_size,\n",
    "    num_heads=num_heads,\n",
    "    D=latent_vector,\n",
    "    L=transformer_blocks,\n",
    "    num_classes=num_classes,\n",
    "    mlp_width=mlp_width,\n",
    "    head_width=head_width\n",
    ")\n",
    "\n",
    "cnn = CNN(in_channels=channels, num_classes=num_classes)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Training Parameters for ViT"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import torch.optim as optim\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "lr = 1e-3\n",
    "num_epochs = 30\n",
    "batch_size = 64\n",
    "weight_decay = 0\n",
    "checkpoint_every_th_epoch = None\n",
    "vit_optimizer = optim.Adam(vit.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "train_data = DataLoader(train_download, batch_size=batch_size, shuffle=True)\n",
    "val_data = DataLoader(val_download, batch_size=batch_size, shuffle=True)\n",
    "test_data = DataLoader(test_download, batch_size=len(test_download), shuffle=False)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# ViT Training Loop"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/30, Train Loss: 0.7271, Val Loss: 0.7325, Val Accuracy: 0.7188\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[112], line 6\u001B[0m\n\u001B[1;32m      4\u001B[0m train_loss \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m0.0\u001B[39m\n\u001B[1;32m      5\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m images, labels \u001B[38;5;129;01min\u001B[39;00m train_data:\n\u001B[0;32m----> 6\u001B[0m     images, labels \u001B[38;5;241m=\u001B[39m images, \u001B[43mlabels\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43msqueeze\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m1\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[1;32m      8\u001B[0m     optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[1;32m     10\u001B[0m     outputs \u001B[38;5;241m=\u001B[39m vit(images)\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    # Training Phase\n",
    "    vit.train()\n",
    "    train_loss = 0.0\n",
    "    for images, labels in train_data:\n",
    "        images, labels = images, labels.squeeze(1)\n",
    "\n",
    "        vit_optimizer.zero_grad()\n",
    "\n",
    "        outputs = vit(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        vit_optimizer.step()\n",
    "\n",
    "        train_loss += loss.item() * images.size(0)\n",
    "\n",
    "    train_loss /= len(train_data.dataset)\n",
    "\n",
    "    # Validation Phase\n",
    "    vit.eval()\n",
    "    val_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in val_data:\n",
    "            images, labels = images, labels.squeeze(1)\n",
    "\n",
    "            outputs = vit(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            val_loss += loss.item() * images.size(0)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    val_loss /= len(val_data.dataset)\n",
    "    val_accuracy = correct / total\n",
    "\n",
    "    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}')\n",
    "\n",
    "    # Save checkpoint\n",
    "    if checkpoint_every_th_epoch:\n",
    "        if (epoch + 1) % checkpoint_every_th_epoch == 0:\n",
    "            torch.save(vit.state_dict(), f'checkpoint_epoch_{epoch+1}.pth')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Training Parameters for CNN"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "lr = 1e-3\n",
    "num_epochs = 30\n",
    "batch_size = 64\n",
    "weight_decay = 0\n",
    "checkpoint_every_th_epoch = None\n",
    "cnn_optimizer = optim.Adam(cnn.parameters(), lr=lr, weight_decay=weight_decay)\n",
    "\n",
    "train_data = DataLoader(train_download, batch_size=batch_size, shuffle=True)\n",
    "val_data = DataLoader(val_download, batch_size=batch_size, shuffle=True)\n",
    "test_data = DataLoader(test_download, batch_size=len(test_download), shuffle=False)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# CNN Training Loop"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/30, Train Loss: 1.0117, Val Loss: 0.7841, Val Accuracy: 0.7099\n",
      "Epoch 2/30, Train Loss: 0.7944, Val Loss: 0.7290, Val Accuracy: 0.7428\n",
      "Epoch 3/30, Train Loss: 0.7456, Val Loss: 0.7284, Val Accuracy: 0.7348\n",
      "Epoch 4/30, Train Loss: 0.7148, Val Loss: 0.8066, Val Accuracy: 0.6670\n",
      "Epoch 5/30, Train Loss: 0.7111, Val Loss: 0.6817, Val Accuracy: 0.7537\n",
      "Epoch 6/30, Train Loss: 0.6854, Val Loss: 0.6805, Val Accuracy: 0.7567\n",
      "Epoch 7/30, Train Loss: 0.6717, Val Loss: 0.6717, Val Accuracy: 0.7637\n",
      "Epoch 8/30, Train Loss: 0.6564, Val Loss: 0.6445, Val Accuracy: 0.7697\n",
      "Epoch 9/30, Train Loss: 0.6427, Val Loss: 0.6447, Val Accuracy: 0.7627\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m                         Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[119], line 10\u001B[0m\n\u001B[1;32m      6\u001B[0m images, labels \u001B[38;5;241m=\u001B[39m images, labels\u001B[38;5;241m.\u001B[39msqueeze(\u001B[38;5;241m1\u001B[39m)\n\u001B[1;32m      8\u001B[0m cnn_optimizer\u001B[38;5;241m.\u001B[39mzero_grad()\n\u001B[0;32m---> 10\u001B[0m outputs \u001B[38;5;241m=\u001B[39m \u001B[43mcnn\u001B[49m\u001B[43m(\u001B[49m\u001B[43mimages\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m     11\u001B[0m loss \u001B[38;5;241m=\u001B[39m criterion(outputs, labels)\n\u001B[1;32m     12\u001B[0m loss\u001B[38;5;241m.\u001B[39mbackward()\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "Cell \u001B[0;32mIn[95], line 22\u001B[0m, in \u001B[0;36mCNN.forward\u001B[0;34m(self, X)\u001B[0m\n\u001B[1;32m     21\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, X):\n\u001B[0;32m---> 22\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mmodel\u001B[49m\u001B[43m(\u001B[49m\u001B[43mX\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/container.py:215\u001B[0m, in \u001B[0;36mSequential.forward\u001B[0;34m(self, input)\u001B[0m\n\u001B[1;32m    213\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, \u001B[38;5;28minput\u001B[39m):\n\u001B[1;32m    214\u001B[0m     \u001B[38;5;28;01mfor\u001B[39;00m module \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m:\n\u001B[0;32m--> 215\u001B[0m         \u001B[38;5;28minput\u001B[39m \u001B[38;5;241m=\u001B[39m \u001B[43mmodule\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[1;32m    216\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28minput\u001B[39m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1518\u001B[0m, in \u001B[0;36mModule._wrapped_call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1516\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_compiled_call_impl(\u001B[38;5;241m*\u001B[39margs, \u001B[38;5;241m*\u001B[39m\u001B[38;5;241m*\u001B[39mkwargs)  \u001B[38;5;66;03m# type: ignore[misc]\u001B[39;00m\n\u001B[1;32m   1517\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m-> 1518\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_call_impl\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/module.py:1527\u001B[0m, in \u001B[0;36mModule._call_impl\u001B[0;34m(self, *args, **kwargs)\u001B[0m\n\u001B[1;32m   1522\u001B[0m \u001B[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001B[39;00m\n\u001B[1;32m   1523\u001B[0m \u001B[38;5;66;03m# this function, and just call forward.\u001B[39;00m\n\u001B[1;32m   1524\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m (\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_forward_pre_hooks\n\u001B[1;32m   1525\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_backward_pre_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_backward_hooks\n\u001B[1;32m   1526\u001B[0m         \u001B[38;5;129;01mor\u001B[39;00m _global_forward_hooks \u001B[38;5;129;01mor\u001B[39;00m _global_forward_pre_hooks):\n\u001B[0;32m-> 1527\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mforward_call\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43margs\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[38;5;241;43m*\u001B[39;49m\u001B[43mkwargs\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m   1529\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m   1530\u001B[0m     result \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mNone\u001B[39;00m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/conv.py:460\u001B[0m, in \u001B[0;36mConv2d.forward\u001B[0;34m(self, input)\u001B[0m\n\u001B[1;32m    459\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mforward\u001B[39m(\u001B[38;5;28mself\u001B[39m, \u001B[38;5;28minput\u001B[39m: Tensor) \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m>\u001B[39m Tensor:\n\u001B[0;32m--> 460\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43m_conv_forward\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mbias\u001B[49m\u001B[43m)\u001B[49m\n",
      "File \u001B[0;32m/opt/homebrew/lib/python3.11/site-packages/torch/nn/modules/conv.py:456\u001B[0m, in \u001B[0;36mConv2d._conv_forward\u001B[0;34m(self, input, weight, bias)\u001B[0m\n\u001B[1;32m    452\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mpadding_mode \u001B[38;5;241m!=\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mzeros\u001B[39m\u001B[38;5;124m'\u001B[39m:\n\u001B[1;32m    453\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m F\u001B[38;5;241m.\u001B[39mconv2d(F\u001B[38;5;241m.\u001B[39mpad(\u001B[38;5;28minput\u001B[39m, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_reversed_padding_repeated_twice, mode\u001B[38;5;241m=\u001B[39m\u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mpadding_mode),\n\u001B[1;32m    454\u001B[0m                     weight, bias, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mstride,\n\u001B[1;32m    455\u001B[0m                     _pair(\u001B[38;5;241m0\u001B[39m), \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mdilation, \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39mgroups)\n\u001B[0;32m--> 456\u001B[0m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[43mF\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mconv2d\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;28;43minput\u001B[39;49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mweight\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mbias\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mstride\u001B[49m\u001B[43m,\u001B[49m\n\u001B[1;32m    457\u001B[0m \u001B[43m                \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mpadding\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mdilation\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;28;43mself\u001B[39;49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mgroups\u001B[49m\u001B[43m)\u001B[49m\n",
      "\u001B[0;31mKeyboardInterrupt\u001B[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    # Training Phase\n",
    "    cnn.train()\n",
    "    train_loss = 0.0\n",
    "    for images, labels in train_data:\n",
    "        images, labels = images, labels.squeeze(1)\n",
    "\n",
    "        cnn_optimizer.zero_grad()\n",
    "\n",
    "        outputs = cnn(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        cnn_optimizer.step()\n",
    "\n",
    "        train_loss += loss.item() * images.size(0)\n",
    "\n",
    "    train_loss /= len(train_data.dataset)\n",
    "\n",
    "    # Validation Phase\n",
    "    cnn.eval()\n",
    "    val_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in val_data:\n",
    "            images, labels = images, labels.squeeze(1)\n",
    "\n",
    "            outputs = cnn(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            val_loss += loss.item() * images.size(0)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    val_loss /= len(val_data.dataset)\n",
    "    val_accuracy = correct / total\n",
    "\n",
    "    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}')\n",
    "\n",
    "    # Save checkpoint\n",
    "    if checkpoint_every_th_epoch:\n",
    "        if (epoch + 1) % checkpoint_every_th_epoch == 0:\n",
    "            torch.save(cnn.state_dict(), f'checkpoint_epoch_{epoch+1}.pth')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Evaluating our Models"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Vision Transformer Performance ---\n",
      "Test Loss: 0.6937, Test Accuracy: 0.7406\n",
      "\n",
      "Class-wise Accuracy:\n",
      "{0: 33.333333333333336, 1: 37.86407766990291, 2: 41.36363636363637, 3: 0.0, 4: 10.762331838565023, 5: 96.42058165548099, 6: 55.172413793103445}\n",
      "\n",
      "--- CNN Performance ---\n",
      "Test Loss: 0.6621, Test Accuracy: 0.7566\n",
      "\n",
      "Class-wise Accuracy:\n",
      "{0: 27.272727272727273, 1: 62.13592233009709, 2: 43.18181818181818, 3: 13.043478260869565, 4: 24.2152466367713, 5: 94.33258762117822, 6: 62.06896551724138}\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "def evaluate_model(model, test_loader, criterion, num_classes):\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    # Initialize a confusion matrix\n",
    "    confusion_matrix = np.zeros((num_classes, num_classes), dtype=np.int64)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images, labels = images, labels.squeeze(1)\n",
    "\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            test_loss += loss.item() * images.size(0)\n",
    "            _, predicted = torch.max(outputs, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "            # Update confusion matrix\n",
    "            for t, p in zip(labels.view(-1), predicted.view(-1)):\n",
    "                confusion_matrix[t.long(), p.long()] += 1\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    test_accuracy = correct / total\n",
    "\n",
    "    # Calculate accuracy for each class\n",
    "    class_accuracy = 100 * confusion_matrix.diagonal() / confusion_matrix.sum(1)\n",
    "    class_accuracy_dict = {idx: acc for idx, acc in enumerate(class_accuracy)}\n",
    "\n",
    "    return test_loss, test_accuracy, class_accuracy_dict\n",
    "\n",
    "\n",
    "\n",
    "# Evaluate Vision Transformer\n",
    "\n",
    "vit_test_loss, vit_test_accuracy, vit_class_accuracy = evaluate_model(vit, test_data, criterion, 7)\n",
    "print(f'--- Vision Transformer Performance ---')\n",
    "print(f'Test Loss: {vit_test_loss:.4f}, Test Accuracy: {vit_test_accuracy:.4f}\\n')\n",
    "print(f'Class-wise Accuracy:\\n{vit_class_accuracy}\\n')\n",
    "\n",
    "# Evaluate CNN\n",
    "print(f'--- CNN Performance ---')\n",
    "cnn_test_loss, cnn_test_accuracy, cnn_class_accuracy = evaluate_model(cnn, test_data, criterion, 7)\n",
    "print(f'Test Loss: {cnn_test_loss:.4f}, Test Accuracy: {cnn_test_accuracy:.4f}\\n')\n",
    "print(f'Class-wise Accuracy:\\n{cnn_class_accuracy}')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "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": 0
}