{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","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"},{"sourceId":9203638,"sourceType":"datasetVersion","datasetId":4522892},{"sourceId":132818,"sourceType":"modelInstanceVersion","modelInstanceId":112260,"modelId":134029},{"sourceId":132978,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":112396,"modelId":134029},{"sourceId":132999,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":112414,"modelId":134029}],"dockerImageVersionId":30761,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Vision Transformer","metadata":{}},{"cell_type":"markdown","source":"## Overview\n\nThis notebook is my PyTorch implementation of a Vision Transformer (ViT) model, based on a research paper \"[An Image Is Worth 16x16 Words:\nTransformers For Image Recognition At Scale](https://arxiv.org/abs/2010.11929)\".\\\nViTs divide images into a sequence of patches and apply transformer architectures, originally designed for text, to image recognition.","metadata":{}},{"cell_type":"markdown","source":"### Models and parameters\n|Model|Layers|Embedding size|MLP size|Heads|\n|:--:|:---:|:---:|:---:|:---:|\n|ViT-T|12|192|768|4|\n|ViT-S|12|384|1536|6|\n|ViT-S-12|12|384|1536|12","metadata":{}},{"cell_type":"markdown","source":"## Install additional dependencies","metadata":{"_uuid":"44749934-0e25-437c-871c-4ce75663f31e","_cell_guid":"2abd5426-0af2-4cdb-ba3d-9d148c41e567","trusted":true}},{"cell_type":"code","source":"!pip install torchtyping","metadata":{"_uuid":"8e598397-0557-439f-a078-63a4b2d77e20","_cell_guid":"49673c87-7f1d-4ff8-88ee-1734d99d6668","execution":{"iopub.status.busy":"2024-10-29T08:40:59.906674Z","iopub.execute_input":"2024-10-29T08:40:59.907340Z","iopub.status.idle":"2024-10-29T08:41:14.649110Z","shell.execute_reply.started":"2024-10-29T08:40:59.907282Z","shell.execute_reply":"2024-10-29T08:41:14.648038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Imports","metadata":{"_uuid":"55dcd9d4-7a09-4fe8-a809-dea59bf3180c","_cell_guid":"0e54622c-3dc2-4746-86e3-db211ffe066b","trusted":true}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport cv2","metadata":{"_uuid":"1653acba-b489-4cf5-ad43-6caf432f666e","_cell_guid":"1589177f-4b15-487b-91a0-cb9c04e44d97","execution":{"iopub.status.busy":"2024-10-29T08:41:14.651405Z","iopub.execute_input":"2024-10-29T08:41:14.651818Z","iopub.status.idle":"2024-10-29T08:41:15.307581Z","shell.execute_reply.started":"2024-10-29T08:41:14.651774Z","shell.execute_reply":"2024-10-29T08:41:15.306780Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nfrom torchtyping import TensorType\n\nfrom torchvision.datasets import ImageFolder\nfrom torchvision.transforms import v2","metadata":{"_uuid":"e36702ab-d96a-4ff1-ab58-fcf271c26dc9","_cell_guid":"ed0d9026-0281-45bf-a5b3-d87d79610330","execution":{"iopub.status.busy":"2024-10-29T08:41:15.308993Z","iopub.execute_input":"2024-10-29T08:41:15.309407Z","iopub.status.idle":"2024-10-29T08:41:21.428012Z","shell.execute_reply.started":"2024-10-29T08:41:15.309373Z","shell.execute_reply":"2024-10-29T08:41:21.426927Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datetime import date","metadata":{"_uuid":"c5ed2218-a110-458a-bf3c-87c5bc9fa26d","_cell_guid":"58e0e0b1-1d97-4652-bcae-44f05d800d95","execution":{"iopub.status.busy":"2024-10-29T08:41:21.430540Z","iopub.execute_input":"2024-10-29T08:41:21.431490Z","iopub.status.idle":"2024-10-29T08:41:21.436273Z","shell.execute_reply.started":"2024-10-29T08:41:21.431453Z","shell.execute_reply":"2024-10-29T08:41:21.435385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Vision Transformer Architecture","metadata":{"_uuid":"1b9abcda-cc7b-43b6-8fc8-db31346baa41","_cell_guid":"8b9023e0-8442-4107-8688-ea0e54ce3561","trusted":true}},{"cell_type":"markdown","source":"### Patch and position embedding\n\nTransformers receive a sequence as input. To feed an image into a transformer, the following steps are followed:\n 1. We first use `unfold()` to split the image into non-overlapping `16x16` patches, each flattened into a `1D` vector. For an image of size `224x224`, this results in `n_patches = (224/16) * (224/16) = 196` patches.\n 2. Next, a trainable `linear_projection()` layer is used to map each flattened patch into an embedding of size `embed_dim`.\n 3. A class token is prepended to the sequence, so the input to the transformer has the shape `(batch_size, n_patches + 1, embed_dim)`, where the additional token is used to aggregate information during classification.\n 4. Finally, we add trainable position encodings to each patch and the class token, which provides information about the relative position of each patch in the original image.","metadata":{}},{"cell_type":"code","source":"class PatchAndPositionEmbedding(nn.Module):\n    \"\"\"\n    Combines patch and position embeddings for Vision Transformer (ViT) models.\n\n    Args:\n        image_size (int, optional): Input image size (assumed square). Default is 224.\n        patch_size (int, optional): Patch size. Default is 16.\n        embed_dim (int, optional): Dimensionality of patch embeddings. Default is 192.\n        n_channels (int, optional): Number of input channels (e.g., 3 for RGB). Default is 3.\n        dropout (float, optional): Dropout rate applied after embeddings. Default is 0.1.\n\n    Attributes:\n        n_patches (int): Number of patches in the input image.\n        linear_projection (nn.Linear): Linear layer to project flattened patches to the embedding dimension.\n        cls_token (nn.Parameter): Learnable class token appended to the patch embeddings.\n        pos_embedding (nn.Parameter): Learnable positional embeddings added to the patch and class token embeddings.\n        unfold (nn.Unfold): Unfolding layer to extract non-overlapping patches from the image.\n        dropout (nn.Dropout): Dropout layer applied to the final patch embeddings.\n    \"\"\"\n    \n    def __init__(\n        self,\n        image_size: int = 224,\n        patch_size: int = 16,\n        embed_dim: int = 192,\n        n_channels: int = 3,\n        dropout: float = 0.1\n    ) -> None:\n        \"\"\"\n        Initializes the PatchAndPositionEmbedding module.\n\n        Args:\n            image_size (int, optional): Size of the input image (assumed square). Default is 224.\n            patch_size (int, optional): Size of each patch to be extracted from the image. Default is 16.\n            embed_dim (int, optional): Dimensionality of the patch embeddings. Default is 192.\n            n_channels (int, optional): Number of input channels (e.g., 3 for RGB). Default is 3.\n            dropout (float, optional): Dropout rate applied to the patch and position embeddings. Default is 0.1.\n        \"\"\"\n        \n        super(PatchAndPositionEmbedding, self).__init__()\n        \n        assert image_size % patch_size == 0, f\"image_size of {image_size} is not divisible by patch_size of {patch_size}\"\n        \n        self.n_patches = image_size * image_size // patch_size ** 2\n        self.image_size = image_size\n        self.patch_size = patch_size\n        self.embed_dim = embed_dim\n        self.n_channels = n_channels\n        \n        self._unfold = nn.Unfold(kernel_size=self.patch_size, stride=self.patch_size)\n        self._linear_projection = nn.Linear(self.patch_size ** 2 * self.n_channels, self.embed_dim)\n        self._cls_token = nn.Parameter(torch.zeros(1, self.embed_dim))\n        self._pos_embedding = nn.Parameter(torch.zeros(self.n_patches + 1, self.embed_dim))\n        self._dropout = nn.Dropout(dropout)\n\n        \n    def forward(self, x: TensorType[torch.float32]) -> TensorType[torch.float32]:\n        \"\"\"\n        Forward pass that computes patch embeddings, prepends the class token, and adds positional embeddings.\n\n        Args:\n            x (TensorType[torch.float32]): Input tensor of shape (N, C, H, W), where N is the batch size, \n                                           C is the number of channels, and H, W are the height and width of the image.\n\n        Returns:\n            TensorType[torch.float32]: Tensor of shape (batch_size, n_patches + 1, embed_dim), where `n_patches + 1` includes \n                                       the class token, and `embed_dim` is the embedding dimension.\n        \"\"\"\n        \n        # patch embedding\n        x = self._unfold(x).transpose(1, 2)\n        x = self._linear_projection(x)\n        \n        # prepending class token\n        batch_size = x.shape[0]\n        \n        cls_tokens = self._cls_token.expand(batch_size, -1, -1)\n        x = torch.cat((cls_tokens, x), dim=1)\n        \n        # positional embedding\n        pos_embeddings = self._pos_embedding.expand(batch_size, -1, -1)\n        x = x + pos_embeddings\n        \n        x = self._dropout(x)\n\n        return x","metadata":{"_uuid":"3fc513dd-1e01-4827-8e8b-de63d5d9bbcb","_cell_guid":"b31cef00-8677-4392-b88d-e9dc7883c291","execution":{"iopub.status.busy":"2024-10-29T08:41:21.437867Z","iopub.execute_input":"2024-10-29T08:41:21.438524Z","iopub.status.idle":"2024-10-29T08:41:21.457146Z","shell.execute_reply.started":"2024-10-29T08:41:21.438480Z","shell.execute_reply":"2024-10-29T08:41:21.456155Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Vision Transformer (ViT) model\n\nTo build a Vision Transformer (ViT) model, we'll use our `PatchAndPositionEmbedding` class along with the `nn.TransformerEncode` module from PyTorch.\n1. **Patch and Position Embeddings**: The patch embeddings generated by `PatchAndPositionEmbedding` are fed into 12 layers of Transformer Encoder.\n2. **Transformer Encoder Layers**: Each Transformer Encoder layer processes the sequence of patch embeddings, and the output shape is `(n_patches + 1, embed_dim)`, where `n_patches + 1` includes the class token, and `embed_dim` is the embedding dimension.\n3. **Output**: After passing through the Transformer Encoder, the output has the shape `(n_patches + 1, n_classes)` (ignoring the batch_size for simplicity). The class token, located in the first position of the sequence, aggregates information from all patches.\n4. **MLP Head**:\n  - We select the class token and pass it through an MLP Head.\n  - The MLP Head consists of a `nn.LayerNorm` layer, followed by a fully connected layers.","metadata":{}},{"cell_type":"code","source":"class VisionTransformer(nn.Module):\n    \"\"\"\n    Vision Transformer (ViT) model for image classification.\n\n    Args:\n        image_size (int, optional): Input image size (square). Default is 224.\n        patch_size (int, optional): Size of image patches. Default is 16.\n        embed_dim (int, optional): Dimensionality of patch embeddings. Default is 192.\n        n_layers (int, optional): Number of transformer encoder layers. Default is 12.\n        n_heads (int, optional): Number of attention heads. Default is 4.\n        mlp_size (int, optional): Size of the MLP layer. Default is 768.\n        n_classes (int, optional): Number of output classes. Default is 10.\n        n_channels (int, optional): Number of input channels (e.g., 3 for RGB). Default is 3.\n        dropout (float, optional): Dropout rate. Default is 0.1.\n        batch_first (bool, optional): If True, the batch dimension comes first. Default is False.\n\n    Attributes:\n        embedding (PatchAndPositionEmbedding): Patch and position embedding layer.\n        transformer_encoder (nn.TransformerEncoder): Transformer encoder with multiple layers.\n        MLP (nn.Sequential): Final classification layer with normalization and linear layers.\n    \"\"\"\n    \n    def __init__(\n        self,\n        image_size: int = 224,\n        patch_size: int = 16,\n        embed_dim: int = 192,\n        n_layers: int = 12,\n        n_heads: int = 4,\n        mlp_size: int = 768,\n        n_classes: int = 10,\n        n_channels: int = 3,\n        dropout: float = 0.1,\n        batch_first: bool = False\n    ) -> None:\n        \"\"\"\n        Initializes the VisionTransformer model.\n\n        Args:\n            image_size (int, optional): Size of the input image. Default is 224.\n            patch_size (int, optional): Size of image patches. Default is 16.\n            embed_dim (int, optional): Dimensionality of patch embeddings. Default is 192.\n            n_layers (int, optional): Number of transformer encoder layers. Default is 12.\n            n_heads (int, optional): Number of attention heads. Default is 4.\n            mlp_size (int, optional): Size of the MLP layer. Default is 768.\n            n_classes (int, optional): Number of output classes. Default is 10.\n            n_channels (int, optional): Number of input channels. Default is 3.\n            dropout (float, optional): Dropout rate. Default is 0.1.\n            batch_first (bool, optional): If True, batch dimension comes first. Default is False.\n        \"\"\"\n        \n        super(VisionTransformer, self).__init__()\n        \n        self._batch_first = batch_first\n        \n        self._embedding = PatchAndPositionEmbedding(image_size, patch_size, embed_dim, n_channels)\n        \n        encoder_layer = nn.TransformerEncoderLayer(\n            d_model=embed_dim,\n            nhead=n_heads,\n            dim_feedforward=mlp_size,\n            dropout=dropout,\n            activation='gelu',\n            batch_first=self._batch_first\n        )\n        self._transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)\n        \n        self._MLP = nn.Sequential(\n            nn.LayerNorm(embed_dim),\n            nn.Linear(embed_dim, embed_dim),\n            nn.ReLU(),\n            nn.Linear(embed_dim, n_classes)\n        )\n        \n        self._init_weights()\n                \n    def _init_weights(self):\n        \"\"\"\n        Initializes the weights of the model, including the class token, positional embedding, \n        and the linear projection layer.\n        \"\"\"\n        \n        nn.init.trunc_normal_(self._embedding._cls_token, std=0.02)\n        nn.init.trunc_normal_(self._embedding._pos_embedding, std=0.02)\n        nn.init.xavier_normal_(self._embedding._linear_projection.weight)\n        nn.init.zeros_(self._embedding._linear_projection.bias)\n    \n                    \n    def forward(self, x: TensorType[torch.float32]) -> TensorType[torch.float32]:\n        \"\"\"\n        Forward pass of the VisionTransformer model.\n\n        Args:\n            x (TensorType[torch.float32]): Input image tensor of shape (C, H, W), where C is the number of channels, \n                                           and H, W are the height and width of the image.\n\n        Returns:\n            TensorType[torch.float32]: Output logits of shape (batch_size, n_classes), representing the class scores.\n        \"\"\"\n        \n        x = self._embedding(x)\n        \n        if not self._batch_first: \n            x = x.transpose(0, 1)\n        \n        x = self._transformer_encoder(x)\n        \n        cls_token = x[:, 0] if self._batch_first else x[0]\n        \n        logits = self._MLP(cls_token)\n        \n        return logits","metadata":{"_uuid":"3b7d6edb-d3fc-4b93-9cd2-d814bc3abed3","_cell_guid":"e6daa07b-4bbe-4e92-8461-f35f47f122e8","execution":{"iopub.status.busy":"2024-10-29T08:41:21.458630Z","iopub.execute_input":"2024-10-29T08:41:21.459576Z","iopub.status.idle":"2024-10-29T08:41:21.475959Z","shell.execute_reply.started":"2024-10-29T08:41:21.459491Z","shell.execute_reply":"2024-10-29T08:41:21.475285Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Trainer class\n\nTo train our model, I’ve implemented a Trainer class with the following features:\n1. **Multi-GPU Support**: The class utilizes `nn.DataParallel` for multi-GPU training. This allows the model to be trained on multiple GPUs in parallel. Note that `torch.nn.parallel.DistributedDataParallel` isn't used here because it requires multiprocessing, which doesn’t work well in Jupyter environments.\n2. **Loss Function**: The class employs `nn.CrossEntropyLoss`, a suitable loss function for classification tasks, to compute the loss during training.\n3. **Training Process**:\n  - The model is trained for a user-specified number of epochs.\n  - After each epoch, the model is evaluated to track performance over time.\n4. **Model Saving**: Users can provide a path to save the trained model once the training process is complete.\n5. **Statistics Plotting**: Users can plot the model's training loss, validation loss, and validation accuracy.","metadata":{}},{"cell_type":"code","source":"class Trainer():\n    \"\"\"\n    Trainer class for managing the training and testing of a PyTorch model.\n\n    Args:\n        model (nn.Module): The neural network model to be trained.\n        dataloaders (tuple[DataLoader, DataLoader]): A tuple containing the training and testing DataLoaders.\n        optimizer (optim.Optimizer): The optimizer used for training the model.\n        device (str | torch.device): The device (e.g., 'cpu', 'cuda') to which the model and data should be moved.\n        log_interval (int, optional): The interval (in batches) for logging training loss. Default is 10.\n\n    Attributes:\n        device (str | torch.device): The device used for training.\n        _model (nn.DataParallel): The model wrapped in a DataParallel module for multi-GPU training.\n        _train_dataloader (DataLoader): The DataLoader for the training dataset.\n        _test_dataloader (DataLoader): The DataLoader for the testing dataset.\n        _optimizer (optim.Optimizer): The optimizer used for updating the model weights.\n        _loss_fn (nn.CrossEntropyLoss): The loss function used for training.\n        _softmax (nn.Softmax): The softmax layer for converting logits to probabilities.\n        _train_loss (list[float]): A list storing the training loss for each epoch.\n        _test_loss (list[float]): A list storing the test loss for each epoch.\n        _log_interval (int): The interval (in batches) for logging training loss. Default is 10.\n        _test_accuracy (list[float]): A list storing the test accuracy for each epoch.\n    \"\"\"\n    \n    def __init__(\n        self,\n        model: nn.Module,\n        dataloaders: tuple[DataLoader, DataLoader],\n        optimizer: optim.Optimizer,\n        device: str | torch.device,\n        log_interval: int = 10\n    ) -> None:\n        \"\"\"\n        Initializes the Trainer with the given model, dataloaders, optimizer, and device.\n\n        Args:\n            model (nn.Module): The neural network model to be trained.\n            dataloaders (tuple[DataLoader, DataLoader]): A tuple containing the training and testing DataLoaders.\n            optimizer (optim.Optimizer): The optimizer used for training the model.\n            device (str | torch.device): The device (e.g., 'cpu', 'cuda') to which the model and data should be moved.\n            log_interval (int, optional): The interval (in batches) for logging training loss. Default is 10.\n        \"\"\"\n        \n        self.device = device\n        self._model = nn.DataParallel(model.to(self.device))\n        self._train_dataloader, self._test_dataloader = dataloaders\n        self._optimizer = optimizer\n        self._loss_fn = nn.CrossEntropyLoss()\n        self._softmax = nn.Softmax(dim=1)\n        \n        self._log_interval = log_interval\n        self._train_loss = []\n        self._test_loss = []\n        self._test_accuracy = []\n        \n    def _train_epoch(self) -> None:\n        \"\"\"\n        Executes one epoch of training on the training dataset.\n\n        This method iterates over the training DataLoader, computes the loss, \n        performs backpropagation, and updates the model weights. It logs the \n        loss every 'log_interval' batches and appends the average loss for the \n        epoch to the training loss list.\n        \"\"\"\n        \n        avg_loss = 0.0\n        avg_epoch_loss = 0.0\n        n_batches = len(self._train_dataloader)\n        n_samples = len(self._train_dataloader.dataset)\n\n        self._model.train()\n        for batch_index, (X, y) in enumerate(self._train_dataloader, 1):\n            X, y = X.to(self.device), y.to(self.device)\n\n            prediction = self._model(X)\n            loss = self._loss_fn(prediction, y)\n\n            avg_loss += loss.item()\n            avg_epoch_loss += loss.item()\n\n            self._optimizer.zero_grad()\n            loss.backward()\n            self._optimizer.step()\n            \n            \n\n            if batch_index % self._log_interval == 0:\n                loss = loss.item()\n                current_sample = batch_index * len(X)\n\n                print(\n                    f'Current loss: {loss:.5f}, ' \\\n                    f'Average loss across 10 batches: {avg_loss / self._log_interval:.5f} ' \\\n                    f'[{current_sample} / {n_samples}]'\n                )\n                avg_loss = 0.0\n                \n        self._train_loss.append(avg_epoch_loss / n_batches)\n                \n    def _test_epoch(self) -> None:\n        \"\"\"\n        Evaluates the model on the testing dataset.\n\n        This method sets the model to evaluation mode, computes the average test loss, \n        and calculates the accuracy over the entire test set. The average loss and \n        accuracy are appended to the respective lists for plotting and analysis.\n        \"\"\"\n        \n        self._model.eval()\n        \n        avg_test_loss = 0.0\n        n_correct = 0\n        n_batches = len(self._test_dataloader)\n        n_samples = len(self._test_dataloader.dataset)\n\n        for X, y in self._test_dataloader:\n            X, y = X.to(self.device), y.to(self.device)\n\n            with torch.no_grad():\n                prediction = self._model(X)\n\n                loss = self._loss_fn(prediction, y)\n                avg_test_loss += loss.item()\n\n                prediction = self._softmax(prediction)\n                n_correct += (torch.argmax(prediction, dim=1) == y).sum().item()\n                \n        avg_test_loss /= n_batches\n        accuracy = n_correct / n_samples * 100\n\n        print(\n            f'Average test loss: {avg_test_loss:.5f}, ' \\\n            f'Accuracy: {accuracy:.2f}%'\n        )\n        \n        self._test_loss.append(avg_test_loss)\n        self._test_accuracy.append(accuracy)\n\n        \n    def train(self, n_epochs: int, path: str | None, plot: bool = True) -> None:\n        \"\"\"\n        Trains the model for a specified number of epochs.\n\n        This method trains the model for the given number of epochs, evaluates \n        it after each epoch, and optionally saves the model to a file. It can \n        also plot the training statistics after training if requested.\n\n        Args:\n            n_epochs (int): The number of epochs to train the model.\n            path (str | None): The file path to save the model after training. If None, the model is not saved.\n            plot (bool, optional): Whether to plot the training statistics after training. Default is True.\n        \"\"\"\n        \n        for epoch in range(n_epochs):\n            print(f\"|--------------------------{epoch + 1}/{n_epochs}--------------------------|\")\n            self._train_epoch()\n            self._test_epoch()\n            \n        if path:\n            self.save_model(path)\n            \n        if plot:\n            self.plot_statistics()\n            \n            \n    def save_model(self, path: str) -> None:\n        \"\"\"\n        Saves the model state dictionary to the specified file path.\n\n        Args:\n            path (str): The file path where the model will be saved.\n        \"\"\"\n        \n        state = self._model.module.state_dict()\n        torch.save(state, path)\n        print(f'Model saved in: {path}')\n        \n    \n    def plot_statistics(self) -> None:\n        \"\"\"\n        Plots the training and testing loss over epochs, as well as test accuracy.\n\n        This method creates two plots:\n        1. A line plot showing both the training and testing loss over each epoch.\n        2. A line plot showing the testing accuracy over each epoch.\n\n        These plots provide insight into the model's performance during training.\n\n        Attributes used:\n            _train_loss (list[float]): List of average training loss values for each epoch.\n            _test_loss (list[float]): List of average testing loss values for each epoch.\n            _test_accuracy (list[float]): List of testing accuracy values for each epoch.\n        \"\"\"\n        \n        fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))\n        fig.suptitle(\"Train Statstics\")\n                \n        #plot train and test loss\n        line1, = ax1.plot(self._train_loss, label=\"train loss\")\n        line2, = ax1.plot(self._test_loss, label=\"test loss\")\n        ax1.legend(handles=[line1, line2])\n        ax1.set_title(\"Train/Test Loss\")\n        ax1.set_ylabel(\"Loss\")\n        ax1.set_xlabel(\"Epoch\")\n\n        \n        #plot accuracy\n        ax2.plot(self._test_accuracy)\n        ax2.set_title(\"Test Accuracy\")\n        ax2.set_ylabel(\"Accuracy (%)\")\n        ax2.set_xlabel(\"Epoch\")\n        \n        plt.show()\n        ","metadata":{"_uuid":"54f3e327-0a70-47e0-a15f-00fc427ca7ce","_cell_guid":"0da97fc4-3404-4f83-ab65-18d8d78fe794","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:21.477535Z","iopub.execute_input":"2024-10-29T08:41:21.478275Z","iopub.status.idle":"2024-10-29T08:41:21.505773Z","shell.execute_reply.started":"2024-10-29T08:41:21.478232Z","shell.execute_reply":"2024-10-29T08:41:21.504764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preparing the data","metadata":{"_uuid":"25ed9490-4494-480d-abd0-525d8072af17","_cell_guid":"4c3e250e-1b93-4af4-a9e9-5e42e4f6f95d","trusted":true}},{"cell_type":"markdown","source":"### Defining Transforms\n\nWe define transformations for both the training and testing datasets to ensure consistency and prepare the images for input into the model:\n1. **Resizing**: All images are resized to `224x224` pixels to match the input size required by the Vision Transformer (ViT) model.\n2. **Training Data Augmentation**: For the training data, we apply a 50% chance of a random vertical flip to introduce data augmentation and increase model robustness to variations.\n3. **Common Transforms**:\n  - Both the training and testing images are converted from `PIL` Image objects into tensors.\n  - The tensors are cast to `torch.float32` to ensure compatibility with the model.\n  - Finally, pixel values are scaled to the range `[0, 1]` for both training and testing datasets.","metadata":{}},{"cell_type":"code","source":"train_transforms = v2.Compose([\n    v2.Resize((224, 224)),\n    v2.RandomVerticalFlip(p=0.5),\n    v2.RandomHorizontalFlip(p=0.5),\n    v2.RandomRotation((-90, 90)),\n    v2.RandomGrayscale(p=0.1),\n    v2.PILToTensor(),\n    v2.ToDtype(torch.float32, scale=True)\n])\n\ntest_transforms = v2.Compose([\n    v2.Resize((224, 224)),\n    v2.PILToTensor(),\n    v2.ToDtype(torch.float32, scale=True)\n])","metadata":{"_uuid":"dbb97c7c-fbb0-4220-864d-2566eeac1743","_cell_guid":"16eb8a69-173a-4c90-8f8e-7d26b255049c","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:21.506989Z","iopub.execute_input":"2024-10-29T08:41:21.507613Z","iopub.status.idle":"2024-10-29T08:41:21.520943Z","shell.execute_reply.started":"2024-10-29T08:41:21.507564Z","shell.execute_reply":"2024-10-29T08:41:21.519875Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Data Loading\n\nWe load the training and testing data using PyTorch's `ImageFolder`, which is suitable because each class of images is organized into its own directory. We pass the paths to the train and test directories along with the previously defined transformations.\n\nThe dataset used to train the model is the [Packed Fruits and Vegetables Recognition Benchmark](https://www.kaggle.com/datasets/sergeynesteruk/packed-fruits-and-vegetables-recognition-benchmark) dataset.","metadata":{}},{"cell_type":"code","source":"train_dataset = ImageFolder(\n    root='/kaggle/input/packed-fruits-and-vegetables-recognition-benchmark/train/train',\n    transform=train_transforms\n)\ntest_dataset = ImageFolder(\n    root='/kaggle/input/packed-fruits-and-vegetables-recognition-benchmark/test/test',\n    transform=test_transforms\n)","metadata":{"_uuid":"226d0461-d8ed-4d7a-b03b-67f4b17c94ad","_cell_guid":"47623a93-0955-4bb4-86ca-9bf2648bf283","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:21.522409Z","iopub.execute_input":"2024-10-29T08:41:21.523072Z","iopub.status.idle":"2024-10-29T08:41:50.816833Z","shell.execute_reply.started":"2024-10-29T08:41:21.522955Z","shell.execute_reply":"2024-10-29T08:41:50.815810Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_dataset), len(test_dataset)","metadata":{"execution":{"iopub.status.busy":"2024-10-29T08:41:50.820744Z","iopub.execute_input":"2024-10-29T08:41:50.821084Z","iopub.status.idle":"2024-10-29T08:41:50.827957Z","shell.execute_reply.started":"2024-10-29T08:41:50.821050Z","shell.execute_reply":"2024-10-29T08:41:50.826976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Preparing the dataloaders\n\nWe define a simple `get_dataloaders()` function to create and return PyTorch `DataLoader`s for both the training and testing datasets. The function sets the following parameters:\n\n1. **Batch Size**: Set to `256` for efficient batch processing.\n2. **Number of Workers**: Set to `4` to parallelize data loading and speed up the process.\n3. **Memory Pinning**: `pin_memory=True` is enabled to improve performance when transferring data to the GPU.\n4. **Shuffling**: For the training data, `shuffle=True` is used to randomize the data at the start of each epoch, ensuring better generalization.\n\nThe function returns a tuple containing the train and test DataLoaders.","metadata":{}},{"cell_type":"code","source":"def get_dataloaders(batch_size: int = 256, shuffle_test: bool = False) -> tuple[DataLoader, DataLoader]:\n    \"\"\"\n    Returns the DataLoader objects for the training and testing datasets.\n\n    Args:\n        batch_size (int, optional): The number of samples per batch to load. Default is 256.\n        shuffle_test (bool, optional): If True, shuffles the test dataset. Default is False.\n\n    Returns:\n        tuple[DataLoader, DataLoader]: A tuple containing:\n            - DataLoader for the training dataset (with shuffling enabled).\n            - DataLoader for the testing dataset (with optional shuffling based on 'shuffle_test').\n\n    The DataLoader for the training set has shuffle enabled, and both DataLoaders \n    use pin_memory for faster data transfer to GPU and 4 worker threads for loading the data.\n    \"\"\"\n    \n    return (\n        DataLoader(\n            train_dataset,\n            batch_size=batch_size,\n            shuffle=True,\n            pin_memory=True,\n            num_workers=4,\n        ),\n        DataLoader(\n            test_dataset,\n            batch_size=batch_size,\n            shuffle=shuffle_test,\n            pin_memory=True,\n            num_workers=4,\n        )\n    )","metadata":{"_uuid":"3687cbcf-7b2e-4445-ad34-e3f78a1f606a","_cell_guid":"087bd2ae-566e-4107-8f32-7cea3ff1f7c0","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:50.829375Z","iopub.execute_input":"2024-10-29T08:41:50.829943Z","iopub.status.idle":"2024-10-29T08:41:50.837497Z","shell.execute_reply.started":"2024-10-29T08:41:50.829909Z","shell.execute_reply":"2024-10-29T08:41:50.836761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X, y = test_dataset[0]\nX.shape","metadata":{"_uuid":"4600e623-1bfa-43a5-8262-28fcbd71e78d","_cell_guid":"ceb99711-d668-4a1e-96e4-14dc7d97be5c","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:50.838622Z","iopub.execute_input":"2024-10-29T08:41:50.838904Z","iopub.status.idle":"2024-10-29T08:41:50.949107Z","shell.execute_reply.started":"2024-10-29T08:41:50.838874Z","shell.execute_reply":"2024-10-29T08:41:50.948194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here we can see the first `5` classes from our dataset:","metadata":{}},{"cell_type":"code","source":"idx_to_class_train = train_dataset.find_classes(train_dataset.root)[0]\nidx_to_class_test = test_dataset.find_classes(test_dataset.root)[0]\nidx_to_class_train[:5]","metadata":{"_uuid":"58d2a62c-1de6-4778-9fe4-45a79af12855","_cell_guid":"89e50a17-015f-427c-98c9-b3bfc9dc271e","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:50.950258Z","iopub.execute_input":"2024-10-29T08:41:50.950665Z","iopub.status.idle":"2024-10-29T08:41:50.958247Z","shell.execute_reply.started":"2024-10-29T08:41:50.950620Z","shell.execute_reply":"2024-10-29T08:41:50.957278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Patch embedding visualization\n\nThe following code demonstrates how patch embedding works by visualizing the process. On the left, we display the original image, and on the right, the image is divided into `16x16` patches, showing how the Vision Transformer (ViT) splits the input into smaller, fixed-size sections for further processing. This visualization helps illustrate how ViT operates on image patches to capture visual features.","metadata":{}},{"cell_type":"code","source":"fig = plt.figure(layout='constrained', figsize=(10, 5))\nsubfig1, subfig2 = fig.subfigures(1, 2)\n\n# display original image\nax = subfig1.add_axes((0, 0, 0.95, 0.95))\nax.imshow(X.permute(1, 2, 0))\nax.set_xticks([])\nax.set_yticks([])\nsubfig1.suptitle('Original')\n\n# get patches\npatches = PatchAndPositionEmbedding()._unfold(X).transpose(0, 1)\nn_patches = 224 // 16\n\n# display patches\naxs = subfig2.subplots(nrows=n_patches, ncols=n_patches).flatten()\nfor i, patch in enumerate(patches):\n    axs[i].imshow(patch.reshape(3, 16, 16).permute(1, 2, 0))\n    axs[i].set_xticks([])\n    axs[i].set_yticks([])\n\nsubfig2.suptitle('Divided into patches')\n\nfig.suptitle(f\"{idx_to_class_test[y]} before and after patch embedding\")\n\nplt.show()","metadata":{"_uuid":"85db156d-1c8c-4d7d-8a98-fa730b3e1985","_cell_guid":"e0baf018-1b85-4992-bcbb-7f0975999a3d","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:41:50.959543Z","iopub.execute_input":"2024-10-29T08:41:50.959817Z","iopub.status.idle":"2024-10-29T08:42:05.446818Z","shell.execute_reply.started":"2024-10-29T08:41:50.959781Z","shell.execute_reply":"2024-10-29T08:42:05.445953Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training the ViT","metadata":{"_uuid":"0de2e9bb-b004-4541-a704-3d1caaf3d9d7","_cell_guid":"0e38558d-7c2d-4c7d-82cb-6f7851c60eed","trusted":true}},{"cell_type":"markdown","source":"### Hyperparameters\n\nFinally, it's time to train the model. We begin by creating an instance of the `VisionTransformer` class, configuring it with the desired parameters. We set the following training settings:\n\n1. **Number of Epochs**: Set to `10`.\n2. **Learning Rate**: Set to `8e-5`.\n3. **Batch Size**: Set to `256`.\n4. **Optimizer**: We use the AdamW optimizer, which is well-suited for transformer models.\n5. **Weight Decay**: A weight decay of `0.1`, as mentioned in the original Vision Transformer paper, is applied.","metadata":{}},{"cell_type":"code","source":"ViT = VisionTransformer(embed_dim=384, n_heads=6, mlp_size=1536, batch_first=True, n_classes=len(idx_to_class_train))","metadata":{"_uuid":"6f21412a-537e-47e6-bed3-03806a15c078","_cell_guid":"65ab7497-95db-4299-bc13-6a4089a8a11c","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:42:05.447848Z","iopub.execute_input":"2024-10-29T08:42:05.448119Z","iopub.status.idle":"2024-10-29T08:42:05.574585Z","shell.execute_reply.started":"2024-10-29T08:42:05.448088Z","shell.execute_reply":"2024-10-29T08:42:05.573673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"N_EPOCHS = 10\nLEARNING_RATE = 8e-5\nWEIGHT_DECAY = 0.1\n\noptimizer = optim.AdamW(ViT.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)","metadata":{"_uuid":"b0dbdf84-30f6-4089-a2b3-55e80e61e761","_cell_guid":"53be63ef-9faf-4db7-9b55-3ecd201cbdb0","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2024-10-29T08:42:05.575675Z","iopub.execute_input":"2024-10-29T08:42:05.575979Z","iopub.status.idle":"2024-10-29T08:42:05.581927Z","shell.execute_reply.started":"2024-10-29T08:42:05.575947Z","shell.execute_reply":"2024-10-29T08:42:05.580614Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = \"cuda\"\ntrainer = Trainer(\n    model=ViT,\n    dataloaders=get_dataloaders(),\n    optimizer=optimizer,\n    device=device\n)\ntrainer.train(n_epochs=N_EPOCHS, path=f'ViT-Sa-{date.today().isoformat()}.pth')","metadata":{"execution":{"iopub.status.busy":"2024-10-29T08:42:05.583165Z","iopub.execute_input":"2024-10-29T08:42:05.583484Z","iopub.status.idle":"2024-10-29T11:02:03.845114Z","shell.execute_reply.started":"2024-10-29T08:42:05.583451Z","shell.execute_reply":"2024-10-29T11:02:03.843937Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Evaluating Models Performance\n\nNow let's compare the models using the `get_models_accuracy()` function. This function takes a dictionary of models, evaluates each on the test dataset using cross-entropy loss, and returns a dictionary containing the average loss (computed across the entire dataset) and accuracy for each model.","metadata":{}},{"cell_type":"code","source":"def get_models_accuracy(\n    dataloaders: tuple[DataLoader, DataLoader],\n    models: dict[str, nn.Module],\n    device: str | torch.device\n) -> dict[str, dict[str, float]]:\n    \n    \"\"\"\n    Evaluates the accuracy and average test loss of multiple models on a test dataset.\n\n    Args:\n        dataloaders (tuple[DataLoader, DataLoader]): A tuple containing the training and testing DataLoaders. \n                                                     Only the testing DataLoader is used.\n        models (dict[str, nn.Module]): A dictionary where keys are model names (str) and values are PyTorch models (nn.Module).\n        device (str | torch.device): The device on which the models will be evaluated (e.g., 'cuda' or torch.device object).\n\n    Returns:\n        dict[str, dict[str, float]]: A dictionary where each key is a model name (str) and the corresponding value \n                                     is another dictionary with the model's average test loss ('loss') and accuracy ('accuracy').\n\n    The function computes the loss and accuracy for each model by iterating over the test dataset, and stores the results in a dictionary.\n    \"\"\"\n    \n    results: dict[str, dict[str, float]] = {}\n    \n    _, test_dataloader = dataloaders\n    \n    loss_fn = nn.CrossEntropyLoss()\n    softmax = nn.Softmax(dim=1)\n    \n    with torch.no_grad():\n        for model_name in models.keys():\n            \n            model = nn.DataParallel(models[model_name].to(device))\n            model.eval()\n        \n            avg_test_loss = 0.0\n            n_correct = 0\n            n_samples = len(test_dataloader.dataset)\n            n_batches = len(test_dataloader)\n\n            for X, y in test_dataloader:\n                X, y = X.to(device), y.to(device)\n\n                prediction = model(X)\n\n                loss = loss_fn(prediction, y)\n                avg_test_loss += loss.item()\n\n                prediction = softmax(prediction)\n                n_correct += (torch.argmax(prediction, dim=1) == y).sum().item()\n\n            results[model_name] = {\"loss\": avg_test_loss / n_batches, \"accuracy\": n_correct / n_samples * 100}\n    \n    return results\n    ","metadata":{"execution":{"iopub.status.busy":"2024-10-09T21:01:34.786098Z","iopub.execute_input":"2024-10-09T21:01:34.786861Z","iopub.status.idle":"2024-10-09T21:01:34.797856Z","shell.execute_reply.started":"2024-10-09T21:01:34.786819Z","shell.execute_reply":"2024-10-09T21:01:34.796882Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Loading the Models\n\nBefore comparing the models, we need to load their pre-trained weights from the appropriate files using PyTorch's `load_state_dict()`.","metadata":{}},{"cell_type":"code","source":"ViT_T = VisionTransformer(batch_first=True, n_classes=len(idx_to_class_train))\nViT_S = VisionTransformer(embed_dim=384, n_heads=6, mlp_size=1536, batch_first=True, n_classes=len(idx_to_class_train))\nViT_S_12 = VisionTransformer(embed_dim=384, n_heads=12, mlp_size=1536, batch_first=True, n_classes=len(idx_to_class_train))","metadata":{"execution":{"iopub.status.busy":"2024-10-13T16:52:42.260876Z","iopub.execute_input":"2024-10-13T16:52:42.261372Z","iopub.status.idle":"2024-10-13T16:52:42.541850Z","shell.execute_reply.started":"2024-10-13T16:52:42.261322Z","shell.execute_reply":"2024-10-13T16:52:42.540719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ViT_T.load_state_dict(torch.load(\"/kaggle/input/vision-transformer/pytorch/vit-t/1/ViT-T-2024-10-09.pth\", weights_only=True, map_location=torch.device(\"cpu\")))\nViT_S.load_state_dict(torch.load(\"/kaggle/input/vision-transformer/pytorch/vit-s/1/ViT-S-2024-10-09.pth\", weights_only=True, map_location=torch.device(\"cpu\")))\nViT_S_12.load_state_dict(torch.load(\"/kaggle/input/vision-transformer/pytorch/vit-s-12/1/ViT-S-12-2024-10-08.pth\", weights_only=True, map_location=torch.device(\"cpu\")))","metadata":{"execution":{"iopub.status.busy":"2024-10-13T16:56:31.693698Z","iopub.execute_input":"2024-10-13T16:56:31.694415Z","iopub.status.idle":"2024-10-13T16:56:31.960337Z","shell.execute_reply.started":"2024-10-13T16:56:31.694362Z","shell.execute_reply":"2024-10-13T16:56:31.959131Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Results\n\nWith the pre-trained weights loaded, we can now evaluate the models on the test dataset. Below are the results:","metadata":{}},{"cell_type":"code","source":"models = {\"ViT-T\": ViT_T, \"ViT-S\": ViT_S, \"ViT-S-12\": ViT_S_12}\nresults = get_models_accuracy(get_dataloaders(), models=models, device=\"cuda\")","metadata":{"execution":{"iopub.status.busy":"2024-10-09T21:01:38.027867Z","iopub.execute_input":"2024-10-09T21:01:38.028986Z","iopub.status.idle":"2024-10-09T21:09:51.327919Z","shell.execute_reply.started":"2024-10-09T21:01:38.028935Z","shell.execute_reply":"2024-10-09T21:09:51.323512Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for model_name, metrics in results.items():\n    print(f\"{model_name} - Loss: {metrics['loss']:.5f}, Accuracy: {metrics['accuracy']:.2f}%\")","metadata":{"execution":{"iopub.status.busy":"2024-10-09T21:12:06.727918Z","iopub.execute_input":"2024-10-09T21:12:06.728342Z","iopub.status.idle":"2024-10-09T21:12:06.734020Z","shell.execute_reply.started":"2024-10-09T21:12:06.728301Z","shell.execute_reply":"2024-10-09T21:12:06.732976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualizing Attention\n\nTo understand which parts of an image contribute the most to a model's decision-making, we can generate an attention map. The steps are as follows:\n1. Extract the attention weights from each multihead attention layer and compute the average across the heads.\n2. Calculate the attention rollout using the following formula:\n$$A^l_{rollout} = (A^l + I)A^{l-1}_{rollout}, \\text{for}\\,i > 1\\\\\nA^1_{rollout} = A^1 + I$$\nwhere:\n  - $A^l$ represents the average attention weights across heads for the $l$-th layer\n  - $I$ is the identity matrix, used to account for residual connections\n3. Select the first column of the final attention rollout, which estimates the self-attention weights of the CLS token for all tokens in the sequence, excluding the first row (the attention of the CLS token). This serves as the mask we'll apply to our image.\n4. Normalize the mask by dividing each element by its maximum value, then resize it to match the image dimensions.\n5. Apply the mask to the original image to generate an attention-highlighted output.","metadata":{}},{"cell_type":"code","source":"def get_attention_map(model: nn.Module, image: TensorType[torch.float32]) -> np.ndarray:\n    \"\"\"\n    Computes the attention map for a given image using the attention layers of a Vision Transformer (ViT) model.\n\n    This function calculates the attention rollout, which aggregates the attention weights across all layers \n    to generate a heatmap representing the importance of different image regions.\n\n    Args:\n        model (nn.Module): The Vision Transformer (ViT) model containing a transformer encoder.\n        image (TensorType[torch.float32]): The input image tensor of shape (C, H, W), where C is the number of channels, \n                                           and H and W are the height and width of the image, respectively.\n\n    Returns:\n        np.ndarray: An attention map overlayed on the input image, resized to match the original image dimensions \n                    and with the same number of channels as the input image.\n\n    Process:\n        1. Extracts the transformer encoder layers and embedding from the model.\n        2. Computes attention weights for each layer of the transformer using the self-attention mechanism.\n        3. Calculates the attention rollout, incorporating skip connections (identity matrix) to account \n           for the layer-wise impact of attention.\n        4. Generates an attention mask, which is resized to match the input image size.\n        5. Applies the mask to the original image to produce an attention-highlighted output.\n\n    Returns an attention-weighted image, visualizing which parts of the input image contributed most \n    to the model's decision-making.\n    \"\"\"\n    \n    image_size = image.shape[1]\n    encoder_layers = model._transformer_encoder.layers\n    embedding = model._embedding\n    attention_weights = []\n\n    for layer in encoder_layers:\n        X_embeded = embedding(X.unsqueeze(0))\n        _, attention_output_weights = layer.self_attn(X_embeded, X_embeded, X_embeded)\n        attention_weights.append(attention_output_weights.squeeze())\n        \n    A_1 = attention_weights[0]\n    I = torch.eye(len(A_1))\n    A_1_rollout = A_1 + I\n    \n    A_l_prev_rollout = A_1_rollout\n    for i in range(1, len(attention_weights)):\n        A = attention_weights[i]\n        A_l_rollout = (A + I) @ A_l_prev_rollout\n        A_l_prev_rollout = A_l_rollout\n        \n    mask_size = int(np.sqrt(A_1.shape[0]))\n    mask = A_l_rollout[0, 1:].reshape(mask_size, mask_size)\n    \n    mask = mask / mask.max()\n    \n    return (cv2.resize(mask.detach().numpy(), (image_size, image_size)) * X.numpy()).transpose(1, 2, 0)","metadata":{"execution":{"iopub.status.busy":"2024-10-13T19:00:51.450280Z","iopub.execute_input":"2024-10-13T19:00:51.450747Z","iopub.status.idle":"2024-10-13T19:00:51.463852Z","shell.execute_reply.started":"2024-10-13T19:00:51.450703Z","shell.execute_reply":"2024-10-13T19:00:51.462547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X, y = test_dataset[0]\n\nfig, (ax1, ax2, ax3, ax4) = plt.subplots(nrows=1, ncols=4, figsize=(16, 6))\n\nfig.suptitle(\"Attention map for each model\")\n\nax1.set_title(\"Original Image\")\nax1.imshow(X.permute(1, 2, 0))\n\nax2.set_title(\"ViT-T attention map\")\nax2.imshow(get_attention_map(ViT_T, X))\n\nax3.set_title(\"ViT-S attention map\")\nax3.imshow(get_attention_map(ViT_S, X))\n\nax4.set_title(\"ViT-S-12 attention map\")\nax4.imshow(get_attention_map(ViT_S_12, X))\n\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-10-13T19:32:37.764479Z","iopub.execute_input":"2024-10-13T19:32:37.765338Z","iopub.status.idle":"2024-10-13T19:32:39.114267Z","shell.execute_reply.started":"2024-10-13T19:32:37.765291Z","shell.execute_reply":"2024-10-13T19:32:39.113080Z"},"trusted":true},"execution_count":null,"outputs":[]}]}