Transformer.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jd3k3cSF_Fgw",
        "outputId": "70329f22-5df0-4586-807c-32a5694fcfd6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting timm\n",
            "  Downloading timm-1.0.9-py3-none-any.whl.metadata (42 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/42.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.4/42.4 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from timm) (2.4.1+cu121)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from timm) (0.19.1+cu121)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.10/dist-packages (from timm) (6.0.2)\n",
            "Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.10/dist-packages (from timm) (0.24.7)\n",
            "Requirement already satisfied: safetensors in /usr/local/lib/python3.10/dist-packages (from timm) (0.4.5)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (3.16.1)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (2024.6.1)\n",
            "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (24.1)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (4.66.5)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface_hub->timm) (4.12.2)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->timm) (1.13.3)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->timm) (3.3)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->timm) (3.1.4)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torchvision->timm) (1.26.4)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.10/dist-packages (from torchvision->timm) (10.4.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->timm) (2.1.5)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub->timm) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub->timm) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub->timm) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface_hub->timm) (2024.8.30)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->timm) (1.3.0)\n",
            "Downloading timm-1.0.9-py3-none-any.whl (2.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m36.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: timm\n",
            "Successfully installed timm-1.0.9\n"
          ]
        }
      ],
      "source": [
        "!pip install timm\n",
        "from timm.models.layers import to_2tuple, trunc_normal_, DropPath\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from functools import partial\n"
      ],
      "metadata": {
        "id": "T7z4JQnRJLPD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class PatchEmbed(nn.Module):\n",
        "    def __init__(self, img_size=224, patch_size=16, in_chan=3, embed_dim=768, multi_conv=False, use_simple_conv=False):\n",
        "        super().__init__()\n",
        "        img_size = to_2tuple(img_size)\n",
        "        patch_size = to_2tuple(patch_size)\n",
        "        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])\n",
        "        self.img_size = img_size\n",
        "        self.patch_size = patch_size\n",
        "        self.num_patches = num_patches\n",
        "\n",
        "        if multi_conv and not use_simple_conv:\n",
        "            if patch_size[0] == 24:\n",
        "                self.proj = nn.Sequential(\n",
        "                    nn.Conv2d(in_chan, embed_dim // 4, kernel_size=7, stride=3, padding=2),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 4, embed_dim // 2, kernel_size=3, stride=2, padding=1),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 2, embed_dim, kernel_size=3, stride=1, padding=1),\n",
        "                )\n",
        "            elif patch_size[0] == 12:\n",
        "                self.proj = nn.Sequential(\n",
        "                    nn.Conv2d(in_chan, embed_dim // 4, kernel_size=7, stride=4, padding=3),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 4, embed_dim // 2, kernel_size=3, stride=2, padding=1),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 2, embed_dim, kernel_size=3, stride=1, padding=1),\n",
        "                )\n",
        "            elif patch_size[0] == 4:\n",
        "                self.proj = nn.Sequential(\n",
        "                    nn.Conv2d(in_chan, embed_dim // 4, kernel_size=7, stride=2, padding=3),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 4, embed_dim // 2, kernel_size=3, stride=2, padding=1),\n",
        "                    nn.ReLU(inplace=True),\n",
        "                    nn.Conv2d(embed_dim // 2, embed_dim, kernel_size=3, stride=1, padding=1),\n",
        "                )\n",
        "            else:\n",
        "                raise ValueError(f\"Unsupported patch size {patch_size[0]}\")\n",
        "        elif use_simple_conv:\n",
        "            # Simple convolutional embedding\n",
        "            self.proj = nn.Conv2d(in_chan, embed_dim, kernel_size=patch_size, stride=patch_size)\n",
        "        else:\n",
        "            # Default behavior if multi_conv is not used\n",
        "            self.proj = nn.Conv2d(in_chan, embed_dim, kernel_size=patch_size, stride=patch_size)\n",
        "\n",
        "    def forward(self, x, extra_padding=False):\n",
        "        B, C, H, W = x.shape\n",
        "        if extra_padding and (H % self.patch_size[0] != 0 or W % self.patch_size[1] != 0):\n",
        "            p_l = (self.patch_size[1] - W % self.patch_size[1]) // 2\n",
        "            p_r = (self.patch_size[1] - W % self.patch_size[1]) - p_l\n",
        "            p_t = (self.patch_size[0] - H % self.patch_size[0]) // 2\n",
        "            p_b = (self.patch_size[0] - H % self.patch_size[0]) - p_t\n",
        "            x = F.pad(x, (p_l, p_r, p_t, p_b))\n",
        "\n",
        "        x = self.proj(x)  # Apply the projection (whether multi-conv or simple-conv)\n",
        "        x = x.flatten(2).transpose(1, 2)  # Flatten and transpose for transformer input\n",
        "\n",
        "        return x\n",
        "\n",
        "# Utility function to convert values to tuple if they aren't already\n",
        "def to_2tuple(x):\n",
        "    if isinstance(x, tuple):\n",
        "        return x\n",
        "    return (x, x)\n",
        "\n",
        "class CrossAttention(nn.Module):\n",
        "    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., patch_size=16):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        head_dim = dim // num_heads\n",
        "        self.scale = qk_scale or head_dim ** -0.5\n",
        "\n",
        "        self.wq = nn.Linear(dim, dim, bias=qkv_bias)\n",
        "        self.wk = nn.Linear(dim, dim, bias=qkv_bias)\n",
        "        self.wv = nn.Linear(dim, dim, bias=qkv_bias)\n",
        "        self.attn_drop = nn.Dropout(attn_drop)\n",
        "        self.proj = nn.Linear(dim, dim)\n",
        "        self.proj_drop = nn.Dropout(proj_drop)\n",
        "        self.patch_size = patch_size\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, N, C = x.shape\n",
        "\n",
        "        # Compute Q, K, V matrices\n",
        "        q = self.wq(x[:, 0:1, ...]).reshape(B, 1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)\n",
        "        k = self.wk(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)\n",
        "        v = self.wv(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)\n",
        "\n",
        "        # Compute attention scores\n",
        "        attn = (q @ k.transpose(-2, -1)) * self.scale\n",
        "        attn = attn.softmax(dim=-1)\n",
        "        attn = self.attn_drop(attn)\n",
        "\n",
        "        # Compute attention output\n",
        "        x = (attn @ v).transpose(1, 2).reshape(B, 1, C)\n",
        "        x = self.proj(x)\n",
        "        x = self.proj_drop(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "class CrossAttentionBlock(nn.Module):\n",
        "    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,\n",
        "                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, has_mlp=True, patch_size=16):\n",
        "        super().__init__()\n",
        "        self.norm1 = norm_layer(dim)\n",
        "        self.attn = CrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
        "                                   attn_drop=attn_drop, proj_drop=drop, patch_size=patch_size)\n",
        "        self.drop_path = nn.Identity() if drop_path == 0 else nn.Dropout(drop_path)\n",
        "        self.has_mlp = has_mlp\n",
        "        if self.has_mlp:\n",
        "            self.norm2 = norm_layer(dim)\n",
        "            mlp_hidden_dim = int(dim * mlp_ratio)\n",
        "            self.mlp = nn.Sequential(\n",
        "                nn.Linear(dim, mlp_hidden_dim),\n",
        "                act_layer(),\n",
        "                nn.Dropout(drop),\n",
        "                nn.Linear(mlp_hidden_dim, dim),\n",
        "                nn.Dropout(drop),\n",
        "            )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x[:, 0:1, ...] + self.drop_path(self.attn(self.norm1(x)))\n",
        "        if self.has_mlp:\n",
        "            x = x + self.drop_path(self.mlp(self.norm2(x)))\n",
        "        return x\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):\n",
        "        super().__init__()\n",
        "        self.norm1 = norm_layer(dim)\n",
        "        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=num_heads, dropout=attn_drop)\n",
        "        self.drop_path = nn.Identity() if drop_path == 0 else nn.Dropout(drop_path)\n",
        "        self.norm2 = norm_layer(dim)\n",
        "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
        "        self.mlp = nn.Sequential(\n",
        "            nn.Linear(dim, mlp_hidden_dim),\n",
        "            act_layer(),\n",
        "            nn.Dropout(drop),\n",
        "            nn.Linear(mlp_hidden_dim, dim),\n",
        "            nn.Dropout(drop),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.drop_path(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])\n",
        "        x = x + self.drop_path(self.mlp(self.norm2(x)))\n",
        "        return x\n",
        "\n",
        "class MultiScaleBlock(nn.Module):\n",
        "    def __init__(self, dim, patches, depth, num_heads, mlp_ratio, act_layer=nn.GELU, qkv_bias=False,\n",
        "                 qk_scale=None, attn_drop=0., drop=0., drop_path=0 , norm_layer=nn.LayerNorm):\n",
        "        super().__init__()\n",
        "\n",
        "        num_branches = len(dim)\n",
        "        self.num_branches = num_branches\n",
        "\n",
        "        self.blocks = nn.ModuleList()\n",
        "        for d in range(num_branches):\n",
        "            temp = []\n",
        "            for i in range(depth[d]):\n",
        "                temp.append(\n",
        "                    Block(dim[d], num_heads[d], mlp_ratio[d], qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
        "                          drop=drop, attn_drop=attn_drop, drop_path=drop_path, norm_layer=norm_layer)\n",
        "                )\n",
        "            self.blocks.append(nn.ModuleList(temp))\n",
        "\n",
        "        self.proj = nn.ModuleList()\n",
        "        self.patches = patches\n",
        "\n",
        "        for i in range(len(dim)):\n",
        "            self.proj.append(nn.Linear(dim[i], dim[i + 1]) if i < len(dim) - 1 else nn.Identity())\n",
        "\n",
        "        self.proj.append(nn.Linear(dim[1], dim[0]))\n",
        "\n",
        "    def forward(self, x):\n",
        "        cross_input = x\n",
        "        for i, block in enumerate(self.blocks):\n",
        "            for layer in block:\n",
        "                x = layer(x)\n",
        "\n",
        "            if i + 1 != self.num_branches:\n",
        "                cross_input += self.proj[i](x)\n",
        "\n",
        "        return cross_input\n",
        "\n",
        "\n",
        "class VisionTransformer(nn.Module):\n",
        "    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,\n",
        "                 num_heads=12, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n",
        "                 drop_path_rate=0., norm_layer=nn.LayerNorm, act_layer=nn.GELU):\n",
        "        super().__init__()\n",
        "        self.patch_size = patch_size\n",
        "        self.num_classes = num_classes\n",
        "\n",
        "        self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chan=in_chans, embed_dim=embed_dim)\n",
        "\n",
        "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n",
        "        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim))\n",
        "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
        "\n",
        "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule\n",
        "        self.blocks = nn.ModuleList([\n",
        "            MultiScaleBlock([embed_dim] * 2, patches=[1] * 2, depth=[depth] * 2, num_heads=[num_heads] * 2,\n",
        "                            mlp_ratio=[mlp_ratio] * 2, norm_layer=norm_layer)\n",
        "            for _ in range(depth)\n",
        "        ])\n",
        "        self.norm = norm_layer(embed_dim)\n",
        "\n",
        "        # Classifier head\n",
        "        self.head = nn.Linear(embed_dim, num_classes)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B = x.shape[0]\n",
        "        x = self.patch_embed(x)\n",
        "\n",
        "        cls_tokens = self.cls_token.expand(B, -1, -1)\n",
        "        x = torch.cat((cls_tokens, x), dim=1)\n",
        "        x = self.pos_drop(x + self.pos_embed)\n",
        "\n",
        "        for block in self.blocks:\n",
        "            x = block(x)\n",
        "\n",
        "        x = self.norm(x)\n",
        "        x = self.head(x[:, 0])\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "Yb20D6FzRqyv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch\n",
        "from torchvision import transforms, datasets\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "from PIL import Image\n",
        "# Define the paths to your datasets\n",
        "data_dir = \"/content/gdrive/MyDrive/Colab Notebooks/PlantCLEF_Subset\"\n",
        "train_dir = \"/content/gdrive/MyDrive/Colab Notebooks/PlantCLEF_Subset/train\"\n",
        "\n",
        "# Define the custom dataset and transformations\n",
        "class TestDataset(Dataset):\n",
        "    def __init__(self, root, transform=None):\n",
        "        self.root = root\n",
        "        self.transform = transform\n",
        "        self.image_paths = [os.path.join(root, fname) for fname in os.listdir(root)]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.image_paths)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        img_path = self.image_paths[idx]\n",
        "        image = Image.open(img_path).convert(\"RGB\")\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        return image, -1  # return a dummy label since the test set doesn't have labels\n",
        "\n",
        "train_transform = transforms.Compose([\n",
        "    transforms.Resize((224, 224)),\n",
        "    transforms.RandomHorizontalFlip(),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
        "])\n",
        "\n",
        "val_test_transform = transforms.Compose([\n",
        "    transforms.Resize((224, 224)),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
        "])\n",
        "\n",
        "train_dataset = datasets.ImageFolder(root=train_dir, transform=train_transform)\n",
        "val_dataset = datasets.ImageFolder(root=os.path.join(data_dir, 'val'), transform=val_test_transform)\n",
        "test_dataset = TestDataset(root=os.path.join(data_dir, 'test'), transform=val_test_transform)\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)\n",
        "val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)\n",
        "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)\n",
        "\n",
        "# Training and validation functions\n",
        "def train_one_epoch(model, data_loader, criterion, optimizer, device):\n",
        "    model.train()\n",
        "    running_loss = 0.0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "\n",
        "    for images, labels in data_loader:\n",
        "        images, labels = images.to(device), labels.to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(images)\n",
        "        loss = criterion(outputs, labels)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        running_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",
        "    epoch_loss = running_loss / len(data_loader.dataset)\n",
        "    epoch_accuracy = correct / total\n",
        "    return epoch_loss, epoch_accuracy\n",
        "\n",
        "def validate(model, data_loader, criterion, device):\n",
        "    model.eval()\n",
        "    running_loss = 0.0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for images, labels in data_loader:\n",
        "            images, labels = images.to(device), labels.to(device)\n",
        "\n",
        "            outputs = model(images)\n",
        "            loss = criterion(outputs, labels)\n",
        "\n",
        "            running_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",
        "    epoch_loss = running_loss / len(data_loader.dataset)\n",
        "    epoch_accuracy = correct / total\n",
        "    return epoch_loss, epoch_accuracy\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w5uxzEFd_LXc",
        "outputId": "1293cc29-9ef8-445d-ce16-e16e794dd372"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:557: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
            "  warnings.warn(_create_warning_msg(\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.optim as optim\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "model = VisionTransformer(\n",
        "    img_size=224, patch_size=16, in_chans=3, num_classes=len(train_dataset.classes),\n",
        "    embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.\n",
        ").to(device)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
        "num_epochs=2\n",
        "# Training Loop\n",
        "for epoch in range(num_epochs):\n",
        "    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)\n",
        "\n",
        "    val_loss, val_acc = validate(model, val_loader, criterion, device)\n",
        "\n",
        "    print(f\"Epoch {epoch+1}/{num_epochs}\")\n",
        "    print(f\"Train Loss: {train_loss:.4f}, Train Accuracy: {train_acc:.4f}\")\n",
        "    print(f\"Val Loss: {val_loss:.4f}, Val Accuracy: {val_acc:.4f}\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 408
        },
        "id": "1KKmywaN_OOJ",
        "outputId": "ca5b32c6-f37d-4ae6-f372-7504397667c2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "error",
          "ename": "OutOfMemoryError",
          "evalue": "CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacity of 14.75 GiB of which 19.06 MiB is free. Process 6848 has 14.73 GiB memory in use. Of the allocated memory 13.46 GiB is allocated by PyTorch, and 1.15 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-5-04411b789068>\u001b[0m in \u001b[0;36m<cell line: 13>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;31m# Training Loop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnum_epochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m     \u001b[0mtrain_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_one_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m     \u001b[0mval_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvalidate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-4-50edbd3c68c4>\u001b[0m in \u001b[0;36mtrain_one_epoch\u001b[0;34m(model, data_loader, criterion, optimizer, device)\u001b[0m\n\u001b[1;32m     57\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     58\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 59\u001b[0;31m         \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimages\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     60\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     61\u001b[0m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1560\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1561\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1563\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1564\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-2-b27f7c6f3aef>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    221\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    222\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mblock\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mblocks\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 223\u001b[0;31m             \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mblock\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    224\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    225\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1560\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1561\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1563\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1564\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-2-b27f7c6f3aef>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    179\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mblock\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mblocks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    180\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mlayer\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mblock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 181\u001b[0;31m                 \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlayer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    182\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    183\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mi\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m1\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_branches\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1560\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1561\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1563\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1564\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-2-b27f7c6f3aef>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    145\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    146\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdrop_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdrop_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmlp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1560\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1561\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1562\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1563\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1564\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/activation.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, query, key, value, key_padding_mask, need_weights, attn_mask, average_attn_weights, is_causal)\u001b[0m\n\u001b[1;32m   1273\u001b[0m                 is_causal=is_causal)\n\u001b[1;32m   1274\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1275\u001b[0;31m             attn_output, attn_output_weights = F.multi_head_attention_forward(\n\u001b[0m\u001b[1;32m   1276\u001b[0m                 \u001b[0mquery\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membed_dim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_heads\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1277\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_proj_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0min_proj_bias\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mmulti_head_attention_forward\u001b[0;34m(query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight, out_proj_bias, training, key_padding_mask, need_weights, attn_mask, use_separate_proj_weight, q_proj_weight, k_proj_weight, v_proj_weight, static_k, static_v, average_attn_weights, is_causal)\u001b[0m\n\u001b[1;32m   5418\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0muse_separate_proj_weight\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5419\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0min_proj_weight\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"use_separate_proj_weight is False but in_proj_weight is None\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 5420\u001b[0;31m         \u001b[0mq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_in_projection_packed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mquery\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_proj_weight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_proj_bias\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   5421\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5422\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mq_proj_weight\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"use_separate_proj_weight is True but q_proj_weight is None\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36m_in_projection_packed\u001b[0;34m(q, k, v, w, b)\u001b[0m\n\u001b[1;32m   4940\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4941\u001b[0m             \u001b[0mb_q\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_k\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_v\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchunk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 4942\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mw_q\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_q\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mw_k\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_k\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlinear\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mw_v\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mb_v\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   4943\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   4944\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacity of 14.75 GiB of which 19.06 MiB is free. Process 6848 has 14.73 GiB memory in use. Of the allocated memory 13.46 GiB is allocated by PyTorch, and 1.15 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install einops"
      ],
      "metadata": {
        "id": "uyehRnZ8_UZ1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b1079fef-be03-467f-d538-00d371e1d0b4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: einops in /usr/local/lib/python3.10/dist-packages (0.8.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from einops import rearrange\n",
        "\n",
        "class R2LAttentionPlusFFN(nn.Module):\n",
        "    def __init__(self, input_channels, output_channels, kernel_size, num_heads, mlp_ratio=1., qkv_bias=False, qk_scale=None,\n",
        "                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, drop_path=0., attn_drop=0., drop=0.,\n",
        "                 cls_attn=True):\n",
        "        super().__init__()\n",
        "\n",
        "        if not isinstance(kernel_size, (tuple, list)):\n",
        "            kernel_size = [(kernel_size, kernel_size), (kernel_size, kernel_size), 0]\n",
        "        self.kernel_size = kernel_size\n",
        "\n",
        "        if cls_attn:\n",
        "            self.norm0 = norm_layer(input_channels)\n",
        "        else:\n",
        "            self.norm0 = None\n",
        "\n",
        "        self.norm1 = norm_layer(input_channels)\n",
        "        self.attn = AttentionWithRelPos(\n",
        "            input_channels, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,\n",
        "            attn_map_dim=(kernel_size[0][0], kernel_size[0][1]), num_cls_tokens=1)\n",
        "        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here\n",
        "        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n",
        "        self.norm2 = norm_layer(input_channels)\n",
        "        self.mlp = Mlp(in_features=input_channels, hidden_features=int(output_channels * mlp_ratio), out_features=output_channels, act_layer=act_layer, drop=drop)\n",
        "\n",
        "        self.expand = nn.Sequential(\n",
        "            norm_layer(input_channels),\n",
        "            act_layer(),\n",
        "            nn.Linear(input_channels, output_channels)\n",
        "        ) if input_channels != output_channels else None\n",
        "        self.output_channels = output_channels\n",
        "        self.input_channels = input_channels\n",
        "\n",
        "    def forward(self, xs):\n",
        "        # xs is expected to be a tuple of (out, B, H, W, mask, additional_patch_tokens)\n",
        "        out, B, H, W, mask, additional_patch_tokens = xs\n",
        "\n",
        "        # Process the output from Cross ViT model\n",
        "        cls_tokens = out[:, 0:1, ...]\n",
        "        C = cls_tokens.shape[-1]\n",
        "        cls_tokens = cls_tokens.reshape(B, -1, C)  # (B)x(H/sxW/s)xC\n",
        "\n",
        "        if self.norm0 is not None:\n",
        "            cls_tokens = cls_tokens + self.drop_path(self.attn(self.norm0(cls_tokens)))  # (B)x(H/sxW/s)xC\n",
        "\n",
        "        # Reshape class tokens for concatenation\n",
        "        cls_tokens = cls_tokens.reshape(-1, 1, C)  # (BxH/sxW/s)x1xC\n",
        "\n",
        "        # Concatenate the processed Cross ViT output with additional 4x4 patch tokens\n",
        "        out = torch.cat((cls_tokens, out[:, 1:, ...]), dim=1)\n",
        "\n",
        "        # Add additional 4x4 patch tokens to the output\n",
        "        out = torch.cat((out, additional_patch_tokens), dim=1)\n",
        "\n",
        "        # Apply attention and feed-forward network on concatenated output\n",
        "        tmp = out\n",
        "        tmp = tmp + self.drop_path(self.attn(self.norm1(tmp), patch_attn=True, mask=mask))\n",
        "        identity = self.expand(tmp) if self.expand is not None else tmp\n",
        "        tmp = identity + self.drop_path(self.mlp(self.norm2(tmp)))\n",
        "\n",
        "        return tmp\n",
        "\n",
        "class AttentionWithRelPos(nn.Module):\n",
        "    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.,\n",
        "                 attn_map_dim=None, num_cls_tokens=1):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        head_dim = dim // num_heads\n",
        "        self.scale = qk_scale or head_dim ** -0.5\n",
        "\n",
        "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
        "        self.attn_drop = nn.Dropout(attn_drop)\n",
        "        self.proj = nn.Linear(dim, dim)\n",
        "        self.proj_drop = nn.Dropout(proj_drop)\n",
        "\n",
        "        self.num_cls_tokens = num_cls_tokens\n",
        "        if attn_map_dim is not None:\n",
        "            one_dim = attn_map_dim[0]\n",
        "            rel_pos_dim = (2 * one_dim - 1)\n",
        "            self.rel_pos = nn.Parameter(torch.zeros(num_heads, rel_pos_dim ** 2))\n",
        "            tmp = torch.arange(rel_pos_dim ** 2).reshape((rel_pos_dim, rel_pos_dim))\n",
        "            out = []\n",
        "            offset_x = offset_y = one_dim // 2\n",
        "            for y in range(one_dim):\n",
        "                for x in range(one_dim):\n",
        "                    for dy in range(one_dim):\n",
        "                        for dx in range(one_dim):\n",
        "                            out.append(tmp[dy - y + offset_y, dx - x + offset_x])\n",
        "            self.rel_pos_index = torch.tensor(out, dtype=torch.long)\n",
        "            trunc_normal_(self.rel_pos, std=.02)\n",
        "        else:\n",
        "            self.rel_pos = None\n",
        "\n",
        "    def forward(self, x, patch_attn=False, mask=None):\n",
        "        B, N, C = x.shape\n",
        "        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
        "        q, k, v = qkv[0], qkv[1], qkv[2]\n",
        "\n",
        "        attn = (q @ k.transpose(-2, -1)) * self.scale\n",
        "\n",
        "        if self.rel_pos is not None and patch_attn:\n",
        "            # Apply relative positional encoding\n",
        "            rel_pos = self.rel_pos[:, self.rel_pos_index.to(attn.device)].reshape(self.num_heads, N - self.num_cls_tokens, N - self.num_cls_tokens)\n",
        "            attn[:, :, self.num_cls_tokens:, self.num_cls_tokens:] = attn[:, :, self.num_cls_tokens:, self.num_cls_tokens:] + rel_pos\n",
        "\n",
        "        if mask is not None:\n",
        "            mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)\n",
        "            attn = attn.masked_fill(mask == 0, torch.finfo(attn.dtype).min)\n",
        "\n",
        "        attn = attn.softmax(dim=-1)\n",
        "        attn = self.attn_drop(attn)\n",
        "\n",
        "        x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n",
        "        x = self.proj(x)\n",
        "        x = self.proj_drop(x)\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "bAEn5icxH15w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class ConvAttBlock(nn.Module):\n",
        "    def __init__(self, input_channels, output_channels, kernel_size, num_blocks, num_heads, mlp_ratio=1., qkv_bias=False, qk_scale=None, pool='sc',\n",
        "                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, drop_path_rate=(0.,), attn_drop_rate=0., drop_rate=0.,\n",
        "                 cls_attn=True, peg=False):\n",
        "        super().__init__()\n",
        "        tmp = []\n",
        "        if pool:\n",
        "            tmp.append(Projection(input_channels, output_channels, act_layer=act_layer, mode=pool))\n",
        "\n",
        "        for i in range(num_blocks):\n",
        "            kernel_size_ = kernel_size\n",
        "            tmp.append(R2LAttentionPlusFFN(output_channels, output_channels, kernel_size_, num_heads, mlp_ratio, qkv_bias, qk_scale,\n",
        "                                           act_layer=act_layer, norm_layer=norm_layer, drop_path=drop_path_rate[i], attn_drop=attn_drop_rate, drop=drop_rate,\n",
        "                                           cls_attn=cls_attn))\n",
        "\n",
        "        self.block = nn.ModuleList(tmp)\n",
        "        self.output_channels = output_channels\n",
        "        self.ws = kernel_size\n",
        "        if not isinstance(kernel_size, (tuple, list)):\n",
        "            kernel_size = [(kernel_size, kernel_size), (kernel_size, kernel_size), 0]\n",
        "        self.kernel_size = kernel_size\n",
        "\n",
        "        self.peg = nn.Conv2d(output_channels, output_channels, kernel_size=3, padding=1, groups=output_channels, bias=False) if peg else None\n",
        "\n",
        "    def forward(self, xs):\n",
        "        cls_tokens, patch_tokens = xs\n",
        "        cls_tokens, patch_tokens = self.block[0]((cls_tokens, patch_tokens))\n",
        "        out, mask, p_l, p_r, p_t, p_b, B, C, H, W = convert_to_flatten_layout(cls_tokens, patch_tokens, self.ws)\n",
        "        for i in range(1, len(self.block)):\n",
        "            blk = self.block[i]\n",
        "\n",
        "            out = blk((out, B, H, W, mask))\n",
        "            if self.peg is not None and i == 1:\n",
        "                cls_tokens, patch_tokens = convert_to_spatial_layout(out, self.output_channels, B, H, W, self.kernel_size, mask, p_l, p_r, p_t, p_b)\n",
        "                cls_tokens = cls_tokens + self.peg(cls_tokens)\n",
        "                patch_tokens = patch_tokens + self.peg(patch_tokens)\n",
        "                out, mask, p_l, p_r, p_t, p_b, B, C, H, W = convert_to_flatten_layout(cls_tokens, patch_tokens, self.ws)\n",
        "\n",
        "        cls_tokens, patch_tokens = convert_to_spatial_layout(out, self.output_channels, B, H, W, self.kernel_size, mask, p_l, p_r, p_t, p_b)\n",
        "        return cls_tokens, patch_tokens\n"
      ],
      "metadata": {
        "id": "XrtTevviH_ug"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Projection(nn.Module):\n",
        "    def __init__(self, input_channels, output_channels, act_layer, mode='sc'):\n",
        "        super().__init__()\n",
        "        tmp = []\n",
        "        #adds stride and padding to the kernel according to the description\n",
        "        if 'c' in mode:\n",
        "            ks = 2 if 's' in mode else 1\n",
        "            if ks == 2:\n",
        "                stride = ks\n",
        "                ks = ks + 1\n",
        "                padding = ks // 2\n",
        "            else:\n",
        "                stride = ks\n",
        "                padding = 0\n",
        "\n",
        "            if input_channels == output_channels and ks == 1:\n",
        "                tmp.append(nn.Identity())\n",
        "            else:\n",
        "                tmp.extend([\n",
        "                    LayerNorm2d(input_channels),\n",
        "                    act_layer(),\n",
        "                ])\n",
        "                tmp.append(nn.Conv2d(in_channels=input_channels, out_channels=output_channels, kernel_size=ks, stride=stride, padding=padding, groups=input_channels))\n",
        "\n",
        "        self.proj = nn.Sequential(*tmp)\n",
        "        self.proj_cls = self.proj\n",
        "\n",
        "    def forward(self, xs):\n",
        "        cls_tokens, patch_tokens = xs\n",
        "        # x: BxCxHxW\n",
        "        cls_tokens = self.proj_cls(cls_tokens)\n",
        "        patch_tokens = self.proj(patch_tokens)\n",
        "        return cls_tokens, patch_tokens\n",
        "\n",
        "\n",
        "def convert_to_flatten_layout(cls_tokens, patch_tokens, ws):\n",
        "    # padding if needed, and all paddings are happened at bottom and right.\n",
        "    B, C, H, W = patch_tokens.shape\n",
        "    _, _, H_ks, W_ks = cls_tokens.shape\n",
        "    need_mask = False\n",
        "    p_l, p_r, p_t, p_b = 0, 0, 0, 0\n",
        "    if H % (H_ks * ws) != 0 or W % (W_ks * ws) != 0:\n",
        "        p_l, p_r = 0, W_ks * ws - W\n",
        "        p_t, p_b = 0, H_ks * ws - H\n",
        "        patch_tokens = F.pad(patch_tokens, (p_l, p_r, p_t, p_b))\n",
        "        need_mask = True\n",
        "\n",
        "    B, C, H, W = patch_tokens.shape\n",
        "    kernel_size = (H // H_ks, W // W_ks)\n",
        "    tmp = F.unfold(patch_tokens, kernel_size=kernel_size, stride=kernel_size, padding=(0, 0))\n",
        "    patch_tokens = tmp.transpose(1, 2).reshape(-1, C, kernel_size[0] * kernel_size[1]).transpose(-2, -1)\n",
        "    # If the image is not in correct size then we add 0's to make it same and we have to make sure that the kernel does not use those values so this\n",
        "    #makes all the positions where 0's were added to 0 in keernel as well.\n",
        "    if need_mask:\n",
        "        BH_sK_s, ksks, C = patch_tokens.shape\n",
        "        H_s, W_s = H // ws, W // ws\n",
        "        mask = torch.ones(BH_sK_s // B, 1 + ksks, 1 + ksks, device=patch_tokens.device, dtype=torch.float)\n",
        "        right = torch.zeros(1 + ksks, 1 + ksks, device=patch_tokens.device, dtype=torch.float)\n",
        "        tmp = torch.zeros(ws, ws, device=patch_tokens.device, dtype=torch.float)\n",
        "        tmp[0:(ws - p_r), 0:(ws - p_r)] = 1.\n",
        "        tmp = tmp.repeat(ws, ws)\n",
        "        right[1:, 1:] = tmp\n",
        "        right[0, 0] = 1\n",
        "        right[0, 1:] = torch.tensor([1.] * (ws - p_r) + [0.] * p_r).repeat(ws).to(right.device)\n",
        "        right[1:, 0] = torch.tensor([1.] * (ws - p_r) + [0.] * p_r).repeat(ws).to(right.device)\n",
        "        bottom = torch.zeros_like(right)\n",
        "        bottom[0:ws * (ws - p_b) + 1, 0:ws * (ws - p_b) + 1] = 1.\n",
        "        bottom_right = copy.deepcopy(right)\n",
        "        bottom_right[0:ws * (ws - p_b) + 1, 0:ws * (ws - p_b) + 1] = 1.\n",
        "\n",
        "        mask[W_s - 1:(H_s - 1) * W_s:W_s, ...] = right\n",
        "        mask[(H_s - 1) * W_s:, ...] = bottom\n",
        "        mask[-1, ...] = bottom_right\n",
        "        mask = mask.repeat(B, 1, 1)\n",
        "    else:\n",
        "        mask = None\n",
        "\n",
        "    cls_tokens = cls_tokens.flatten(2).transpose(-2, -1)\n",
        "    cls_tokens = cls_tokens.reshape(-1, 1, cls_tokens.size(-1))\n",
        "\n",
        "    out = torch.cat((cls_tokens, patch_tokens), dim=1)\n",
        "\n",
        "    return out, mask, p_l, p_r, p_t, p_b, B, C, H, W\n",
        "\n",
        "\n",
        "def convert_to_spatial_layout(out, output_channels, B, H, W, kernel_size, mask, p_l, p_r, p_t, p_b):\n",
        "    cls_tokens = out[:, 0:1, ...]\n",
        "    patch_tokens = out[:, 1:, ...]\n",
        "    C = output_channels\n",
        "    kernel_size = kernel_size[0]\n",
        "    H_ks = H // kernel_size[0]\n",
        "    W_ks = W // kernel_size[1]\n",
        "    cls_tokens = cls_tokens.reshape(B, -1, C).transpose(-2, -1).reshape(B, C, H_ks, W_ks)\n",
        "    patch_tokens = patch_tokens.transpose(1, 2).reshape((B, -1, kernel_size[0] * kernel_size[1] * C)).transpose(1, 2)\n",
        "    patch_tokens = F.fold(patch_tokens, (H, W), kernel_size=kernel_size, stride=kernel_size, padding=(0, 0))\n",
        "\n",
        "    if mask is not None:\n",
        "        if p_b > 0:\n",
        "            patch_tokens = patch_tokens[:, :, :-p_b, :]\n",
        "        if p_r > 0:\n",
        "            patch_tokens = patch_tokens[:, :, :, :-p_r]\n",
        "\n",
        "    return cls_tokens, patch_tokens"
      ],
      "metadata": {
        "id": "4oKbAW6ZIF2Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class LayerNorm2d(nn.Module):\n",
        "    def __init__(self, channels, eps: float = 1e-5, elementwise_affine: bool = True) -> None:\n",
        "        super(LayerNorm2d, self).__init__()\n",
        "\n",
        "        self.channels = channels\n",
        "        self.eps = torch.tensor(eps)\n",
        "        self.elementwise_affine = elementwise_affine\n",
        "        if self.elementwise_affine:\n",
        "            self.weight = nn.Parameter(torch.zeros(1, channels, 1, 1))\n",
        "            self.bias = nn.Parameter(torch.zeros(1, channels, 1, 1))\n",
        "        else:\n",
        "            self.register_parameter('weight', None)\n",
        "            self.register_parameter('bias', None)\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self) -> None:\n",
        "        if self.elementwise_affine:\n",
        "            nn.init.ones_(self.weight)\n",
        "            nn.init.zeros_(self.bias)\n",
        "\n",
        "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
        "        mean = input.mean(1, keepdim=True)\n",
        "        std = torch.sqrt(input.var(1, unbiased=False, keepdim=True) + self.eps)\n",
        "        out = (input - mean) / std\n",
        "        if self.elementwise_affine:\n",
        "            out = out * self.weight + self.bias\n",
        "        return out\n",
        "\n",
        "    def extra_repr(self):\n",
        "        return '{channels}, eps={eps}, ' \\\n",
        "            'elementwise_affine={elementwise_affine}'.format(**self.__dict__)\n"
      ],
      "metadata": {
        "id": "0Ckk7mPEPEY8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class RegionViT(nn.Module):\n",
        "    \"\"\"\n",
        "    Note:\n",
        "        The variable naming mapping between codes and papers:\n",
        "        - cls_tokens -> regional tokens\n",
        "        - patch_tokens -> local tokens\n",
        "    \"\"\"\n",
        "    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=(768,), depth=(12,),\n",
        "                 num_heads=(12,), mlp_ratio=4., qkv_bias=True, qk_scale=None, drop_rate=0., attn_drop_rate=0.,\n",
        "                 drop_path_rate=0., norm_layer=partial(nn.LayerNorm, eps=1e-6),\n",
        "                 # regionvit parameters\n",
        "                 kernel_sizes=None, downsampling=None,\n",
        "                 patch_conv_type='3conv',\n",
        "                 computed_cls_token=True, peg=False,\n",
        "                 det_norm=False):\n",
        "\n",
        "        super().__init__()\n",
        "        self.num_classes = num_classes\n",
        "        self.kernel_sizes = kernel_sizes\n",
        "        self.num_features = embed_dim[-1]  # num_features for consistency with other models\n",
        "        self.embed_dim = embed_dim\n",
        "        self.patch_size = patch_size\n",
        "        self.img_size = img_size\n",
        "        self.patch_embed = PatchEmbed(\n",
        "            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim[0],\n",
        "            patch_conv_type=patch_conv_type)\n",
        "\n",
        "        if not isinstance(mlp_ratio, (list, tuple)):\n",
        "            mlp_ratio = [mlp_ratio] * len(depth)\n",
        "\n",
        "        self.computed_cls_token = computed_cls_token\n",
        "        self.cls_token = PatchEmbed(\n",
        "            img_size=img_size, patch_size=patch_size * kernel_sizes[0], in_chans=in_chans, embed_dim=embed_dim[0],\n",
        "            patch_conv_type='linear'\n",
        "        )\n",
        "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
        "        total_depth = sum(depth)\n",
        "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, total_depth)]  # stochastic depth decay rule\n",
        "        dpr_ptr = 0\n",
        "        self.layers = nn.ModuleList()\n",
        "        for i in range(len(embed_dim) - 1):\n",
        "            curr_depth = depth[i]\n",
        "            dpr_ = dpr[dpr_ptr:dpr_ptr + curr_depth]\n",
        "\n",
        "            self.layers.append(\n",
        "                ConvAttBlock(embed_dim[i], embed_dim[i + 1], kernel_size=kernel_sizes[i], num_blocks=depth[i], drop_path_rate=dpr_,\n",
        "                             num_heads=num_heads[i], mlp_ratio=mlp_ratio[i], qkv_bias=qkv_bias, qk_scale=qk_scale,\n",
        "                             pool=downsampling[i], norm_layer=norm_layer, attn_drop_rate=attn_drop_rate, drop_rate=drop_rate,\n",
        "                             cls_attn=True, peg=peg)\n",
        "            )\n",
        "            dpr_ptr += curr_depth\n",
        "        self.norm = norm_layer(embed_dim[-1])\n",
        "\n",
        "        # Classifier head\n",
        "        self.head = nn.Linear(embed_dim[-1], num_classes) if num_classes > 0 else nn.Identity()\n",
        "\n",
        "        if not computed_cls_token:\n",
        "            trunc_normal_(self.cls_token, std=.02)\n",
        "\n",
        "        self.det_norm = det_norm\n",
        "        if self.det_norm:\n",
        "            # add a norm layer for the outputs at each stage, for detection\n",
        "            for i in range(4):\n",
        "                layer = LayerNorm2d(embed_dim[1 + i])\n",
        "                layer_name = f'norm{i}'\n",
        "                self.add_module(layer_name, layer)\n",
        "\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, m):\n",
        "        if isinstance(m, nn.Linear):\n",
        "            trunc_normal_(m.weight, std=.02)\n",
        "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "        elif isinstance(m, nn.LayerNorm):\n",
        "            nn.init.constant_(m.bias, 0)\n",
        "            nn.init.constant_(m.weight, 1.0)\n",
        "\n",
        "    @torch.jit.ignore\n",
        "    def no_weight_decay(self):\n",
        "        if not self.computed_cls_token:\n",
        "            return {'cls_token'}\n",
        "        else:\n",
        "            return {}\n",
        "\n",
        "    def get_classifier(self):\n",
        "        return self.head\n",
        "\n",
        "    def reset_classifier(self, num_classes, global_pool=''):\n",
        "        self.num_classes = num_classes\n",
        "        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()\n",
        "\n",
        "    def forward_features(self, x, detection=False):\n",
        "        o_x = x\n",
        "        x = self.patch_embed(x)\n",
        "        # B x branches x classes\n",
        "        cls_tokens = self.cls_token(o_x, extra_padding=True)\n",
        "        x = self.pos_drop(x)  # N C H W\n",
        "        tmp_out = []\n",
        "\n",
        "        for idx, layer in enumerate(self.layers):\n",
        "            cls_tokens, x = layer((cls_tokens, x))\n",
        "            if self.det_norm:\n",
        "                norm_layer = getattr(self, f'norm{idx}')\n",
        "                x = norm_layer(x)\n",
        "            tmp_out.append(x)\n",
        "\n",
        "        if detection:\n",
        "            return tmp_out\n",
        "\n",
        "        N, C, H, W = cls_tokens.shape\n",
        "        cls_tokens = cls_tokens.reshape(N, C, -1).transpose(1, 2)\n",
        "        cls_tokens = self.norm(cls_tokens)\n",
        "        out = torch.mean(cls_tokens, dim=1)\n",
        "\n",
        "        return out\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.forward_features(x)\n",
        "        x = self.head(x)\n",
        "\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "kULSvjPyQR-B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from einops import rearrange\n",
        "\n",
        "class RegionalLocalAttention(nn.Module):\n",
        "    def __init__(self, dim_regional, dim_local, num_heads):\n",
        "        super().__init__()\n",
        "        self.num_heads = num_heads\n",
        "        self.scale = dim_regional ** -0.5\n",
        "\n",
        "        # Regional self-attention\n",
        "        self.qkv_regional = nn.Linear(dim_regional, dim_regional * 3, bias=True)\n",
        "        self.proj_regional = nn.Linear(dim_regional, dim_regional)\n",
        "\n",
        "        # Local self-attention\n",
        "        self.qkv_local = nn.Linear(dim_local, dim_local * 3, bias=True)\n",
        "        self.proj_local = nn.Linear(dim_local, dim_local)\n",
        "\n",
        "        # Fusion\n",
        "        self.fusion = nn.Linear(dim_regional + dim_local, dim_regional)\n",
        "\n",
        "    def forward(self, regional_tokens, local_tokens):\n",
        "        B, N_regional, C_regional = regional_tokens.shape\n",
        "        _, N_local, C_local = local_tokens.shape\n",
        "\n",
        "        # Regional self-attention\n",
        "        qkv_regional = self.qkv_regional(regional_tokens).reshape(B, N_regional, 3, self.num_heads, C_regional // self.num_heads).permute(2, 0, 3, 1, 4)\n",
        "        q_regional, k_regional, v_regional = qkv_regional[0], qkv_regional[1], qkv_regional[2]\n",
        "        attn_regional = (q_regional @ k_regional.transpose(-2, -1)) * self.scale\n",
        "        attn_regional = attn_regional.softmax(dim=-1)\n",
        "        regional_attn_out = (attn_regional @ v_regional).transpose(1, 2).reshape(B, N_regional, C_regional)\n",
        "        regional_attn_out = self.proj_regional(regional_attn_out)\n",
        "\n",
        "        # Local self-attention (within each regional token's receptive field)\n",
        "        # (Implementation details for local attention may vary based on receptive field definition)\n",
        "        # For simplicity, assume each regional token attends to all local tokens\n",
        "        qkv_local = self.qkv_local(local_tokens).reshape(B, N_local, 3, self.num_heads, C_local // self.num_heads).permute(2, 0, 3, 1, 4)\n",
        "        q_local, k_local, v_local = qkv_local[0], qkv_local[1], qkv_local[2]\n",
        "        attn_local = (q_local @ k_local.transpose(-2, -1)) * self.scale\n",
        "        attn_local = attn_local.softmax(dim=-1)\n",
        "        local_attn_out = (attn_local @ v_local).transpose(1, 2).reshape(B, N_local, C_local)\n",
        "        local_attn_out = self.proj_local(local_attn_out)\n",
        "\n",
        "        # Fusion\n",
        "        fused_features = self.fusion(torch.cat([regional_attn_out, local_attn_out], dim=-1))\n",
        "\n",
        "        return fused_features\n",
        "\n",
        "# ... (In your VisionTransformer)\n",
        "\n",
        "# Assuming patch_size_3 is your third patch size\n",
        "patch_embed_3 = PatchEmbed(img_size=img_size, patch_size=patch_size_3, in_chan=in_chans, embed_dim=embed_dim)\n",
        "patch_features_3 = patch_embed_3(x)\n",
        "\n",
        "# Regional-local attention\n",
        "regional_local_attn = RegionalLocalAttention(dim_regional=embed_dim, dim_local=embed_dim, num_heads=num_heads)\n",
        "fused_features = regional_local_attn(crossvit_output, patch_features_3)\n",
        "\n",
        "# ... (Further processing using fused_features)"
      ],
      "metadata": {
        "id": "hVc1InrsI_9J"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}