{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":3362,"databundleVersionId":31148,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **Extract File**","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport torch\nfrom torchvision import models\nimport zipfile","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:26:34.927675Z","iopub.execute_input":"2025-08-21T02:26:34.928166Z","iopub.status.idle":"2025-08-21T02:26:34.931860Z","shell.execute_reply.started":"2025-08-21T02:26:34.928144Z","shell.execute_reply":"2025-08-21T02:26:34.931176Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train = \"/kaggle/input/dogs-vs-cats/train.zip\"\ntest  = \"/kaggle/input/dogs-vs-cats/test1.zip\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:26:36.299166Z","iopub.execute_input":"2025-08-21T02:26:36.299455Z","iopub.status.idle":"2025-08-21T02:26:36.303056Z","shell.execute_reply.started":"2025-08-21T02:26:36.299432Z","shell.execute_reply":"2025-08-21T02:26:36.302472Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with zipfile.ZipFile(train, 'r') as zip_ref:\n    zip_ref.extractall(\"/kaggle/working/train\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:26:38.202835Z","iopub.execute_input":"2025-08-21T02:26:38.203486Z","iopub.status.idle":"2025-08-21T02:26:47.541848Z","shell.execute_reply.started":"2025-08-21T02:26:38.203464Z","shell.execute_reply":"2025-08-21T02:26:47.541018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"with zipfile.ZipFile(test, 'r') as zip_ref:\n    zip_ref.extractall(\"/kaggle/working/test\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:26:49.118924Z","iopub.execute_input":"2025-08-21T02:26:49.119601Z","iopub.status.idle":"2025-08-21T02:26:53.404445Z","shell.execute_reply.started":"2025-08-21T02:26:49.119566Z","shell.execute_reply":"2025-08-21T02:26:53.403846Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_data = \"/kaggle/working/train/train\" #it contain images with name like cat.1.jpg, cat.2.jpg, dog.1.jpg, etc\ntest_data = \"/kaggle/working/test/test1\" #the images naming is only 1.jpg, 2.jpg, etc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:26:54.899799Z","iopub.execute_input":"2025-08-21T02:26:54.900667Z","iopub.status.idle":"2025-08-21T02:26:54.904073Z","shell.execute_reply.started":"2025-08-21T02:26:54.900636Z","shell.execute_reply":"2025-08-21T02:26:54.903350Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Here, we will map the attention map **from [CLS] token**.\n\n[CLS] token is some kind a decider, a global representative for the entire sequence. You can say it is somekind of a captain to other image patches.","metadata":{}},{"cell_type":"markdown","source":"# **Simple Attention Map Visualization**","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T03:31:48.515121Z","iopub.execute_input":"2025-08-21T03:31:48.515817Z","iopub.status.idle":"2025-08-21T03:31:48.519826Z","shell.execute_reply.started":"2025-08-21T03:31:48.515786Z","shell.execute_reply":"2025-08-21T03:31:48.519069Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Imagine 3 pixel with 2 channel:\n\nRow 1 = Pixel A's 2 channel features\n\nRow 2 = Pixel B's 2 channel features\n\nRow 3 = Pixel C's 2 channel features","metadata":{}},{"cell_type":"code","source":"# Toy feature map: 3 pixels, 2 channels\nX_pixels = torch.tensor([[1., 0.],   # Pixel A\n                         [0., 1.],   # Pixel B\n                         [1., 1.]])  # Pixel C\n\n# Weight matrices (kept the same for comparison)\nW_Q = torch.tensor([[1., 0.],\n                    [0., 1.]])\nW_K = torch.tensor([[1., 1.],\n                    [0., 1.]])\nW_V = torch.tensor([[1., 0.],\n                    [0., 1.]])\n\n# --- 1. Add the [CLS] Token ---\n# In a real model, this would be a learnable nn.Parameter.\n# Here, we'll just define it as a tensor for demonstration.\ncls_token = torch.tensor([[-0.5, 0.5]]) # A distinct starting feature\n\n# Prepend the CLS token to the sequence of pixels\nX_with_cls = torch.cat([cls_token, X_pixels], dim=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T03:31:50.584163Z","iopub.execute_input":"2025-08-21T03:31:50.584427Z","iopub.status.idle":"2025-08-21T03:31:50.590012Z","shell.execute_reply.started":"2025-08-21T03:31:50.584406Z","shell.execute_reply":"2025-08-21T03:31:50.589366Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"--- Input with [CLS] Token ---\")\nprint(X_with_cls)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T03:31:53.039575Z","iopub.execute_input":"2025-08-21T03:31:53.039844Z","iopub.status.idle":"2025-08-21T03:31:53.088762Z","shell.execute_reply.started":"2025-08-21T03:31:53.039825Z","shell.execute_reply":"2025-08-21T03:31:53.088122Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 2. Run the Attention Mechanism ---\n# The process is the same, but now on a 4x2 input tensor\nQ = X_with_cls @ W_Q\nK = X_with_cls @ W_K\nV = X_with_cls @ W_V\n\nscores = Q @ K.T\nattn = F.softmax(scores, dim=-1)\n\nout = attn @ V\n\nprint(\"\\n--- New 4x4 Attention Weights ---\")\nprint(attn)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T03:32:01.968276Z","iopub.execute_input":"2025-08-21T03:32:01.968544Z","iopub.status.idle":"2025-08-21T03:32:02.027861Z","shell.execute_reply.started":"2025-08-21T03:32:01.968528Z","shell.execute_reply":"2025-08-21T03:32:02.026987Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 3. Visualize the New Attention Map ---\nattn_np = attn.detach().numpy()\n\nplt.figure(figsize=(7, 6))\nsns.heatmap(attn_np, annot=True, fmt=\".4f\", cmap=\"viridis\",\n            xticklabels=['[CLS]', 'Pixel A', 'Pixel B', 'Pixel C'],\n            yticklabels=['[CLS]', 'Pixel A', 'Pixel B', 'Pixel C'])\n\nplt.title(\"[CLS] Token Self-Attention Map\")\nplt.xlabel(\"Tokens Being Looked At (Keys)\")\nplt.ylabel(\"Tokens Doing the Looking (Queries)\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T03:32:07.423328Z","iopub.execute_input":"2025-08-21T03:32:07.423877Z","iopub.status.idle":"2025-08-21T03:32:07.633158Z","shell.execute_reply.started":"2025-08-21T03:32:07.423855Z","shell.execute_reply":"2025-08-21T03:32:07.632593Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Code above is just a very basic [CLS] token implementation. In a real, trained model, the [CLS] token would learn to pay high attention to the pixels (or image patches) that are most critical for the classification task. The other rows show what the individual pixels are paying attention to, but for the final decision, it's the [CLS] token's perspective that matters.","metadata":{}},{"cell_type":"markdown","source":"So in the cat vs dog below, we only see from the [CLS] query corresponds with every image patch. The brightest area means that [CLS] token is pay more attention on that specific image patch. We will try Vision Transformer from scratch and with pretrained.","metadata":{}},{"cell_type":"markdown","source":"# **Vision Transformer From Scratch**","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision import transforms\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport os\nfrom PIL import Image\nimport pandas as pd\nfrom tqdm import tqdm\nimport numpy as np\nimport math","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:10.849785Z","iopub.execute_input":"2025-08-21T00:58:10.850518Z","iopub.status.idle":"2025-08-21T00:58:12.003084Z","shell.execute_reply.started":"2025-08-21T00:58:10.850490Z","shell.execute_reply":"2025-08-21T00:58:12.002550Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 1. Configuration\n# ======================================================================================\nTRAIN_DIR = \"/kaggle/working/train/train\"\nTEST_DIR = \"/kaggle/working/test/test1\"\nSUBMISSION_FILE = \"submission.csv\"\n\nIMG_SIZE = 224\nPATCH_SIZE = 16\nBATCH_SIZE = 32\nLEARNING_RATE = 1e-4\nEPOCHS = 10 \nEMBED_DIM = 256\nNUM_HEADS = 8\nNUM_LAYERS = 6\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:14.424970Z","iopub.execute_input":"2025-08-21T00:58:14.425409Z","iopub.status.idle":"2025-08-21T00:58:14.516083Z","shell.execute_reply.started":"2025-08-21T00:58:14.425385Z","shell.execute_reply":"2025-08-21T00:58:14.515530Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 2. Custom Dataset for Cats vs. Dogs (Unchanged)\n# ======================================================================================\nclass CatsDogsDataset(Dataset):\n    def __init__(self, data_dir, transform=None, is_test=False):\n        self.data_dir = data_dir\n        self.transform = transform\n        self.is_test = is_test\n        self.image_files = sorted(os.listdir(self.data_dir), key=lambda x: int(x.split('.')[0]) if self.is_test else (x.split('.')[0], int(x.split('.')[1])))\n\n    def __len__(self):\n        return len(self.image_files)\n\n    def __getitem__(self, idx):\n        img_name = self.image_files[idx]\n        img_path = os.path.join(self.data_dir, img_name)\n        image = Image.open(img_path).convert(\"RGB\")\n        if self.transform:\n            image = self.transform(image)\n        if not self.is_test:\n            label = 0 if 'cat' in img_name else 1\n            return image, torch.tensor(label, dtype=torch.long)\n        else:\n            image_id = int(img_name.split('.')[0])\n            return image, image_id","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:16.729606Z","iopub.execute_input":"2025-08-21T00:58:16.730286Z","iopub.status.idle":"2025-08-21T00:58:16.736406Z","shell.execute_reply.started":"2025-08-21T00:58:16.730258Z","shell.execute_reply":"2025-08-21T00:58:16.735692Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 3. Image Transformations (Unchanged)\n# ======================================================================================\ndata_transforms = {\n    'train': transforms.Compose([\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomRotation(10),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n    'val': transforms.Compose([\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:19.473539Z","iopub.execute_input":"2025-08-21T00:58:19.474086Z","iopub.status.idle":"2025-08-21T00:58:19.478655Z","shell.execute_reply.started":"2025-08-21T00:58:19.474063Z","shell.execute_reply":"2025-08-21T00:58:19.477995Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 4. Create DataLoaders (Unchanged)\n# ======================================================================================\nfull_train_dataset = CatsDogsDataset(TRAIN_DIR, transform=data_transforms['train'])\ntrain_size = int(0.8 * len(full_train_dataset))\nval_size = len(full_train_dataset) - train_size\ntrain_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])\nval_dataset.dataset.transform = data_transforms['val'] \ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\nprint(f\"Training samples: {len(train_dataset)}\")\nprint(f\"Validation samples: {len(val_dataset)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:21.896210Z","iopub.execute_input":"2025-08-21T00:58:21.896899Z","iopub.status.idle":"2025-08-21T00:58:21.950960Z","shell.execute_reply.started":"2025-08-21T00:58:21.896874Z","shell.execute_reply":"2025-08-21T00:58:21.950406Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 5. Vision Transformer from Scratch (Returns Attention)\n# ======================================================================================\nclass PatchEmbedding(nn.Module):\n    def __init__(self, img_size, patch_size, in_channels, embed_dim):\n        super().__init__()\n        self.n_patches = (img_size // patch_size) ** 2\n        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)\n    def forward(self, x): return self.proj(x).flatten(2).transpose(1, 2)\n\nclass MultiHeadSelfAttention(nn.Module):\n    def __init__(self, embed_dim, num_heads):\n        super().__init__()\n        self.num_heads, self.head_dim = num_heads, embed_dim // num_heads\n        self.qkv, self.out_proj = nn.Linear(embed_dim, embed_dim * 3), nn.Linear(embed_dim, embed_dim)\n    def forward(self, x):\n        B, N, C = x.shape\n        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n        q, k, v = qkv[0], qkv[1], qkv[2]\n        attn_scores = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)\n        attn = attn_scores.softmax(dim=-1)\n        context = (attn @ v).transpose(1, 2).reshape(B, N, C)\n        return self.out_proj(context), attn\n\nclass TransformerEncoderBlock(nn.Module):\n    def __init__(self, embed_dim, num_heads):\n        super().__init__()\n        self.norm1, self.norm2 = nn.LayerNorm(embed_dim), nn.LayerNorm(embed_dim)\n        self.attn = MultiHeadSelfAttention(embed_dim, num_heads)\n        self.mlp = nn.Sequential(nn.Linear(embed_dim, embed_dim * 4), nn.GELU(), nn.Linear(embed_dim * 4, embed_dim))\n    def forward(self, x):\n        attn_output, attn_weights = self.attn(self.norm1(x))\n        x = x + attn_output\n        x = x + self.mlp(self.norm2(x))\n        return x, attn_weights\n\nclass VisionTransformer(nn.Module):\n    def __init__(self, img_size, patch_size, in_channels, num_classes, embed_dim, num_heads, num_layers):\n        super().__init__()\n        num_patches = (img_size // patch_size) ** 2\n        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)\n        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))\n        self.encoder_layers = nn.ModuleList([TransformerEncoderBlock(embed_dim, num_heads) for _ in range(num_layers)])\n        self.norm = nn.LayerNorm(embed_dim)\n        self.head = nn.Linear(embed_dim, num_classes)\n    def forward(self, x):\n        B = x.shape[0]\n        x = self.patch_embed(x)\n        cls_tokens = self.cls_token.expand(B, -1, -1)\n        x = torch.cat((cls_tokens, x), dim=1)\n        x = x + self.pos_embed\n        for layer in self.encoder_layers: x, attn_weights = layer(x)\n        cls_output = self.norm(x[:, 0])\n        return self.head(cls_output), attn_weights","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:23.913846Z","iopub.execute_input":"2025-08-21T00:58:23.914734Z","iopub.status.idle":"2025-08-21T00:58:23.927965Z","shell.execute_reply.started":"2025-08-21T00:58:23.914706Z","shell.execute_reply":"2025-08-21T00:58:23.927181Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 6. Instantiate the Model\n# ======================================================================================\nmodel = VisionTransformer(\n    img_size=IMG_SIZE, patch_size=PATCH_SIZE, in_channels=3, num_classes=2,\n    embed_dim=EMBED_DIM, num_heads=NUM_HEADS, num_layers=NUM_LAYERS\n).to(DEVICE)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:26.802392Z","iopub.execute_input":"2025-08-21T00:58:26.802940Z","iopub.status.idle":"2025-08-21T00:58:27.067513Z","shell.execute_reply.started":"2025-08-21T00:58:26.802916Z","shell.execute_reply":"2025-08-21T00:58:27.066908Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 7. NEW: Epoch-by-Epoch Visualization Function\n# ======================================================================================\ndef visualize_attention_by_epoch(model, dataloader, device, epoch):\n    print(f\"\\n--- Visualizing Attention for Epoch {epoch} ---\")\n    model.eval()\n    \n    correct_samples, incorrect_samples = [], []\n    \n    # Find 5 correct and 5 incorrect samples\n    with torch.no_grad():\n        for images, labels in dataloader:\n            images, labels = images.to(device), labels.to(device)\n            logits, attn_weights = model(images)\n            preds = torch.argmax(logits, dim=1)\n            \n            for i in range(images.size(0)):\n                is_correct = (preds[i] == labels[i]).item()\n                sample = {\n                    'image': images[i], 'label': labels[i], \n                    'pred': preds[i], 'attn': attn_weights[i]\n                }\n                if is_correct and len(correct_samples) < 5:\n                    correct_samples.append(sample)\n                elif not is_correct and len(incorrect_samples) < 5:\n                    incorrect_samples.append(sample)\n            \n            if len(correct_samples) >= 5 and len(incorrect_samples) >= 5:\n                break\n\n    # If not enough incorrect samples, fill with correct ones\n    samples_to_show = incorrect_samples + correct_samples\n    samples_to_show = samples_to_show[:10]\n    \n    if not samples_to_show:\n        print(\"Not enough samples to visualize.\")\n        return\n\n    inv_normalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225])\n    \n    fig, axes = plt.subplots(2, len(samples_to_show), figsize=(20, 5))\n    \n    for i, sample in enumerate(samples_to_show):\n        img = inv_normalize(sample['image']).cpu().permute(1, 2, 0).numpy()\n        \n        # Process attention map\n        cls_attn = sample['attn'][:, 0, 1:].mean(dim=0) # Avg heads, CLS token -> patches\n        num_patches_side = int(math.sqrt(cls_attn.shape[0]))\n        attn_map = cls_attn.cpu().reshape(num_patches_side, num_patches_side).numpy()\n        \n        # Plot original image\n        axes[0, i].imshow(img)\n        axes[0, i].axis('off')\n        true_label = \"Dog\" if sample['label'].item() == 1 else \"Cat\"\n        pred_label = \"Dog\" if sample['pred'].item() == 1 else \"Cat\"\n        color = \"green\" if true_label == pred_label else \"red\"\n        axes[0, i].set_title(f\"True: {true_label}\\nPred: {pred_label}\", color=color)\n        \n        # Plot attention map\n        axes[1, i].imshow(attn_map, cmap='viridis')\n        axes[1, i].axis('off')\n\n    plt.suptitle(f\"Attention Maps - Epoch {epoch}\", fontsize=16)\n    plt.tight_layout(rect=[0, 0, 1, 0.96])\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:29.194424Z","iopub.execute_input":"2025-08-21T00:58:29.194696Z","iopub.status.idle":"2025-08-21T00:58:29.205514Z","shell.execute_reply.started":"2025-08-21T00:58:29.194680Z","shell.execute_reply":"2025-08-21T00:58:29.204660Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 8. Training Loop (MODIFIED to call visualization)\n# ======================================================================================\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)\n\nfor epoch in range(EPOCHS):\n    print(f\"\\n--- Epoch {epoch+1}/{EPOCHS} ---\")\n    model.train()\n    for inputs, labels in tqdm(train_loader, desc=\"Training\"):\n        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n        optimizer.zero_grad()\n        logits, _ = model(inputs)\n        loss = criterion(logits, labels)\n        loss.backward()\n        optimizer.step()\n    \n    # --- Validation and Visualization after each epoch ---\n    model.eval()\n    val_corrects = 0\n    with torch.no_grad():\n        for inputs, labels in val_loader:\n            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n            logits, _ = model(inputs)\n            preds = torch.argmax(logits, 1)\n            val_corrects += torch.sum(preds == labels.data)\n    \n    val_acc = val_corrects.double() / len(val_dataset)\n    print(f\"Validation Accuracy: {val_acc:.4f}\")\n    \n    # Call the new visualization function\n    visualize_attention_by_epoch(model, val_loader, DEVICE, epoch + 1)\n\nprint(\"\\n--- Training Finished ---\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T00:58:31.754248Z","iopub.execute_input":"2025-08-21T00:58:31.754582Z","iopub.status.idle":"2025-08-21T01:11:16.361528Z","shell.execute_reply.started":"2025-08-21T00:58:31.754561Z","shell.execute_reply":"2025-08-21T01:11:16.360733Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **Pretrained Vision Transformer**","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchvision import transforms\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport timm # PyTorch Image Models library\n\nimport os\nfrom PIL import Image\nimport pandas as pd\nfrom tqdm import tqdm\nimport numpy as np\nimport math","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.408230Z","iopub.execute_input":"2025-08-21T02:27:01.408495Z","iopub.status.idle":"2025-08-21T02:27:01.413606Z","shell.execute_reply.started":"2025-08-21T02:27:01.408476Z","shell.execute_reply":"2025-08-21T02:27:01.412731Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 1. Configuration\n# ======================================================================================\nTRAIN_DIR = \"/kaggle/working/train/train\"\nTEST_DIR = \"/kaggle/working/test/test1\"\nSUBMISSION_FILE = \"submission.csv\"\n\n# Model and training parameters\nIMG_SIZE = 224\nBATCH_SIZE = 32\nLEARNING_RATE = 1e-4\nEPOCHS = 5 # Pre-trained models need fewer epochs for fine-tuning\nMODEL_NAME = 'vit_base_patch16_224' # Pre-trained Vision Transformer\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.415139Z","iopub.execute_input":"2025-08-21T02:27:01.415395Z","iopub.status.idle":"2025-08-21T02:27:01.434136Z","shell.execute_reply.started":"2025-08-21T02:27:01.415374Z","shell.execute_reply":"2025-08-21T02:27:01.433584Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 2. Custom Dataset for Cats vs. Dogs (Unchanged)\n# ======================================================================================\nclass CatsDogsDataset(Dataset):\n    def __init__(self, data_dir, transform=None, is_test=False):\n        self.data_dir = data_dir\n        self.transform = transform\n        self.is_test = is_test\n        self.image_files = sorted(os.listdir(self.data_dir), key=lambda x: int(x.split('.')[0]) if self.is_test else (x.split('.')[0], int(x.split('.')[1])))\n\n    def __len__(self):\n        return len(self.image_files)\n\n    def __getitem__(self, idx):\n        img_name = self.image_files[idx]\n        img_path = os.path.join(self.data_dir, img_name)\n        image = Image.open(img_path).convert(\"RGB\")\n        if self.transform:\n            image = self.transform(image)\n        if not self.is_test:\n            label = 0 if 'cat' in img_name else 1\n            return image, torch.tensor(label, dtype=torch.long)\n        else:\n            image_id = int(img_name.split('.')[0])\n            return image, image_id","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.434868Z","iopub.execute_input":"2025-08-21T02:27:01.435098Z","iopub.status.idle":"2025-08-21T02:27:01.456329Z","shell.execute_reply.started":"2025-08-21T02:27:01.435078Z","shell.execute_reply":"2025-08-21T02:27:01.455838Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 3. Image Transformations (Unchanged)\n# ======================================================================================\ndata_transforms = {\n    'train': transforms.Compose([\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.RandomHorizontalFlip(),\n        transforms.RandomRotation(10),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n    'val': transforms.Compose([\n        transforms.Resize((IMG_SIZE, IMG_SIZE)),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n    ]),\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.457014Z","iopub.execute_input":"2025-08-21T02:27:01.457355Z","iopub.status.idle":"2025-08-21T02:27:01.477838Z","shell.execute_reply.started":"2025-08-21T02:27:01.457328Z","shell.execute_reply":"2025-08-21T02:27:01.477317Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 4. Create DataLoaders (Unchanged)\n# ======================================================================================\nfull_train_dataset = CatsDogsDataset(TRAIN_DIR, transform=data_transforms['train'])\ntrain_size = int(0.8 * len(full_train_dataset))\nval_size = len(full_train_dataset) - train_size\ntrain_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size])\nval_dataset.dataset.transform = data_transforms['val'] \ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\nprint(f\"Training samples: {len(train_dataset)}\")\nprint(f\"Validation samples: {len(val_dataset)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.479376Z","iopub.execute_input":"2025-08-21T02:27:01.479855Z","iopub.status.idle":"2025-08-21T02:27:01.548047Z","shell.execute_reply.started":"2025-08-21T02:27:01.479838Z","shell.execute_reply":"2025-08-21T02:27:01.547554Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 5. Instantiate the Pre-trained Model\n# ======================================================================================\nmodel = timm.create_model(MODEL_NAME, pretrained=True, num_classes=2).to(DEVICE)\nmodel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:01.548767Z","iopub.execute_input":"2025-08-21T02:27:01.549391Z","iopub.status.idle":"2025-08-21T02:27:05.323878Z","shell.execute_reply.started":"2025-08-21T02:27:01.549368Z","shell.execute_reply":"2025-08-21T02:27:05.323220Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Model '{MODEL_NAME}' loaded with {sum(p.numel() for p in model.parameters()):,} parameters.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:05.324660Z","iopub.execute_input":"2025-08-21T02:27:05.324897Z","iopub.status.idle":"2025-08-21T02:27:05.329945Z","shell.execute_reply.started":"2025-08-21T02:27:05.324870Z","shell.execute_reply":"2025-08-21T02:27:05.329369Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 6. NEW: Monkey-Patching the Attention Block to Capture Attention\n# ======================================================================================\n\ndef new_forward(self, x):\n    # This is a recreation of the original timm forward pass\n    B, N, C = x.shape\n    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n    q, k, v = qkv.unbind(0)\n\n    attn = (q @ k.transpose(-2, -1)) * self.scale\n    attn = attn.softmax(dim=-1)\n    \n    # *** KEY CHANGE: Save the attention map to an attribute ***\n    self.attn_map = attn\n    \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\n# We use __get__ to bind the new method to the instance of the attention block\nmodel.blocks[-1].attn.forward = new_forward.__get__(model.blocks[-1].attn)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:27:05.330737Z","iopub.execute_input":"2025-08-21T02:27:05.330947Z","iopub.status.idle":"2025-08-21T02:27:05.351387Z","shell.execute_reply.started":"2025-08-21T02:27:05.330932Z","shell.execute_reply":"2025-08-21T02:27:05.350703Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 7. Epoch-by-Epoch Visualization Function (MODIFIED to use new attribute)\n# ======================================================================================\ndef visualize_attention_by_epoch(model, dataloader, device, epoch):\n    print(f\"\\n--- Visualizing Attention for Epoch {epoch} ---\")\n    model.eval()\n    \n    correct_samples, incorrect_samples = [], []\n    \n    with torch.no_grad():\n        for images, labels in dataloader:\n            images, labels = images.to(device), labels.to(device)\n            logits = model(images)\n            preds = torch.argmax(logits, dim=1)\n            \n            # *** KEY CHANGE: Get attention map from our new attribute ***\n            attn_weights = model.blocks[-1].attn.attn_map\n            \n            for i in range(images.size(0)):\n                is_correct = (preds[i] == labels[i]).item()\n                sample = {'image': images[i], 'label': labels[i], 'pred': preds[i], 'attn': attn_weights[i]}\n                if is_correct and len(correct_samples) < 5: correct_samples.append(sample)\n                elif not is_correct and len(incorrect_samples) < 5: incorrect_samples.append(sample)\n            \n            if len(correct_samples) >= 5 and len(incorrect_samples) >= 5: break\n\n    samples_to_show = incorrect_samples + correct_samples\n    samples_to_show = samples_to_show[:10]\n    \n    if not samples_to_show:\n        print(\"Not enough samples to visualize.\")\n        return\n\n    inv_normalize = transforms.Normalize(mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225], std=[1/0.229, 1/0.224, 1/0.225])\n    fig, axes = plt.subplots(2, len(samples_to_show), figsize=(20, 5))\n    \n    for i, sample in enumerate(samples_to_show):\n        img = inv_normalize(sample['image']).cpu().permute(1, 2, 0).numpy()\n        cls_attn = sample['attn'][:, 0, 1:].mean(dim=0)\n        num_patches_side = int(math.sqrt(cls_attn.shape[0]))\n        attn_map = cls_attn.cpu().reshape(num_patches_side, num_patches_side).numpy()\n        \n        axes[0, i].imshow(img)\n        axes[0, i].axis('off')\n        true_label = \"Dog\" if sample['label'].item() == 1 else \"Cat\"\n        pred_label = \"Dog\" if sample['pred'].item() == 1 else \"Cat\"\n        color = \"green\" if true_label == pred_label else \"red\"\n        axes[0, i].set_title(f\"True: {true_label}\\nPred: {pred_label}\", color=color)\n        \n        axes[1, i].imshow(attn_map, cmap='viridis')\n        axes[1, i].axis('off')\n\n    plt.suptitle(f\"Attention Maps - Epoch {epoch}\", fontsize=16)\n    plt.tight_layout(rect=[0, 0, 1, 0.96])\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:28:51.728952Z","iopub.execute_input":"2025-08-21T02:28:51.729224Z","iopub.status.idle":"2025-08-21T02:28:51.739238Z","shell.execute_reply.started":"2025-08-21T02:28:51.729204Z","shell.execute_reply":"2025-08-21T02:28:51.738562Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================================\n# 8. Training Loop (Unchanged)\n# ======================================================================================\ncriterion = nn.CrossEntropyLoss()\noptimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)\n\nfor epoch in range(EPOCHS):\n    print(f\"\\n--- Epoch {epoch+1}/{EPOCHS} ---\")\n    model.train()\n    for inputs, labels in tqdm(train_loader, desc=\"Training\"):\n        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n        optimizer.zero_grad()\n        logits = model(inputs)\n        loss = criterion(logits, labels)\n        loss.backward()\n        optimizer.step()\n    \n    model.eval()\n    val_corrects = 0\n    with torch.no_grad():\n        for inputs, labels in val_loader:\n            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)\n            logits = model(inputs)\n            preds = torch.argmax(logits, 1)\n            val_corrects += torch.sum(preds == labels.data)\n    \n    val_acc = val_corrects.double() / len(val_dataset)\n    print(f\"Validation Accuracy: {val_acc:.4f}\")\n    \n    visualize_attention_by_epoch(model, val_loader, DEVICE, epoch + 1)\n\nprint(\"\\n--- Training Finished ---\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-21T02:28:54.391506Z","iopub.execute_input":"2025-08-21T02:28:54.391880Z","iopub.status.idle":"2025-08-21T03:26:29.751335Z","shell.execute_reply.started":"2025-08-21T02:28:54.391853Z","shell.execute_reply":"2025-08-21T03:26:29.750439Z"}},"outputs":[],"execution_count":null}]}