{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":3004,"databundleVersionId":861823,"sourceType":"competition"}],"dockerImageVersionId":30840,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom einops import rearrange, repeat","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:55.906200Z","iopub.execute_input":"2025-01-24T10:18:55.906502Z","iopub.status.idle":"2025-01-24T10:18:55.910163Z","shell.execute_reply.started":"2025-01-24T10:18:55.906478Z","shell.execute_reply":"2025-01-24T10:18:55.909438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PatchEmbedding(nn.Module):\n    \"\"\"\n    Layer to take in the input image and do the following:\n        1.  Transform grid of image into sequence of patches.\n            No. of patches = (Img_height*Img_width)/(patch_height*patch_width)\n        2.  Add CLS token in the first position\n        3.  Add positional embeddings to the sequence.\n        4.  Dropout if needed\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        # Example config\n        #   Image c, h, w = 3, 224, 224\n        #   Patch    h, w =    16,  16\n\n        # This dictionary contains several key-value pairs, where each key represents a hyperparameter and an optimal value is assigned.\n\n        image_height = config['image_height']\n        image_width = config['image_width']\n        im_channels = config['im_channels']\n        emb_dim = config['emb_dim']\n        patch_emb_drop = config['patch_emb_drop']\n\n        self.patch_height = config['patch_height']\n        self.patch_width = config['patch_width']\n\n        n_patches = (image_height * image_width) // (self.patch_height * self.patch_width)\n        patch_dim = im_channels * self.patch_height * self.patch_width\n\n        self.patch_emb = nn.Sequential(\n            nn.Linear(patch_dim, emb_dim)\n        )\n        \"\"\"\n        Instead of separating patches we could have used a conv layer instead of fc layer \n        with kernel size as (patch_height, patch_width) and same stride.\n        \"\"\"\n\n        # Positional Embeddings\n        self.pos_emb = nn.Parameter(torch.zeros(1, n_patches + 1, emb_dim))\n        self.cls_token = nn.Parameter(torch.randn(emb_dim))\n        self.patch_emb_dropout = nn.Dropout(patch_emb_drop) \n\n\n    def forward(self, x):\n        batch_size = x.shape[0]\n        out = rearrange(x, 'b c (nh ph) (nw pw) -> b (nh nw) (ph pw c)', ph=self.patch_height, pw=self.patch_width)\n         # Flatten spatial dimensions\n        out = self.patch_emb(out)\n\n        # Add cls\n        cls_token = repeat(self.cls_token, 'd -> b 1 d', b=batch_size)\n        out = torch.cat((cls_token, out), dim=1)\n\n        # Add positional embedd and do dropout\n        out += self.pos_emb\n        out = self.patch_emb_dropout(out)\n\n        return out\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:55.917081Z","iopub.execute_input":"2025-01-24T10:18:55.917387Z","iopub.status.idle":"2025-01-24T10:18:55.930325Z","shell.execute_reply.started":"2025-01-24T10:18:55.917348Z","shell.execute_reply":"2025-01-24T10:18:55.929480Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Attention(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.n_heads = config['n_heads']\n        self.head_dim = config['head_dim']\n        self.emb_dim = config['emb_dim']\n        self.drop_prob = config['dropout'] if 'dropout' in config else 0.0\n        self.att_dim = self.n_heads * self.head_dim\n\n        # qkv_proj is the NN whose weight matrix is W matrix multiplied with patch rep. matrix to get q,k,v rep\n        # Thus its a NN to give q,k,v rep on splitting\n        self.qkv_proj = nn.Linear(self.emb_dim, 3 * self.att_dim, bias=False)\n        # FC NN for getting the final output embeddings of each patch (n_patches * emb_dim)\n        self.output_proj = nn.Sequential(\n            nn.Linear(self.att_dim, self.emb_dim),\n            nn.Dropout(self.drop_prob)\n        )\n        self.att_dropout = nn.Dropout(self.drop_prob)\n\n\n    def forward(self, x):\n        # Converting to Attention Dimension\n        B, N = x.shape[:2]\n        # Projecting to 3*att_dim and then splitting to get q, k v(each of att_dim)\n        # qkv -> Batch Size x Number of Patches x (3* Attention Dimension)\n        # q(as well as k and v) -> Batch Size x Number of Patches x Attention Dimension\n        q, k, v = self.qkv_proj(x).split(self.att_dim, dim=-1)\n        # Batch Size x Number of Patches x Attention Dimension\n        # -> Batch Size x Number of Patches x (Heads * Head Dimension)\n        # -> Batch Size x Number of Patches x (Heads * Head Dimension)\n        # -> Batch Size x Heads x Number of Patches x Head Dimension\n        q = rearrange(q, 'b n (h h_dim) -> b h n h_dim', h=self.n_heads, h_dim=self.head_dim)\n        k = rearrange(k, 'b n (h h_dim) -> b h n h_dim', h=self.n_heads, h_dim=self.head_dim)\n        v = rearrange(v, 'b n (h h_dim) -> b h n h_dim', h=self.n_heads, h_dim=self.head_dim)\n\n        # Computing Attention Weights\n        # B x H x N x Head Dimension @ B x H x Head Dimension x N\n        # -> B x H x N x N\n        att = torch.matmul(q, k.transpose(-2, -1) * (self.head_dim**(-0.5)))\n        att = torch.nn.functional.softmax(att, dim=-1)\n        att = self.att_dropout(att)\n\n        # Weighted Value Computation\n        #  B x H x N x N @ B x H x N x Head Dimension\n        # -> B x H x N x Head Dimension\n        out = torch.matmul(att, v)\n\n        # Converting to Transformer dimension\n        # B x N x (Heads * Head Dimension) -> B x N x (Attention Dimension)\n        out = rearrange(out, 'b h n h_dim -> b n (h h_dim)')\n        #  B x N x Dimension\n        out = self.output_proj(out)\n\n        return out\n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:55.931344Z","iopub.execute_input":"2025-01-24T10:18:55.931603Z","iopub.status.idle":"2025-01-24T10:18:55.953367Z","shell.execute_reply.started":"2025-01-24T10:18:55.931582Z","shell.execute_reply":"2025-01-24T10:18:55.952504Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TransformerLayer(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        emb_dim = config['emb_dim']\n        ff_hidden_dim = config['ff_dim'] if 'ff_dim' in config else 4*emb_dim\n        ff_drop_prob = config['ff_drop'] if 'ff_drop' in config else 0.0\n        self.att_norm = nn.LayerNorm(emb_dim)\n        self.att_block = Attention(config)\n        self.ff_norm = nn.LayerNorm(emb_dim)\n        self.ff_block = nn.Sequential(\n            nn.Linear(emb_dim, ff_hidden_dim),\n            nn.GELU(),\n            nn.Dropout(ff_drop_prob),\n            nn.Linear(ff_hidden_dim, emb_dim),\n            nn.Dropout(ff_drop_prob)\n        )\n\n\n    def forward(self, x):\n        out = x\n        out = out + self.att_block(self.att_norm(out))\n        out = out + self.ff_block(self.ff_norm(out))\n        return out","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:55.954493Z","iopub.execute_input":"2025-01-24T10:18:55.954749Z","iopub.status.idle":"2025-01-24T10:18:55.974437Z","shell.execute_reply.started":"2025-01-24T10:18:55.954729Z","shell.execute_reply":"2025-01-24T10:18:55.973663Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class VIT(nn.Module):\n    def __init__(self, config):\n        super(VIT, self).__init__()\n        n_layers = config['n_layers']\n        emb_dim = config['emb_dim']\n        num_classes = config['num_classes']\n        self.patch_emb = PatchEmbedding(config)\n        self.layers = nn.ModuleList([\n            TransformerLayer(config) for _ in range(n_layers)\n        ])\n        self.norm = nn.LayerNorm(emb_dim)\n        self.fc_number = nn.Linear(emb_dim, num_classes)\n\n    def forward(self, x):\n        # Patchify and add CLS token \n        out = self.patch_emb(x)\n\n        # Go through the layers in Transformer\n        for layer in self.layers:\n            out = layer(out)\n        out = self.norm(out)\n\n        return self.fc_number(out[:, 0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:56.078738Z","iopub.execute_input":"2025-01-24T10:18:56.078972Z","iopub.status.idle":"2025-01-24T10:18:56.083821Z","shell.execute_reply.started":"2025-01-24T10:18:56.078953Z","shell.execute_reply":"2025-01-24T10:18:56.082973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import yaml\nimport argparse\nimport torch\nimport cv2\nimport os\nimport numpy as np\nimport pandas as pd\nfrom matplotlib import pyplot as plt\nfrom tqdm import tqdm\nfrom torch.utils.data.dataloader import DataLoader\nimport torch\nfrom torchvision import datasets, transforms, models\nimport torch.optim as optim","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:56.084628Z","iopub.execute_input":"2025-01-24T10:18:56.084913Z","iopub.status.idle":"2025-01-24T10:18:56.104418Z","shell.execute_reply.started":"2025-01-24T10:18:56.084884Z","shell.execute_reply":"2025-01-24T10:18:56.103615Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config = {\n# dataset_params\n    \"root_dir\": \"data\",\n\n# model_params\n    \"n_heads\": 8,\n    \"head_dim\": 64,\n    \"emb_dim\": 256,\n    \"attn_drop\": 0.1,\n    \"ff_dim\": 256,\n    \"ff_drop\": 0.1,\n    \"n_layers\": 6,\n    \"num_classes\": 10,\n    \"image_height\": 224,\n    \"image_width\": 224,\n    \"patch_height\": 16,\n    \"patch_width\": 16,\n    \"patch_emb_drop\": 0.1,\n    \"im_channels\": 1,\n\n# train_params\n    \"task_name\": \"default\",\n    \"batch_size\": 64,\n    \"epochs\": 100,\n    \"lr\": 0.001,\n    \"seed\": 1111,\n    \"ckpt_name\": \"vit_ckpt.pth\"\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:56.106216Z","iopub.execute_input":"2025-01-24T10:18:56.106468Z","iopub.status.idle":"2025-01-24T10:18:56.119776Z","shell.execute_reply.started":"2025-01-24T10:18:56.106436Z","shell.execute_reply":"2025-01-24T10:18:56.118996Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Device configuration\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n# Data preprocessing\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),  # Resize to match ViT input size\n    transforms.ToTensor(),\n    transforms.Normalize((0.5,), (0.5,))\n])\n\n# Load MNIST dataset\ntrain_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)\ntest_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)\n\ntrain_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)\ntest_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)\n\nmodel = VIT(config)\nmodel = model.to(device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:56.120794Z","iopub.execute_input":"2025-01-24T10:18:56.121066Z","iopub.status.idle":"2025-01-24T10:18:56.259396Z","shell.execute_reply.started":"2025-01-24T10:18:56.121038Z","shell.execute_reply":"2025-01-24T10:18:56.258715Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training the model\n\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.Adam(model.parameters(), lr=0.0005)\n\ndef train(model, loader, criterion, optimizer, device):\n    model.train()\n    total_loss = 0\n    correct = 0\n    for images, labels in loader:\n        images, labels = images.to(device), labels.to(device)\n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n        _, predicted = outputs.max(1)\n        correct += predicted.eq(labels).sum().item()\n    return total_loss / len(loader), correct / len(loader.dataset)\n\ndef test(model, loader, device):\n    model.eval()\n    correct_pred = 0\n    total_samples = 0\n    with torch.no_grad():\n        for images, labels in loader:\n            images, labels = images.to(device), labels.to(device)\n            outputs = model(images)\n            predictions = torch.argmax(outputs, dim=1)  # Get predicted class indices\n            correct_pred += (predictions == labels).sum().item()\n            total_samples += labels.size(0)\n        return correct_pred / total_samples\n\nimport matplotlib.pyplot as plt\nimport torch\n\ndef show_test(model, loader, device, rows=8, cols=8):\n    model.eval()\n    with torch.no_grad():\n        for images, labels in loader:\n            images, labels = images.to(device), labels.to(device)\n            y_hat = model(images)\n            predictions = torch.argmax(y_hat, dim=1)\n            \n            # Limit the number of images to display in the grid\n            total_images = min(rows * cols, len(images))\n            \n            fig, axes = plt.subplots(rows, cols, figsize=(cols * 3, rows * 3))\n            axes = axes.flatten()  # Flatten to easily iterate over the axes\n\n            for i in range(total_images):\n                image = images[i].cpu().squeeze()  # Move to CPU and remove channel dimension if grayscale\n                label = labels[i].cpu().item()\n                prediction = predictions[i].cpu().item()\n\n                ax = axes[i]\n                ax.imshow(image, cmap=\"gray\" if image.ndim == 2 else None)\n                ax.set_title(f\"True: {label}, Pred: {prediction}\", fontsize=10)\n                ax.axis(\"off\")\n\n            # Hide any remaining empty subplots\n            for j in range(total_images, len(axes)):\n                axes[j].axis(\"off\")\n\n            plt.tight_layout()\n            plt.show()\n\n            # Display only the first batch (optional)\n            break\n\n        \n        ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:58:16.196321Z","iopub.execute_input":"2025-01-24T10:58:16.196676Z","iopub.status.idle":"2025-01-24T10:58:16.207732Z","shell.execute_reply.started":"2025-01-24T10:58:16.196626Z","shell.execute_reply":"2025-01-24T10:58:16.206864Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"num_epochs = 10\nfor epoch in range(num_epochs):\n    train_loss, train_acc = train(model, train_loader, criterion, optimizer, device)\n    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:18:56.269889Z","iopub.execute_input":"2025-01-24T10:18:56.270134Z","iopub.status.idle":"2025-01-24T10:54:56.881962Z","shell.execute_reply.started":"2025-01-24T10:18:56.270114Z","shell.execute_reply":"2025-01-24T10:54:56.881001Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_acc = test(model, test_loader, device)\nprint(\"Test Accuracy: \", test_acc)\n\nshow_test(model, test_loader, device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T10:58:21.797528Z","iopub.execute_input":"2025-01-24T10:58:21.797871Z","iopub.status.idle":"2025-01-24T10:58:41.015586Z","shell.execute_reply.started":"2025-01-24T10:58:21.797840Z","shell.execute_reply":"2025-01-24T10:58:41.014534Z"}},"outputs":[],"execution_count":null}]}