1_Vision_Transformer.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "name": "08_pytorch_paper_replicating_exercises.ipynb",
   "provenance": [],
   "collapsed_sections": [],
   "toc_visible": true,
   "authorship_tag": "ABX9TyOhoCjGZZxrecbm76R8UJZn",
   "include_colab_link": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3 (ipykernel)",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU",
  "gpuClass": "standard"
 },
 "cells": [
  {
   "cell_type": "code",
   "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)"
   ],
   "metadata": {
    "id": "Y5H5P8EjCNGK",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "4214da9e-f3a8-43e6-a48c-1a33f44a9be9",
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:34.776077Z",
     "start_time": "2024-10-14T18:08:29.992327Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Implementatation"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:34.783105Z",
     "start_time": "2024-10-14T18:08:34.780079Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class PatchEmbedding(nn.Module):\n",
    "    def __init__(self, input_channels, patch_size, embedding_dim):\n",
    "        super().__init__()\n",
    "        self.conv_proj = nn.Conv2d(input_channels, embedding_dim, kernel_size=patch_size, stride=patch_size)\n",
    "        self.flatten = nn.Flatten(start_dim=2, end_dim=3)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv_proj(x)\n",
    "        x = self.flatten(x)\n",
    "        x = x.permute(0, 2, 1)\n",
    "        return x"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:35.164669Z",
     "start_time": "2024-10-14T18:08:35.160531Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, input_dim, embedding_dim):\n",
    "        super().__init__()\n",
    "        self.Q = nn.Linear(input_dim, embedding_dim)\n",
    "        self.K = nn.Linear(input_dim, embedding_dim)\n",
    "        self.V = nn.Linear(input_dim, embedding_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        q = self.Q(x)\n",
    "        k = self.K(x)\n",
    "        v = self.V(x)\n",
    "        \n",
    "        a = (q @ k.permute(0, 2, 1)) / (64 ** 0.5)\n",
    "        x = torch.softmax(a, dim=2) @ v\n",
    "        return x\n",
    "\n",
    "\n",
    "class MultiheadSelfAttention(nn.Module):\n",
    "    def __init__(self, input_dim, n_heads):\n",
    "        super().__init__()\n",
    "        self.attention_heads = nn.ModuleList([SelfAttention(input_dim, input_dim // n_heads) for _ in range(n_heads)])\n",
    "        self.linser = nn.Linear(input_dim, input_dim)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = torch.concat([attention_head(x) for attention_head in self.attention_heads], dim=-1)\n",
    "        x = self.linser(x)\n",
    "        return x"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:35.175156Z",
     "start_time": "2024-10-14T18:08:35.169512Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MSABlock(nn.Module):\n",
    "    def __init__(self, input_dim, n_heads, torch_msa=True):\n",
    "        super().__init__()\n",
    "        self.torch_msa = torch_msa\n",
    "        \n",
    "        self.layer_norm = nn.LayerNorm(input_dim)\n",
    "        self.msa = nn.MultiheadAttention(input_dim, n_heads, batch_first=True) if torch_msa else MultiheadSelfAttention(input_dim, n_heads)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.layer_norm(x)\n",
    "        x = self.msa(x, x, x, need_weights=False)[0] if self.torch_msa else self.msa(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class MLPBlock(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.layer_norm = nn.LayerNorm(input_dim)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(input_dim, mlp_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(mlp_dim, input_dim),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.layer_norm(x)\n",
    "        x = self.mlp(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, n_heads, dropout):\n",
    "        super().__init__()\n",
    "        self.msa_block = MSABlock(input_dim, n_heads)\n",
    "        self.mlp_block = MLPBlock(input_dim, mlp_dim, dropout)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.msa_block(x) + x\n",
    "        x = self.mlp_block(x) + x\n",
    "        return x\n",
    "\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, input_dim, mlp_dim, n_heads, dropout, n_layers):\n",
    "        super().__init__()\n",
    "        self.transformer_layers = nn.Sequential(\n",
    "            *[TransformerBlock(input_dim, mlp_dim, n_heads, dropout) for _ in range(n_layers)]\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.transformer_layers(x)\n",
    "        cls_embedding = x[:, 0, :]\n",
    "        return cls_embedding"
   ],
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:35.184299Z",
     "start_time": "2024-10-14T18:08:35.180155Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class ViT(nn.Module):\n",
    "    def __init__(self, img_size, n_channels, patch_size, embedding_dim, mlp_dim, n_heads, dropout, n_layers, n_classes):\n",
    "        super().__init__()\n",
    "        self.num_patches = (img_size // patch_size) ** 2\n",
    "        self.cls = nn.Parameter(torch.randn(1, 1, embedding_dim))\n",
    "        self.positional_encoding = nn.Parameter(torch.randn(1, self.num_patches + 1, embedding_dim))\n",
    "        self.patch_embedding = PatchEmbedding(n_channels, patch_size, embedding_dim)\n",
    "        \n",
    "        self.transformer_encoder = TransformerEncoder(embedding_dim, mlp_dim, n_heads, dropout, n_layers)\n",
    "        self.classification_head = nn.Sequential(\n",
    "            nn.LayerNorm(embedding_dim),\n",
    "            nn.Linear(embedding_dim, n_classes)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        cls = self.cls.expand(batch_size, -1, -1)\n",
    "        patch_embeddins = self.patch_embedding(x)\n",
    "        patch_position_embeddings = torch.concat([cls, patch_embeddins], dim=1) + self.positional_encoding\n",
    "        cls_embedding = self.transformer_encoder(patch_position_embeddings)\n",
    "        logits = self.classification_head(cls_embedding)\n",
    "        return logits"
   ],
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:39.746131Z",
     "start_time": "2024-10-14T18:08:38.442101Z"
    }
   },
   "cell_type": "code",
   "source": [
    "patch_size = 16\n",
    "embedding_dim = 768\n",
    "mlp_dim = 3072\n",
    "dropout = 0.1\n",
    "n_heads = 12\n",
    "n_layers = 12\n",
    "\n",
    "vit_model = ViT(224, 3, patch_size, embedding_dim, mlp_dim, n_heads, dropout, n_layers, 1000).to(DEVICE)\n",
    "summary(vit_model, input_size=(1, 3, 224, 224), col_names=['input_size', 'output_size', 'num_params'], depth=4, device='cpu')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=============================================================================================================================\n",
       "Layer (type:depth-idx)                             Input Shape               Output Shape              Param #\n",
       "=============================================================================================================================\n",
       "ViT                                                [1, 3, 224, 224]          [1, 1000]                 152,064\n",
       "├─PatchEmbedding: 1-1                              [1, 3, 224, 224]          [1, 196, 768]             --\n",
       "│    └─Conv2d: 2-1                                 [1, 3, 224, 224]          [1, 768, 14, 14]          590,592\n",
       "│    └─Flatten: 2-2                                [1, 768, 14, 14]          [1, 768, 196]             --\n",
       "├─TransformerEncoder: 1-2                          [1, 197, 768]             [1, 768]                  --\n",
       "│    └─Sequential: 2-3                             [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    └─TransformerBlock: 3-1                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-1                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-2                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-2                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-3                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-4                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-3                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-5                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-6                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-4                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-7                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-8                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-5                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-9                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-10                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-6                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-11                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-12                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-7                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-13                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-14                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-8                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-15                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-16                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-9                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-17                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-18                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-10                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-19                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-20                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-11                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-21                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-22                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-12                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-23                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-24                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "├─Sequential: 1-3                                  [1, 768]                  [1, 1000]                 --\n",
       "│    └─LayerNorm: 2-4                              [1, 768]                  [1, 768]                  1,536\n",
       "│    └─Linear: 2-5                                 [1, 768]                  [1, 1000]                 769,000\n",
       "=============================================================================================================================\n",
       "Total params: 86,567,656\n",
       "Trainable params: 86,567,656\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 173.23\n",
       "=============================================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 102.89\n",
       "Params size (MB): 232.27\n",
       "Estimated Total Size (MB): 335.76\n",
       "============================================================================================================================="
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:43.432522Z",
     "start_time": "2024-10-14T18:08:42.345661Z"
    }
   },
   "cell_type": "code",
   "source": [
    "vit_torch_model = models.vit_b_16()\n",
    "summary(vit_model, input_size=(1, 3, 224, 224), col_names=['input_size', 'output_size', 'num_params'], depth=4, device='cpu')"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "=============================================================================================================================\n",
       "Layer (type:depth-idx)                             Input Shape               Output Shape              Param #\n",
       "=============================================================================================================================\n",
       "ViT                                                [1, 3, 224, 224]          [1, 1000]                 152,064\n",
       "├─PatchEmbedding: 1-1                              [1, 3, 224, 224]          [1, 196, 768]             --\n",
       "│    └─Conv2d: 2-1                                 [1, 3, 224, 224]          [1, 768, 14, 14]          590,592\n",
       "│    └─Flatten: 2-2                                [1, 768, 14, 14]          [1, 768, 196]             --\n",
       "├─TransformerEncoder: 1-2                          [1, 197, 768]             [1, 768]                  --\n",
       "│    └─Sequential: 2-3                             [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    └─TransformerBlock: 3-1                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-1                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-2                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-2                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-3                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-4                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-3                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-5                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-6                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-4                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-7                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-8                     [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-5                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-9                     [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-10                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-6                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-11                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-12                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-7                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-13                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-14                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-8                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-15                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-16                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-9                  [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-17                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-18                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-10                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-19                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-20                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-11                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-21                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-22                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "│    │    └─TransformerBlock: 3-12                 [1, 197, 768]             [1, 197, 768]             --\n",
       "│    │    │    └─MSABlock: 4-23                    [1, 197, 768]             [1, 197, 768]             2,363,904\n",
       "│    │    │    └─MLPBlock: 4-24                    [1, 197, 768]             [1, 197, 768]             4,723,968\n",
       "├─Sequential: 1-3                                  [1, 768]                  [1, 1000]                 --\n",
       "│    └─LayerNorm: 2-4                              [1, 768]                  [1, 768]                  1,536\n",
       "│    └─Linear: 2-5                                 [1, 768]                  [1, 1000]                 769,000\n",
       "=============================================================================================================================\n",
       "Total params: 86,567,656\n",
       "Trainable params: 86,567,656\n",
       "Non-trainable params: 0\n",
       "Total mult-adds (Units.MEGABYTES): 173.23\n",
       "=============================================================================================================================\n",
       "Input size (MB): 0.60\n",
       "Forward/backward pass size (MB): 102.89\n",
       "Params size (MB): 232.27\n",
       "Estimated Total Size (MB): 335.76\n",
       "============================================================================================================================="
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# Training"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T18:08:49.448612Z",
     "start_time": "2024-10-14T18:08:48.003099Z"
    }
   },
   "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)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-14T23:11:04.979360Z",
     "start_time": "2024-10-14T18:08:52.361985Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import wandb\n",
    "from src.engine import *\n",
    "\n",
    "config = dict(batch_size=32, lr=5e-4, epochs=20, dataset='CIFAR100')\n",
    "with wandb.init(project='pytorch-study', name='ViT', 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",
    "    vit_model = ViT(224, 3, patch_size, embedding_dim, mlp_dim, n_heads, dropout, n_layers, n_classes).to(DEVICE)\n",
    "    \n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(vit_model.parameters(), lr=w_config.lr)\n",
    "    \n",
    "    loss_history, acc_history = train(vit_model, train_dl, val_dl, criterion, optimizer, w_config.epochs, DEVICE, run) "
   ],
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch=1:   0%|          | 0/20 [00:00<?, ?it/s]C:\\Users\\dk\\.pipenv\\paper_implementation-aJVmDThZ\\Lib\\site-packages\\torch\\nn\\functional.py:5560: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at C:\\actions-runner\\_work\\pytorch\\pytorch\\builder\\windows\\pytorch\\aten\\src\\ATen\\native\\transformers\\cuda\\sdp_utils.cpp:555.)\n",
      "  attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)\n",
      "Epoch=20: 100%|██████████| 20/20 [5:01:59<00:00, 905.99s/it, train_loss=2.396, train_acc=36.48%, val_loss=3.467, val_acc=21.15%]  \n"
     ]
    }
   ],
   "execution_count": 9
  }
 ]
}