{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.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":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport csv\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-09T03:30:25.743744Z","iopub.execute_input":"2024-07-09T03:30:25.744302Z","iopub.status.idle":"2024-07-09T03:30:26.765448Z","shell.execute_reply.started":"2024-07-09T03:30:25.744270Z","shell.execute_reply":"2024-07-09T03:30:26.764330Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Extract zip and separate cat and dog images separately in train directory","metadata":{}},{"cell_type":"code","source":"import os\nimport zipfile\nimport shutil\nimport random\nimport glob\n\n# Paths for train and test zip and working directory\ninput_dir = '/kaggle/input/dogs-vs-cats'\ntrain_zip_path = os.path.join(input_dir, 'train.zip')\ntest_zip_path = os.path.join(input_dir, 'test1.zip')\nworking_dir = '/kaggle/working'\n\n\"\"\"\n    The idea is to first extract train file in some temp folder and\n    then move cat and dog images in train and val folder in separate\n    folder of cats and dogs\n    \n    temp_train/\n    temp_test/\n    \n    train/\n        - cat/\n        - dog/\n\"\"\"\n\n# define paths for unziping\ntemp_train_dir = os.path.join(working_dir, 'temp_train')\n\n\n# create directories for them\nos.makedirs(temp_train_dir, exist_ok=True)\n\n# define cat and dog dir for train and test\ntrain_cat_dir = os.path.join(working_dir, 'train/cat')\ntrain_dog_dir = os.path.join(working_dir, 'train/dog')\n\n\n# Create directories\nos.makedirs(train_cat_dir, exist_ok=True)\nos.makedirs(train_dog_dir, exist_ok=True)\n\n# Extract train.zip\nwith zipfile.ZipFile(train_zip_path, 'r') as zip_ref:\n    zip_ref.extractall(temp_train_dir)\n\nwith zipfile.ZipFile(test_zip_path, 'r') as zip_ref:\n    zip_ref.extractall()\n\n    \n# List cat and dog files for train\ntrain_cat_files = glob.glob(os.path.join(temp_train_dir, 'train/cat.*.jpg'))\ntrain_dog_files = glob.glob(os.path.join(temp_train_dir, 'train/dog.*.jpg'))\n\n# Move files to train and val directories\nfor file in train_cat_files:\n    base_file_name = os.path.basename(file)\n    shutil.move(file, os.path.join(train_cat_dir, base_file_name))\n\nfor file in train_dog_files:\n    base_file_name = os.path.basename(file)\n    shutil.move(file, os.path.join(train_dog_dir, base_file_name))\n\nprint(\"Files have been split and moved successfully.\")\n\n# after moving remove temp_train_dir\nshutil.rmtree(temp_train_dir)\n\nprint(\"Temp train directory has been removed\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:30:26.767449Z","iopub.execute_input":"2024-07-09T03:30:26.767937Z","iopub.status.idle":"2024-07-09T03:30:41.904528Z","shell.execute_reply.started":"2024-07-09T03:30:26.767897Z","shell.execute_reply":"2024-07-09T03:30:41.903625Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Visualize a cat and dog image","metadata":{}},{"cell_type":"code","source":"import random\nfrom PIL import Image\nfrom pathlib import Path\n\nrandom.seed(42)\n\nimage_path = Path(\"/kaggle/working/train/\")\n\nimage_path_list = list(image_path.glob(\"*/*.jpg\"))\n\nrandom_image_path = random.choice(image_path_list)\n\nimage_class = random_image_path.parent.stem\n\nimg = Image.open(random_image_path)\nprint(\"Random image path: \", random_image_path)\nprint(f\"Image Height and Width: {img.height}, {img.width}\")\nprint(\"Image class: \", image_class)\nimg","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:30:41.910123Z","iopub.execute_input":"2024-07-09T03:30:41.910788Z","iopub.status.idle":"2024-07-09T03:30:42.107431Z","shell.execute_reply.started":"2024-07-09T03:30:41.910756Z","shell.execute_reply":"2024-07-09T03:30:42.106658Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualize it using matplotlib","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\n\nimg_as_array = np.asarray(img)\n\nplt.figure(figsize=(10, 7))\nplt.imshow(img_as_array)\n\nplt.title(f\"Image class: {image_class}, Image shape: {img_as_array.shape}\")\nplt.axis(False);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:30:42.108394Z","iopub.execute_input":"2024-07-09T03:30:42.108638Z","iopub.status.idle":"2024-07-09T03:30:42.519733Z","shell.execute_reply.started":"2024-07-09T03:30:42.108617Z","shell.execute_reply":"2024-07-09T03:30:42.518825Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Replicating ViT SOTA model from research paper\n[Here's](https://paperswithcode.com/paper/an-image-is-worth-16x16-words-transformers-1) the link to the paper","metadata":{}},{"cell_type":"markdown","source":"## Setup of data is done, we have separated data into cats and dogs.\n## Now we have to create datasets and dataloaders","metadata":{}},{"cell_type":"code","source":"import torch\nimport torchvision\nfrom torchvision import transforms\n\nIMG_SIZE = 224\n\ntrain_transform = transforms.Compose([\n    transforms.Resize((IMG_SIZE, IMG_SIZE)),\n    transforms.ToTensor(),\n])","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:30:42.520970Z","iopub.execute_input":"2024-07-09T03:30:42.521291Z","iopub.status.idle":"2024-07-09T03:30:47.556973Z","shell.execute_reply.started":"2024-07-09T03:30:42.521264Z","shell.execute_reply":"2024-07-09T03:30:47.555969Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set seeds\ndef set_seeds(seed: int=42):\n    \"\"\"Sets random sets for torch operations.\n\n    Args:\n        seed (int, optional): Random seed to set. Defaults to 42.\n    \"\"\"\n    # Set the seed for general torch operations\n    torch.manual_seed(seed)\n    # Set the seed for CUDA torch operations (ones that happen on the GPU)\n    torch.cuda.manual_seed(seed)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:18:16.090705Z","iopub.execute_input":"2024-07-09T04:18:16.091135Z","iopub.status.idle":"2024-07-09T04:18:16.096726Z","shell.execute_reply.started":"2024-07-09T04:18:16.091098Z","shell.execute_reply":"2024-07-09T04:18:16.095792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import datasets\nfrom torch.utils.data import DataLoader, SubsetRandomSampler\n\nBATCH_SIZE = 32\ntrain_dataset_dir = Path(\"/kaggle/working/train/\")\n\ntrain_dataset = datasets.ImageFolder(\n    root=train_dataset_dir,\n    transform=train_transform,\n    target_transform=None\n)\n\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\n\n# If you want to divide it by 20%\n# split = int(np.floor(0.2 * dataset_size))\n# train_indices, val_indices = indices[split:], indices[:split]\n\nNUM_TRAIN_SAMPLES = 500\nNUM_VAL_SAMPLES = 50\n\n# Select 500 indices for training and 50 for validation\ntrain_indices = indices[:NUM_TRAIN_SAMPLES]\nval_indices = indices[NUM_TRAIN_SAMPLES:NUM_TRAIN_SAMPLES + NUM_VAL_SAMPLES]\n\ntrain_sampler = SubsetRandomSampler(train_indices)\nval_sampler = SubsetRandomSampler(val_indices)\n\nclass_names = train_dataset.classes\ntrain_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=train_sampler,\n    num_workers=os.cpu_count()\n)\n\nval_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=val_sampler,\n    num_workers=os.cpu_count()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:34.205902Z","iopub.execute_input":"2024-07-09T03:31:34.206330Z","iopub.status.idle":"2024-07-09T03:31:34.334411Z","shell.execute_reply.started":"2024-07-09T03:31:34.206296Z","shell.execute_reply":"2024-07-09T03:31:34.333348Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Visualize a single image from dataloader","metadata":{}},{"cell_type":"code","source":"image_batch, label_batch = next(iter(train_dataloader))\n\nimage, label = image_batch[0], label_batch[0]\n\nimage.shape, label","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:39.590690Z","iopub.execute_input":"2024-07-09T03:31:39.591645Z","iopub.status.idle":"2024-07-09T03:31:40.168483Z","shell.execute_reply.started":"2024-07-09T03:31:39.591607Z","shell.execute_reply":"2024-07-09T03:31:40.167276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.imshow(image.permute(1, 2, 0))\nplt.title(class_names[label])\nplt.axis(False);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:43.148843Z","iopub.execute_input":"2024-07-09T03:31:43.149706Z","iopub.status.idle":"2024-07-09T03:31:43.457966Z","shell.execute_reply.started":"2024-07-09T03:31:43.149667Z","shell.execute_reply":"2024-07-09T03:31:43.457055Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Replicating ViT Paper","metadata":{}},{"cell_type":"markdown","source":"### 1. Convert a single image into patches","metadata":{}},{"cell_type":"code","source":"patch_size = 16\nimage_permuted = image.permute(1, 2, 0) # (H, W, C)\nplt.figure(figsize=(patch_size, patch_size))\nplt.imshow(image_permuted[:patch_size, :, :]);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:48.632775Z","iopub.execute_input":"2024-07-09T03:31:48.633643Z","iopub.status.idle":"2024-07-09T03:31:48.833334Z","shell.execute_reply.started":"2024-07-09T03:31:48.633612Z","shell.execute_reply":"2024-07-09T03:31:48.832395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Patches for a single row","metadata":{}},{"cell_type":"code","source":"img_size = 224\npatch_size = 16\nnum_patches = img_size / patch_size\n\nfig, axs = plt.subplots(\n    nrows=1,\n    ncols=img_size // patch_size,\n    figsize=(num_patches, num_patches),\n    sharex=True,\n    sharey=True\n)\n\nfor i, patch in enumerate(range(0, img_size, patch_size)):\n    axs[i].imshow(image_permuted[:patch_size, patch:patch+patch_size, :]);\n    axs[i].set_xlabel(i+1)\n    axs[i].set_xticks([])\n    axs[i].set_yticks([])","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:50.781853Z","iopub.execute_input":"2024-07-09T03:31:50.782687Z","iopub.status.idle":"2024-07-09T03:31:51.511707Z","shell.execute_reply.started":"2024-07-09T03:31:50.782646Z","shell.execute_reply":"2024-07-09T03:31:51.510806Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Pathces for the whole image","metadata":{}},{"cell_type":"code","source":"img_size = 224\npatch_size = 16\nnum_patches = img_size / patch_size\n\nfig, axs = plt.subplots(\n    nrows=img_size // patch_size,\n    ncols=img_size // patch_size,\n    figsize=(num_patches, num_patches),\n    sharex=True,\n    sharey=True\n)\n\nfor i, patch_height in enumerate(range(0, img_size, patch_size)):\n    for j, patch_width in enumerate(range(0, img_size, patch_size)):\n        axs[i, j].imshow(\n            image_permuted[patch_height: patch_height + patch_size,\n                          patch_width:patch_width+patch_size, :]\n        )\n        axs[i, j].set_ylabel(i+1,\n                             rotation=\"horizontal\",\n                             horizontalalignment=\"right\",\n                             verticalalignment=\"center\")\n        axs[i, j].set_xlabel(j+1)\n        axs[i, j].set_xticks([])\n        axs[i, j].set_yticks([])\n        axs[i, j].label_outer()","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:31:53.294542Z","iopub.execute_input":"2024-07-09T03:31:53.295265Z","iopub.status.idle":"2024-07-09T03:32:31.109531Z","shell.execute_reply.started":"2024-07-09T03:31:53.295227Z","shell.execute_reply":"2024-07-09T03:32:31.108294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create image patches with torch.nn.Conv2d()","metadata":{}},{"cell_type":"code","source":"from torch import nn\n\npatch_size = 16\n\nconv2d = nn.Conv2d(\n    in_channels=3,\n    out_channels=768,\n    kernel_size=patch_size,\n    stride=patch_size,\n    padding=0\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:31.111728Z","iopub.execute_input":"2024-07-09T03:32:31.112379Z","iopub.status.idle":"2024-07-09T03:32:31.127632Z","shell.execute_reply.started":"2024-07-09T03:32:31.112341Z","shell.execute_reply":"2024-07-09T03:32:31.126669Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# View single image\nplt.imshow(image.permute(1, 2, 0)) # adjust for matplotlib\nplt.title(class_names[label])\nplt.axis(False);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:31.129074Z","iopub.execute_input":"2024-07-09T03:32:31.129429Z","iopub.status.idle":"2024-07-09T03:32:31.379230Z","shell.execute_reply.started":"2024-07-09T03:32:31.129400Z","shell.execute_reply":"2024-07-09T03:32:31.378213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Pass the image through conv2d\nimage_conv = conv2d(image.unsqueeze(0))\nimage_conv.shape","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:47.202492Z","iopub.execute_input":"2024-07-09T03:32:47.202892Z","iopub.status.idle":"2024-07-09T03:32:47.264858Z","shell.execute_reply.started":"2024-07-09T03:32:47.202856Z","shell.execute_reply":"2024-07-09T03:32:47.263838Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot random 5 convolutional feature maps\nimport random\nrandom_indexes = random.sample(range(0, 758), k=5) # pick 5 numbers between 0 and the embedding size\nprint(f\"Showing random convolutional feature maps from indexes: {random_indexes}\")\n\n# Create plot\nfig, axs = plt.subplots(nrows=1, ncols=5, figsize=(12, 12))\n\n# Plot random image feature maps\nfor i, idx in enumerate(random_indexes):\n    image_conv_feature_map = image_conv[:, idx, :, :] # index on the output tensor of the convolutional layer\n    axs[i].imshow(image_conv_feature_map.squeeze().detach().numpy())\n    axs[i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[]);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:47.980683Z","iopub.execute_input":"2024-07-09T03:32:47.981104Z","iopub.status.idle":"2024-07-09T03:32:48.413497Z","shell.execute_reply.started":"2024-07-09T03:32:47.981073Z","shell.execute_reply":"2024-07-09T03:32:48.412507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Flatten the patch embedding layer","metadata":{}},{"cell_type":"code","source":"# Create flatten layer\nflatten = nn.Flatten(start_dim=2, # flatten feature_map_height (dimension 2)\n                     end_dim=3) # flatten feature_map_width (dimension 3)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:50.322527Z","iopub.execute_input":"2024-07-09T03:32:50.322916Z","iopub.status.idle":"2024-07-09T03:32:50.327846Z","shell.execute_reply.started":"2024-07-09T03:32:50.322885Z","shell.execute_reply":"2024-07-09T03:32:50.326773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_out_of_conv_flattened = flatten(image_conv)\nimage_out_of_conv_flattened.shape","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:51.316685Z","iopub.execute_input":"2024-07-09T03:32:51.317538Z","iopub.status.idle":"2024-07-09T03:32:51.324218Z","shell.execute_reply.started":"2024-07-09T03:32:51.317504Z","shell.execute_reply":"2024-07-09T03:32:51.323201Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> Current shape is 1, 768, 196 and we want the shape of 196, 768","metadata":{}},{"cell_type":"code","source":"# Get flattened image patch embeddings in right shape\nimage_out_of_conv_flattened_reshaped = image_out_of_conv_flattened.permute(0, 2, 1) # [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]\nprint(f\"Patch embedding sequence shape: {image_out_of_conv_flattened_reshaped.shape} -> [batch_size, num_patches, embedding_size]\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:54.027715Z","iopub.execute_input":"2024-07-09T03:32:54.028110Z","iopub.status.idle":"2024-07-09T03:32:54.034287Z","shell.execute_reply.started":"2024-07-09T03:32:54.028079Z","shell.execute_reply":"2024-07-09T03:32:54.033246Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get a single flattened feature map\nsingle_flattened_feature_map = image_out_of_conv_flattened_reshaped[:, :, 0] # index: (batch_size, number_of_patches, embedding_dimension)\n\n# Plot the flattened feature map visually\nplt.figure(figsize=(22, 22))\nplt.imshow(single_flattened_feature_map.detach().numpy())\nplt.title(f\"Flattened feature map shape: {single_flattened_feature_map.shape}\")\nplt.axis(False);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:55.057587Z","iopub.execute_input":"2024-07-09T03:32:55.058120Z","iopub.status.idle":"2024-07-09T03:32:55.237662Z","shell.execute_reply.started":"2024-07-09T03:32:55.058081Z","shell.execute_reply":"2024-07-09T03:32:55.236171Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create a PatchEmbedding class","metadata":{}},{"cell_type":"code","source":"class PatchEmbedding(nn.Module):\n    \"\"\"\n    Turns a 2D input image into a 1D sequence learnable embedding vector\n    \n    Args:\n        in_channels (int): Number of color channels for the input images. Defaults to 3\n        patch_size (int): Size of patches to convert image into.\n        embedding_dim (int): Size of embedding to turn image into. Defaults to 768\n    \"\"\"\n    def __init__(self, \n                in_channels:int=3,\n                patch_size:int=16,\n                embedding_dim:int=768\n                ):\n        super().__init__()\n        self.patcher = nn.Conv2d(in_channels=in_channels,\n                                out_channels=embedding_dim,\n                                kernel_size=patch_size,\n                                stride=patch_size,\n                                padding=0)\n        self.flatten = nn.Flatten(start_dim=2,end_dim=3)\n    \n    def forward(self, x):\n        image_resolution = x.shape[-1]\n        assert image_resolution % patch_size == 0, f\"Input image size must be divisble by patch size, image shape: {image_resolution}, patch size: {patch_size}\"\n\n        # Perform the forward pass\n        x_patched = self.patcher(x)\n        x_flattened = self.flatten(x_patched)\n        # 6. Make sure the output shape has the right order\n        return x_flattened.permute(0, 2, 1) # adjust so the embedding is on the final dimension [batch_size, P^2•C, N] -> [batch_size, N, P^2•C]","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:57.699745Z","iopub.execute_input":"2024-07-09T03:32:57.700571Z","iopub.status.idle":"2024-07-09T03:32:57.708471Z","shell.execute_reply.started":"2024-07-09T03:32:57.700541Z","shell.execute_reply":"2024-07-09T03:32:57.707522Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.manual_seed(42)\ntorch.cuda.manual_seed(42)\n\n\npatchify = PatchEmbedding(in_channels=3,\n                         patch_size=16,\n                         embedding_dim=768)\n\n# Pass a single image through\nprint(f\"Input image shape: {image.unsqueeze(0).shape}\")\npatch_embedded_image = patchify(image.unsqueeze(0)) # add an extra batch dimension on the 0th index, otherwise will error\nprint(f\"Output patch embedding shape: {patch_embedded_image.shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:58.409731Z","iopub.execute_input":"2024-07-09T03:32:58.410607Z","iopub.status.idle":"2024-07-09T03:32:58.426621Z","shell.execute_reply.started":"2024-07-09T03:32:58.410575Z","shell.execute_reply":"2024-07-09T03:32:58.425636Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating the class token embedding","metadata":{}},{"cell_type":"code","source":"# Get the batch size and embedding dimension\nbatch_size = patch_embedded_image.shape[0]\nembedding_dimension = patch_embedded_image.shape[-1]\n\n# Create the class token embedding as a learnable parameter that shares the same size as the embedding dimension (D)\nclass_token = nn.Parameter(torch.ones(batch_size, 1, embedding_dimension), # [batch_size, number_of_tokens, embedding_dimension]\n                           requires_grad=True) # make sure the embedding is learnable\n\n# Show the first 10 examples of the class_token\nprint(class_token[:, :, :10])\n\n# Print the class_token shape\nprint(f\"Class token shape: {class_token.shape} -> [batch_size, number_of_tokens, embedding_dimension]\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:32:59.748677Z","iopub.execute_input":"2024-07-09T03:32:59.749540Z","iopub.status.idle":"2024-07-09T03:32:59.778426Z","shell.execute_reply.started":"2024-07-09T03:32:59.749499Z","shell.execute_reply":"2024-07-09T03:32:59.777475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Add the class token embedding to the front of the patch embedding\npatch_embedded_image_with_class_embedding = torch.cat((class_token, patch_embedded_image),\n                                                      dim=1) # concat on first dimension\n\n# Print the sequence of patch embeddings with the prepended class token embedding\nprint(patch_embedded_image_with_class_embedding)\nprint(f\"Sequence of patch embeddings with class token prepended shape: {patch_embedded_image_with_class_embedding.shape} -> [batch_size, number_of_patches, embedding_dimension]\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:00.557948Z","iopub.execute_input":"2024-07-09T03:33:00.558812Z","iopub.status.idle":"2024-07-09T03:33:00.568479Z","shell.execute_reply.started":"2024-07-09T03:33:00.558779Z","shell.execute_reply":"2024-07-09T03:33:00.567485Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating the position embedding","metadata":{}},{"cell_type":"code","source":"height = 224\nwidth = 224\npatch_size = 16\n# Calculate N (number of patches)\nnumber_of_patches = int((height * width) / patch_size**2)\n\n# Get embedding dimension\nembedding_dimension = patch_embedded_image_with_class_embedding.shape[2]\n\n# Create the learnable 1D position embedding\nposition_embedding = nn.Parameter(torch.ones(1,\n                                             number_of_patches+1,\n                                             embedding_dimension),\n                                  requires_grad=True) # make sure it's learnable\n\n# Show the first 10 sequences and 10 position embedding values and check the shape of the position embedding\nprint(position_embedding[:, :10, :10])\nprint(f\"Position embeddding shape: {position_embedding.shape} -> [batch_size, number_of_patches, embedding_dimension]\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:02.843613Z","iopub.execute_input":"2024-07-09T03:33:02.843984Z","iopub.status.idle":"2024-07-09T03:33:02.853745Z","shell.execute_reply.started":"2024-07-09T03:33:02.843955Z","shell.execute_reply":"2024-07-09T03:33:02.852670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Add the position embedding to the patch and class token embedding\npatch_and_position_embedding = patch_embedded_image_with_class_embedding + position_embedding\nprint(patch_and_position_embedding)\nprint(f\"Patch embeddings, class token prepended and positional embeddings added shape: {patch_and_position_embedding.shape} -> [batch_size, number_of_patches, embedding_dimension]\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:03.707663Z","iopub.execute_input":"2024-07-09T03:33:03.708332Z","iopub.status.idle":"2024-07-09T03:33:03.717108Z","shell.execute_reply.started":"2024-07-09T03:33:03.708300Z","shell.execute_reply":"2024-07-09T03:33:03.716119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Image to embedding with position","metadata":{}},{"cell_type":"code","source":"torch.manual_seed(42)\ntorch.cuda.manual_seed(42)\n\n# 1. Set patch size\npatch_size = 16\n\n# 2. Print shape of original image tensor and get the image dimensions\nprint(f\"Image tensor shape: {image.shape}\")\nheight, width = image.shape[1], image.shape[2]\n\n# 3. Get image tensor and add batch dimension\nx = image.unsqueeze(0)\nprint(f\"Input image with batch dimension shape: {x.shape}\")\n\n# 4. Create patch embedding layer\npatch_embedding_layer = PatchEmbedding(in_channels=3,\n                                       patch_size=patch_size,\n                                       embedding_dim=768)\n\n# 5. Pass image through patch embedding layer\npatch_embedding = patch_embedding_layer(x)\nprint(f\"Patching embedding shape: {patch_embedding.shape}\")\n\n# 6. Create class token embedding\nbatch_size = patch_embedding.shape[0]\nembedding_dimension = patch_embedding.shape[-1]\nclass_token = nn.Parameter(torch.ones(batch_size, 1, embedding_dimension),\n                           requires_grad=True) # make sure it's learnable\nprint(f\"Class token embedding shape: {class_token.shape}\")\n\n# 7. Prepend class token embedding to patch embedding\npatch_embedding_class_token = torch.cat((class_token, patch_embedding), dim=1)\nprint(f\"Patch embedding with class token shape: {patch_embedding_class_token.shape}\")\n\n# 8. Create position embedding\nnumber_of_patches = int((height * width) / patch_size**2)\nposition_embedding = nn.Parameter(torch.ones(1, number_of_patches+1, embedding_dimension),\n                                  requires_grad=True) # make sure it's learnable\n\n# 9. Add position embedding to patch embedding with class token\npatch_and_position_embedding = patch_embedding_class_token + position_embedding\nprint(f\"Patch and position embedding shape: {patch_and_position_embedding.shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:06.411311Z","iopub.execute_input":"2024-07-09T03:33:06.412060Z","iopub.status.idle":"2024-07-09T03:33:06.433794Z","shell.execute_reply.started":"2024-07-09T03:33:06.412030Z","shell.execute_reply":"2024-07-09T03:33:06.432902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Multi-Head Attention","metadata":{}},{"cell_type":"code","source":"# 1. Create a class that inherits from nn.Module\nclass MultiheadSelfAttentionBlock(nn.Module):\n    \"\"\"Creates a multi-head self-attention block (\"MSA block\" for short).\n    \"\"\"\n    # 2. Initialize the class with hyperparameters from Table 1\n    def __init__(self,\n                 embedding_dim:int=768, # Hidden size D from Table 1 for ViT-Base\n                 num_heads:int=12, # Heads from Table 1 for ViT-Base\n                 attn_dropout:float=0): # doesn't look like the paper uses any dropout in MSABlocks\n        super().__init__()\n\n        # 3. Create the Norm layer (LN)\n        self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)\n\n        # 4. Create the Multi-Head Attention (MSA) layer\n        self.multihead_attn = nn.MultiheadAttention(embed_dim=embedding_dim,\n                                                    num_heads=num_heads,\n                                                    dropout=attn_dropout,\n                                                    batch_first=True) # does our batch dimension come first?\n\n    # 5. Create a forward() method to pass the data throguh the layers\n    def forward(self, x):\n        x = self.layer_norm(x)\n        attn_output, _ = self.multihead_attn(query=x, # query embeddings\n                                             key=x, # key embeddings\n                                             value=x, # value embeddings\n                                             need_weights=False) # do we need the weights or just the layer outputs?\n        return attn_output","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:09.264019Z","iopub.execute_input":"2024-07-09T03:33:09.264415Z","iopub.status.idle":"2024-07-09T03:33:09.272307Z","shell.execute_reply.started":"2024-07-09T03:33:09.264384Z","shell.execute_reply":"2024-07-09T03:33:09.271333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of MSABlock\nmultihead_self_attention_block = MultiheadSelfAttentionBlock(embedding_dim=768, # from Table 1\n                                                             num_heads=12) # from Table 1\n\n# Pass patch and position image embedding through MSABlock\npatched_image_through_msa_block = multihead_self_attention_block(patch_and_position_embedding)\nprint(f\"Input shape of MSA block: {patch_and_position_embedding.shape}\")\nprint(f\"Output shape MSA block: {patched_image_through_msa_block.shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:10.117279Z","iopub.execute_input":"2024-07-09T03:33:10.117648Z","iopub.status.idle":"2024-07-09T03:33:10.217145Z","shell.execute_reply.started":"2024-07-09T03:33:10.117620Z","shell.execute_reply":"2024-07-09T03:33:10.216069Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Multi layer perceptron","metadata":{}},{"cell_type":"code","source":"# 1. Create a class that inherits from nn.Module\nclass MLPBlock(nn.Module):\n    \"\"\"Creates a layer normalized multilayer perceptron block (\"MLP block\" for short).\"\"\"\n    # 2. Initialize the class with hyperparameters from Table 1 and Table 3\n    def __init__(self,\n                 embedding_dim:int=768, # Hidden Size D from Table 1 for ViT-Base\n                 mlp_size:int=3072, # MLP size from Table 1 for ViT-Base\n                 dropout:float=0.1): # Dropout from Table 3 for ViT-Base\n        super().__init__()\n\n        # 3. Create the Norm layer (LN)\n        self.layer_norm = nn.LayerNorm(normalized_shape=embedding_dim)\n\n        # 4. Create the Multilayer perceptron (MLP) layer(s)\n        self.mlp = nn.Sequential(\n            nn.Linear(in_features=embedding_dim,\n                      out_features=mlp_size),\n            nn.GELU(), # \"The MLP contains two layers with a GELU non-linearity (section 3.1).\"\n            nn.Dropout(p=dropout),\n            nn.Linear(in_features=mlp_size, # needs to take same in_features as out_features of layer above\n                      out_features=embedding_dim), # take back to embedding_dim\n            nn.Dropout(p=dropout) # \"Dropout, when used, is applied after every dense layer..\"\n        )\n\n    # 5. Create a forward() method to pass the data throguh the layers\n    def forward(self, x):\n        x = self.layer_norm(x)\n        x = self.mlp(x)\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:12.075267Z","iopub.execute_input":"2024-07-09T03:33:12.075909Z","iopub.status.idle":"2024-07-09T03:33:12.087069Z","shell.execute_reply.started":"2024-07-09T03:33:12.075879Z","shell.execute_reply":"2024-07-09T03:33:12.086165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of MLPBlock\nmlp_block = MLPBlock(embedding_dim=768, # from Table 1\n                     mlp_size=3072, # from Table 1\n                     dropout=0.1) # from Table 3\n\n# Pass output of MSABlock through MLPBlock\npatched_image_through_mlp_block = mlp_block(patched_image_through_msa_block)\nprint(f\"Input shape of MLP block: {patched_image_through_msa_block.shape}\")\nprint(f\"Output shape MLP block: {patched_image_through_mlp_block.shape}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:13.305630Z","iopub.execute_input":"2024-07-09T03:33:13.306012Z","iopub.status.idle":"2024-07-09T03:33:13.410082Z","shell.execute_reply.started":"2024-07-09T03:33:13.305982Z","shell.execute_reply":"2024-07-09T03:33:13.409092Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Transformer Encoder","metadata":{}},{"cell_type":"code","source":"# 1. Create a class that inherits from nn.Module\nclass TransformerEncoderBlock(nn.Module):\n    \"\"\"Creates a Transformer Encoder block.\"\"\"\n    # 2. Initialize the class with hyperparameters from Table 1 and Table 3\n    def __init__(self,\n                 embedding_dim:int=768, # Hidden size D from Table 1 for ViT-Base\n                 num_heads:int=12, # Heads from Table 1 for ViT-Base\n                 mlp_size:int=3072, # MLP size from Table 1 for ViT-Base\n                 mlp_dropout:float=0.1, # Amount of dropout for dense layers from Table 3 for ViT-Base\n                 attn_dropout:float=0): # Amount of dropout for attention layers\n        super().__init__()\n\n        # 3. Create MSA block (equation 2)\n        self.msa_block = MultiheadSelfAttentionBlock(embedding_dim=embedding_dim,\n                                                     num_heads=num_heads,\n                                                     attn_dropout=attn_dropout)\n\n        # 4. Create MLP block (equation 3)\n        self.mlp_block =  MLPBlock(embedding_dim=embedding_dim,\n                                   mlp_size=mlp_size,\n                                   dropout=mlp_dropout)\n\n    # 5. Create a forward() method\n    def forward(self, x):\n\n        # 6. Create residual connection for MSA block (add the input to the output)\n        x =  self.msa_block(x) + x\n\n        # 7. Create residual connection for MLP block (add the input to the output)\n        x = self.mlp_block(x) + x\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:15.391812Z","iopub.execute_input":"2024-07-09T03:33:15.392187Z","iopub.status.idle":"2024-07-09T03:33:15.400035Z","shell.execute_reply.started":"2024-07-09T03:33:15.392157Z","shell.execute_reply":"2024-07-09T03:33:15.398997Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create an instance of TransformerEncoderBlock\ntransformer_encoder_block = TransformerEncoderBlock()","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:17.327872Z","iopub.execute_input":"2024-07-09T03:33:17.328649Z","iopub.status.idle":"2024-07-09T03:33:17.402672Z","shell.execute_reply.started":"2024-07-09T03:33:17.328606Z","shell.execute_reply":"2024-07-09T03:33:17.401886Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ViT Architecture","metadata":{}},{"cell_type":"code","source":"# 1. Create a ViT class that inherits from nn.Module\nclass ViT(nn.Module):\n    \"\"\"Creates a Vision Transformer architecture with ViT-Base hyperparameters by default.\"\"\"\n    # 2. Initialize the class with hyperparameters from Table 1 and Table 3\n    def __init__(self,\n                 img_size:int=224, # Training resolution from Table 3 in ViT paper\n                 in_channels:int=3, # Number of channels in input image\n                 patch_size:int=16, # Patch size\n                 num_transformer_layers:int=12, # Layers from Table 1 for ViT-Base\n                 embedding_dim:int=768, # Hidden size D from Table 1 for ViT-Base\n                 mlp_size:int=3072, # MLP size from Table 1 for ViT-Base\n                 num_heads:int=12, # Heads from Table 1 for ViT-Base\n                 attn_dropout:float=0, # Dropout for attention projection\n                 mlp_dropout:float=0.1, # Dropout for dense/MLP layers\n                 embedding_dropout:float=0.1, # Dropout for patch and position embeddings\n                 num_classes:int=1000): # Default for ImageNet but can customize this\n        super().__init__() # don't forget the super().__init__()!\n\n        # 3. Make the image size is divisble by the patch size\n        assert img_size % patch_size == 0, f\"Image size must be divisible by patch size, image size: {img_size}, patch size: {patch_size}.\"\n\n        # 4. Calculate number of patches (height * width/patch^2)\n        self.num_patches = (img_size * img_size) // patch_size**2\n\n        # 5. Create learnable class embedding (needs to go at front of sequence of patch embeddings)\n        self.class_embedding = nn.Parameter(data=torch.randn(1, 1, embedding_dim),\n                                            requires_grad=True)\n\n        # 6. Create learnable position embedding\n        self.position_embedding = nn.Parameter(data=torch.randn(1, self.num_patches+1, embedding_dim),\n                                               requires_grad=True)\n\n        # 7. Create embedding dropout value\n        self.embedding_dropout = nn.Dropout(p=embedding_dropout)\n\n        # 8. Create patch embedding layer\n        self.patch_embedding = PatchEmbedding(in_channels=in_channels,\n                                              patch_size=patch_size,\n                                              embedding_dim=embedding_dim)\n\n        # 9. Create Transformer Encoder blocks (we can stack Transformer Encoder blocks using nn.Sequential())\n        # Note: The \"*\" means \"all\"\n        self.transformer_encoder = nn.Sequential(*[TransformerEncoderBlock(embedding_dim=embedding_dim,\n                                                                            num_heads=num_heads,\n                                                                            mlp_size=mlp_size,\n                                                                            mlp_dropout=mlp_dropout) for _ in range(num_transformer_layers)])\n\n        # 10. Create classifier head\n        self.classifier = nn.Sequential(\n            nn.LayerNorm(normalized_shape=embedding_dim),\n            nn.Linear(in_features=embedding_dim,\n                      out_features=num_classes)\n        )\n\n    # 11. Create a forward() method\n    def forward(self, x):\n\n        # 12. Get batch size\n        batch_size = x.shape[0]\n\n        # 13. Create class token embedding and expand it to match the batch size (equation 1)\n        class_token = self.class_embedding.expand(batch_size, -1, -1) # \"-1\" means to infer the dimension (try this line on its own)\n\n        # 14. Create patch embedding (equation 1)\n        x = self.patch_embedding(x)\n\n        # 15. Concat class embedding and patch embedding (equation 1)\n        x = torch.cat((class_token, x), dim=1)\n\n        # 16. Add position embedding to patch embedding (equation 1)\n        x = self.position_embedding + x\n\n        # 17. Run embedding dropout (Appendix B.1)\n        x = self.embedding_dropout(x)\n\n        # 18. Pass patch, position and class embedding through transformer encoder layers (equations 2 & 3)\n        x = self.transformer_encoder(x)\n\n        # 19. Put 0 index logit through classifier (equation 4)\n        x = self.classifier(x[:, 0]) # run on each sample in a batch at 0 index\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:18.972437Z","iopub.execute_input":"2024-07-09T03:33:18.973168Z","iopub.status.idle":"2024-07-09T03:33:18.987072Z","shell.execute_reply.started":"2024-07-09T03:33:18.973137Z","shell.execute_reply":"2024-07-09T03:33:18.985973Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.manual_seed(42)\ntorch.cuda.manual_seed(42)\n\n# Create a random tensor with same shape as a single image\nrandom_image_tensor = torch.randn(1, 3, 224, 224) # (batch_size, color_channels, height, width)\n\n# Create an instance of ViT with the number of classes we're working with (pizza, steak, sushi)\nvit = ViT(num_classes=len(class_names))\n\n# Pass the random image tensor to our ViT instance\nvit(random_image_tensor)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:20.029899Z","iopub.execute_input":"2024-07-09T03:33:20.030251Z","iopub.status.idle":"2024-07-09T03:33:21.157755Z","shell.execute_reply.started":"2024-07-09T03:33:20.030223Z","shell.execute_reply":"2024-07-09T03:33:21.156823Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Summary of the model","metadata":{}},{"cell_type":"code","source":"from torchinfo import summary\n\n# Print a summary of our custom ViT model using torchinfo (uncomment for actual output)\nsummary(model=vit,\n        input_size=(32, 3, 224, 224), # (batch_size, color_channels, height, width)\n        # col_names=[\"input_size\"], # uncomment for smaller output\n        col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n        col_width=20,\n        row_settings=[\"var_names\"]\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:33:25.363696Z","iopub.execute_input":"2024-07-09T03:33:25.364393Z","iopub.status.idle":"2024-07-09T03:33:26.659586Z","shell.execute_reply.started":"2024-07-09T03:33:25.364359Z","shell.execute_reply":"2024-07-09T03:33:26.658497Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Optimizer and loss function","metadata":{}},{"cell_type":"code","source":"# Setup the optimizer to optimize our ViT model parameters using hyperparameters from the ViT paper\noptimizer = torch.optim.Adam(params=vit.parameters(),\n                             lr=3e-3, # Base LR from Table 3 for ViT-* ImageNet-1k\n                             betas=(0.9, 0.999), # default values but also mentioned in ViT paper section 4.1 (Training & Fine-tuning)\n                             weight_decay=0.3)\nloss_fn = nn.CrossEntropyLoss()","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:41:19.155893Z","iopub.execute_input":"2024-07-09T03:41:19.156525Z","iopub.status.idle":"2024-07-09T03:41:19.162792Z","shell.execute_reply.started":"2024-07-09T03:41:19.156485Z","shell.execute_reply":"2024-07-09T03:41:19.161884Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Creating train and test loop function","metadata":{}},{"cell_type":"code","source":"# train step\ndef train_step(model:torch.nn.Module,\n              dataloader:torch.utils.data.DataLoader,\n              loss_fn:torch.nn.Module,\n              optimizer:torch.optim.Optimizer,\n              device:str=\"gpu\"):\n    # Put model in train mode\n    model.train()\n    \n    # setup train loss and train accuracy values\n    train_loss, train_acc = 0, 0\n    \n    # Loop through data loader data batches\n    for batch, (X, y) in enumerate(dataloader):\n        # Sending data to target device\n        X, y = X.to(device), y.to(device)\n        \n        # 1. Forward pass\n        y_pred = model(X)\n        \n        # 2.Calculate loss and accumulate loss\n        loss = loss_fn(y_pred, y)\n        train_loss += loss.item()\n        \n        # 3. Optimizer zero grad\n        optimizer.zero_grad()\n        \n        # 4. Loss backward\n        loss.backward()\n        \n        # 5. Optimizer step step\n        optimizer.step()\n        \n        # 6. Calculate and accumulate accuracy metric across all batches\n        y_pred_class = torch.argmax(torch.softmax(y_pred, dim=1), dim=1)\n        train_acc += (y_pred_class == y).sum().item() / len(y_pred)\n    \n    # Adjust metrics to get average loss and accuracy per batch\n    train_loss = train_loss / len(dataloader)\n    train_acc = train_acc / len(dataloader)\n    return train_loss, train_acc","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:36:49.182125Z","iopub.execute_input":"2024-07-09T03:36:49.182863Z","iopub.status.idle":"2024-07-09T03:36:49.192467Z","shell.execute_reply.started":"2024-07-09T03:36:49.182815Z","shell.execute_reply":"2024-07-09T03:36:49.191401Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# test step\ndef test_step(model:torch.nn.Module,\n             dataloader: torch.utils.data.DataLoader,\n             loss_fn:torch.nn.Module,\n             device:str):\n    model.eval()\n    \n    test_loss, test_acc = 0, 0\n    with torch.inference_mode():\n        for batch, (X, y) in enumerate(dataloader):\n            # Send data to target device\n            X, y = X.to(device), y.to(device)\n            test_pred_logits = model(X)\n            loss = loss_fn(test_pred_logits, y)\n            test_loss += loss.item()\n\n            test_pred_labels = torch.argmax(torch.softmax(test_pred_logits, dim=1), dim=1)\n            test_acc += ((test_pred_labels == y).sum().item() / len(test_pred_labels))\n    test_loss = test_loss / len(dataloader)\n    test_acc = test_acc / len(dataloader)\n    return test_loss, test_acc","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:37:30.353103Z","iopub.execute_input":"2024-07-09T03:37:30.353732Z","iopub.status.idle":"2024-07-09T03:37:30.361317Z","shell.execute_reply.started":"2024-07-09T03:37:30.353700Z","shell.execute_reply":"2024-07-09T03:37:30.360416Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Combining train and test step","metadata":{}},{"cell_type":"code","source":"!pip install tqdm","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:38:39.654669Z","iopub.execute_input":"2024-07-09T03:38:39.655779Z","iopub.status.idle":"2024-07-09T03:38:53.117616Z","shell.execute_reply.started":"2024-07-09T03:38:39.655734Z","shell.execute_reply":"2024-07-09T03:38:53.116316Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tqdm import tqdm\n\ndef train(\n    model:torch.nn.Module,\n    train_dataloader: torch.utils.data.DataLoader,\n    test_dataloader: torch.utils.data.DataLoader,\n    optimizer: torch.optim.Optimizer,\n    loss_fn: torch.nn.Module=torch.nn.CrossEntropyLoss(),\n    epochs: int=100,\n    device:str=\"cpu\"\n):\n    results = {\n        \"train_loss\": [],\n        \"train_acc\": [],\n        \"test_loss\": [],\n        \"test_acc\": []\n    }\n    for epoch in tqdm(range(epochs)):\n        train_loss, train_acc = train_step(\n            model=model,\n            dataloader=train_dataloader,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            device=device\n        )\n        \n        test_loss, test_acc = test_step(\n            model=model,\n            dataloader=test_dataloader,\n            loss_fn=loss_fn,\n            device=device\n        )\n        \n        print(\n            f\"Epoch: {epoch+1} | \"\n            f\"train_loss: {train_loss:.4f} | \"\n            f\"train_acc: {train_acc:.4f} | \"\n            f\"test_loss: {test_loss:.4f} | \"\n            f\"test_acc: {test_acc:.4f}\"\n        )\n        \n        results[\"train_loss\"].append(train_loss)\n        results[\"train_acc\"].append(train_acc)\n        results[\"test_loss\"].append(test_loss)\n        results[\"test_acc\"].append(test_acc)\n        \n    return results","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:38:15.961478Z","iopub.execute_input":"2024-07-09T03:38:15.961878Z","iopub.status.idle":"2024-07-09T03:38:15.972411Z","shell.execute_reply.started":"2024-07-09T03:38:15.961845Z","shell.execute_reply":"2024-07-09T03:38:15.971233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_dataloader), len(val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:57:56.356504Z","iopub.execute_input":"2024-07-09T03:57:56.357219Z","iopub.status.idle":"2024-07-09T03:57:56.363134Z","shell.execute_reply.started":"2024-07-09T03:57:56.357187Z","shell.execute_reply":"2024-07-09T03:57:56.362237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"32*625, 32 * 157","metadata":{"execution":{"iopub.status.busy":"2024-07-09T03:58:03.878240Z","iopub.execute_input":"2024-07-09T03:58:03.878960Z","iopub.status.idle":"2024-07-09T03:58:03.884705Z","shell.execute_reply.started":"2024-07-09T03:58:03.878927Z","shell.execute_reply":"2024-07-09T03:58:03.883754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Take only some percent of data","metadata":{}},{"cell_type":"code","source":"from torchvision import datasets\nfrom torch.utils.data import DataLoader, SubsetRandomSampler\n\nBATCH_SIZE = 32\ntrain_dataset_dir = Path(\"/kaggle/working/train/\")\n\ntrain_dataset = datasets.ImageFolder(\n    root=train_dataset_dir,\n    transform=train_transform,\n    target_transform=None\n)\n\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\n\n# If you want to divide it by 20%\n# split = int(np.floor(0.2 * dataset_size))\n# train_indices, val_indices = indices[split:], indices[:split]\n\nNUM_TRAIN_SAMPLES = 500\nNUM_VAL_SAMPLES = 50\n\n# Select 500 indices for training and 50 for validation\ntrain_indices = indices[:NUM_TRAIN_SAMPLES]\nval_indices = indices[NUM_TRAIN_SAMPLES:NUM_TRAIN_SAMPLES + NUM_VAL_SAMPLES]\n\ntrain_sampler = SubsetRandomSampler(train_indices)\nval_sampler = SubsetRandomSampler(val_indices)\n\nclass_names = train_dataset.classes\ntrain_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=train_sampler,\n    num_workers=os.cpu_count()\n)\n\nval_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=val_sampler,\n    num_workers=os.cpu_count()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:06:00.932852Z","iopub.execute_input":"2024-07-09T04:06:00.933211Z","iopub.status.idle":"2024-07-09T04:06:01.057200Z","shell.execute_reply.started":"2024-07-09T04:06:00.933184Z","shell.execute_reply":"2024-07-09T04:06:01.056278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"torch.manual_seed(42)\ntorch.cuda.manual_seed(42)\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n# Train the model and save the training results to a dictionary\nmodel_0_results = train(model=vit.to(device),\n                       train_dataloader=train_dataloader,\n                       test_dataloader=val_dataloader,\n                       optimizer=optimizer,\n                       loss_fn=loss_fn,\n                       epochs=10,\n                       device=device)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:06:03.286043Z","iopub.execute_input":"2024-07-09T04:06:03.286753Z","iopub.status.idle":"2024-07-09T04:09:08.170078Z","shell.execute_reply.started":"2024-07-09T04:06:03.286724Z","shell.execute_reply":"2024-07-09T04:09:08.168951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> With just 500 training images, we are getting around 60% accuracy with our **Vision Transformer from Scratch** model.\n>\n> Let's do it with transfer learning","metadata":{}},{"cell_type":"markdown","source":"## ViT model using Transfer Learning","metadata":{}},{"cell_type":"code","source":"pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT\n\npretrained_vit = torchvision.models.vit_b_16(weights=pretrained_vit_weights)\n\n# Freeze the base parameters\nfor parameter in pretrained_vit.parameters():\n    parameter.requires_grad = False\n\nset_seeds()\npretrained_vit.heads = nn.Linear(in_features=768, out_features=len(class_names)).to(device)\n    ","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:20:54.459130Z","iopub.execute_input":"2024-07-09T04:20:54.459814Z","iopub.status.idle":"2024-07-09T04:21:00.423254Z","shell.execute_reply.started":"2024-07-09T04:20:54.459781Z","shell.execute_reply":"2024-07-09T04:21:00.422445Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# # Print a summary using torchinfo (uncomment for actual output)\n# summary(model=pretrained_vit,\n#         input_size=(32, 3, 224, 224), # (batch_size, color_channels, height, width)\n#         # col_names=[\"input_size\"], # uncomment for smaller output\n#         col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n#         col_width=20,\n#         row_settings=[\"var_names\"]\n# )","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:22:04.172967Z","iopub.execute_input":"2024-07-09T04:22:04.173743Z","iopub.status.idle":"2024-07-09T04:22:04.177915Z","shell.execute_reply.started":"2024-07-09T04:22:04.173713Z","shell.execute_reply":"2024-07-09T04:22:04.176904Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get automatic transforms from pretrained ViT weights\npretrained_vit_transforms = pretrained_vit_weights.transforms()\nprint(pretrained_vit_transforms)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:22:23.749868Z","iopub.execute_input":"2024-07-09T04:22:23.750683Z","iopub.status.idle":"2024-07-09T04:22:23.755419Z","shell.execute_reply.started":"2024-07-09T04:22:23.750656Z","shell.execute_reply":"2024-07-09T04:22:23.754536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision import datasets\nfrom torch.utils.data import DataLoader, SubsetRandomSampler\n\nBATCH_SIZE = 32\ntrain_dataset_dir = Path(\"/kaggle/working/train/\")\n\ntrain_dataset = datasets.ImageFolder(\n    root=train_dataset_dir,\n    transform=pretrained_vit_transforms,\n    target_transform=None\n)\n\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\n\n# If you want to divide it by 20%\n# split = int(np.floor(0.2 * dataset_size))\n# train_indices, val_indices = indices[split:], indices[:split]\n\n# NUM_TRAIN_SAMPLES = int(np.floor(0.2*dataset_size))\nNUM_TRAIN_SAMPLES = 500\n\nNUM_VAL_SAMPLES = 50\n\n# Select 500 indices for training and 50 for validation\ntrain_indices = indices[:NUM_TRAIN_SAMPLES]\nval_indices = indices[NUM_TRAIN_SAMPLES:NUM_TRAIN_SAMPLES + NUM_VAL_SAMPLES]\n\ntrain_sampler = SubsetRandomSampler(train_indices)\nval_sampler = SubsetRandomSampler(val_indices)\n\nclass_names = train_dataset.classes\ntrain_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=train_sampler,\n    num_workers=os.cpu_count()\n)\n\nval_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=val_sampler,\n    num_workers=os.cpu_count()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:51:30.931930Z","iopub.execute_input":"2024-07-09T04:51:30.932549Z","iopub.status.idle":"2024-07-09T04:51:31.054184Z","shell.execute_reply.started":"2024-07-09T04:51:30.932510Z","shell.execute_reply":"2024-07-09T04:51:31.053471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"set_seeds()\n\n# its important to provide pretrained_vit.parameters here\noptimizer = torch.optim.Adam(params=pretrained_vit.parameters(),\n                             lr=1e-3)\nloss_fn = torch.nn.CrossEntropyLoss()\n\nmodel_1_results = train(model=pretrained_vit.to(device),\n                       train_dataloader=train_dataloader,\n                       test_dataloader=val_dataloader,\n                       optimizer=optimizer,\n                       loss_fn=loss_fn,\n                       epochs=10,\n                       device=device)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:51:35.753990Z","iopub.execute_input":"2024-07-09T04:51:35.754768Z","iopub.status.idle":"2024-07-09T04:52:57.267090Z","shell.execute_reply.started":"2024-07-09T04:51:35.754737Z","shell.execute_reply":"2024-07-09T04:52:57.265927Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"> With **Transfer Learning** we are getting accuracy of 97% and it trains faster than our model","metadata":{}},{"cell_type":"markdown","source":"## How to improve accuracy?\n- We can increase the dataset\n- We can train it for longer by increasing the epochs","metadata":{}},{"cell_type":"markdown","source":"## Increase the data to 20%","metadata":{}},{"cell_type":"code","source":"from torchvision import datasets\nfrom torch.utils.data import DataLoader, SubsetRandomSampler\n\nBATCH_SIZE = 32\ntrain_dataset_dir = Path(\"/kaggle/working/train/\")\n\ntrain_dataset = datasets.ImageFolder(\n    root=train_dataset_dir,\n    transform=pretrained_vit_transforms,\n    target_transform=None\n)\n\ndataset_size = len(train_dataset)\nindices = list(range(dataset_size))\nnp.random.shuffle(indices)\n\n# If you want to divide it by 20%\n# split = int(np.floor(0.2 * dataset_size))\n# train_indices, val_indices = indices[split:], indices[:split]\n\nNUM_TRAIN_SAMPLES = int(np.floor(0.2*dataset_size))\n\nNUM_VAL_SAMPLES = 50\n\n# Select 500 indices for training and 50 for validation\ntrain_indices = indices[:NUM_TRAIN_SAMPLES]\nval_indices = indices[NUM_TRAIN_SAMPLES:NUM_TRAIN_SAMPLES + NUM_VAL_SAMPLES]\n\ntrain_sampler = SubsetRandomSampler(train_indices)\nval_sampler = SubsetRandomSampler(val_indices)\n\nclass_names = train_dataset.classes\ntrain_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=train_sampler,\n    num_workers=os.cpu_count()\n)\n\nval_dataloader = DataLoader(\n    dataset=train_dataset,\n    batch_size=BATCH_SIZE,\n    sampler=val_sampler,\n    num_workers=os.cpu_count()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:53:07.598419Z","iopub.execute_input":"2024-07-09T04:53:07.598798Z","iopub.status.idle":"2024-07-09T04:53:07.724889Z","shell.execute_reply.started":"2024-07-09T04:53:07.598764Z","shell.execute_reply":"2024-07-09T04:53:07.724118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_indices), len(val_indices)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:53:10.071194Z","iopub.execute_input":"2024-07-09T04:53:10.071951Z","iopub.status.idle":"2024-07-09T04:53:10.078180Z","shell.execute_reply.started":"2024-07-09T04:53:10.071920Z","shell.execute_reply":"2024-07-09T04:53:10.077168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model 2 20%","metadata":{}},{"cell_type":"code","source":"set_seeds()\noptimizer = torch.optim.Adam(params=pretrained_vit.parameters(),\n                             lr=1e-3)\nloss_fn = torch.nn.CrossEntropyLoss()\nmodel_2_results = train(model=pretrained_vit.to(device),\n                       train_dataloader=train_dataloader,\n                       test_dataloader=val_dataloader,\n                       optimizer=optimizer,\n                       loss_fn=loss_fn,\n                       epochs=10,\n                       device=device)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T04:53:14.691303Z","iopub.execute_input":"2024-07-09T04:53:14.692017Z","iopub.status.idle":"2024-07-09T05:04:36.303418Z","shell.execute_reply.started":"2024-07-09T04:53:14.691984Z","shell.execute_reply":"2024-07-09T05:04:36.302265Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Plot loss curves","metadata":{}},{"cell_type":"code","source":"from typing import Tuple, Dict, List\ndef plot_loss_curves(results: Dict[str, List[float]]):\n    \"\"\"Plots training curves of a results dictionary.\n\n    Args:\n        results (dict): dictionary containing list of values, e.g.\n            {\"train_loss\": [...],\n             \"train_acc\": [...],\n             \"test_loss\": [...],\n             \"test_acc\": [...]}\n    \"\"\"\n    \n    # Get the loss values of the results dictionary (training and test)\n    loss = results['train_loss']\n    test_loss = results['test_loss']\n\n    # Get the accuracy values of the results dictionary (training and test)\n    accuracy = results['train_acc']\n    test_accuracy = results['test_acc']\n\n    # Figure out how many epochs there were\n    epochs = range(len(results['train_loss']))\n\n    # Setup a plot \n    plt.figure(figsize=(15, 7))\n\n    # Plot loss\n    plt.subplot(1, 2, 1)\n    plt.plot(epochs, loss, label='train_loss')\n    plt.plot(epochs, test_loss, label='test_loss')\n    plt.title('Loss')\n    plt.xlabel('Epochs')\n    plt.legend()\n\n    # Plot accuracy\n    plt.subplot(1, 2, 2)\n    plt.plot(epochs, accuracy, label='train_accuracy')\n    plt.plot(epochs, test_accuracy, label='test_accuracy')\n    plt.title('Accuracy')\n    plt.xlabel('Epochs')\n    plt.legend();","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:08:18.707138Z","iopub.execute_input":"2024-07-09T05:08:18.707539Z","iopub.status.idle":"2024-07-09T05:08:18.717136Z","shell.execute_reply.started":"2024-07-09T05:08:18.707486Z","shell.execute_reply":"2024-07-09T05:08:18.716222Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_loss_curves(model_2_results)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:08:23.409702Z","iopub.execute_input":"2024-07-09T05:08:23.410585Z","iopub.status.idle":"2024-07-09T05:08:23.992050Z","shell.execute_reply.started":"2024-07-09T05:08:23.410551Z","shell.execute_reply":"2024-07-09T05:08:23.991187Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Testing","metadata":{}},{"cell_type":"markdown","source":"# Custom Dataloader for test images","metadata":{}},{"cell_type":"code","source":"from torch.utils.data import Dataset\nclass CustomTestDataset(Dataset):\n    def __init__(self, img_dir, transform=None):\n        self.img_dir = img_dir\n        self.transform = transform\n        self.img_paths = list(Path(img_dir).rglob(\"*.*\"))  # List all files in the directory\n\n    def __len__(self):\n        return len(self.img_paths)\n\n    def __getitem__(self, idx):\n        img_path = self.img_paths[idx]\n        image = Image.open(img_path).convert(\"RGB\")  # Convert to RGB\n        if self.transform:\n            image = self.transform(image)\n        return image, img_path.name ","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:18:43.966016Z","iopub.execute_input":"2024-07-09T05:18:43.966385Z","iopub.status.idle":"2024-07-09T05:18:43.973645Z","shell.execute_reply.started":"2024-07-09T05:18:43.966357Z","shell.execute_reply":"2024-07-09T05:18:43.972517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PrepareOutputData:\n    def __init__(self, test_dir, model, batch_size, custom_transform):\n        self.model = model\n        self.transform = custom_transform\n        self.test_data = CustomTestDataset(test_dir, transform=self.transform)\n        self.test_data_loader = torch.utils.data.DataLoader(self.test_data, batch_size=batch_size, shuffle=False)\n        self.submission = [['label']]\n        \n    def evaluatePredictionLabels(self):\n        with torch.inference_mode():\n            self.model.eval()\n            for images, _ in self.test_data_loader:\n                images = images.to(device)\n                pred_logits = self.model(images)\n                prediction_batch = torch.argmax(torch.softmax(pred_logits, dim=1), dim=1)\n                for prediction in prediction_batch:\n                    self.submission.append([prediction.item()])\n                    \n    \n    def savePredictedOutput(self):\n        with open('submission.csv', 'w') as submissionFile:\n            writer = csv.writer(submissionFile)\n            writer.writerows(self.submission)\n\n        print('Output saved to submission.csv')","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:20:04.789091Z","iopub.execute_input":"2024-07-09T05:20:04.789460Z","iopub.status.idle":"2024-07-09T05:20:04.798520Z","shell.execute_reply.started":"2024-07-09T05:20:04.789430Z","shell.execute_reply":"2024-07-09T05:20:04.797653Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:24:52.685172Z","iopub.execute_input":"2024-07-09T05:24:52.685567Z","iopub.status.idle":"2024-07-09T05:24:52.689888Z","shell.execute_reply.started":"2024-07-09T05:24:52.685536Z","shell.execute_reply":"2024-07-09T05:24:52.688966Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_dir = Path(os.path.join(working_dir, \"test1/\"))\noutput_instance = PrepareOutputData(\n    test_dir=test_dir, \n    model=pretrained_vit.to(device), \n    batch_size=BATCH_SIZE,\n    custom_transform=pretrained_vit_transforms\n)\n\n# evaluate images\noutput_instance.evaluatePredictionLabels()","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:20:06.983243Z","iopub.execute_input":"2024-07-09T05:20:06.983614Z","iopub.status.idle":"2024-07-09T05:23:05.731697Z","shell.execute_reply.started":"2024-07-09T05:20:06.983583Z","shell.execute_reply":"2024-07-09T05:23:05.730889Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"output_instance.savePredictedOutput()","metadata":{"execution":{"iopub.status.busy":"2024-07-09T05:24:58.090125Z","iopub.execute_input":"2024-07-09T05:24:58.090809Z","iopub.status.idle":"2024-07-09T05:24:58.103520Z","shell.execute_reply.started":"2024-07-09T05:24:58.090774Z","shell.execute_reply":"2024-07-09T05:24:58.102696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from typing import List, Tuple\n\nfrom PIL import Image\n\n# 1. Take in a trained model, class names, image path, image size, a transform and target device\ndef pred_and_plot_image(model: torch.nn.Module,\n                        image_path: str, \n                        class_names: List[str],\n                        image_size: Tuple[int, int] = (224, 224),\n                        transform: torchvision.transforms = None,\n                        device: torch.device=device):\n    \n    \n    # 2. Open image\n    img = Image.open(image_path)\n\n    # 3. Create transformation for image (if one doesn't exist)\n    if transform is not None:\n        image_transform = transform\n    else:\n        image_transform = transforms.Compose([\n            transforms.Resize(image_size),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                                 std=[0.229, 0.224, 0.225]),\n        ])\n\n    ### Predict on image ### \n\n    # 4. Make sure the model is on the target device\n    model.to(device)\n\n    # 5. Turn on model evaluation mode and inference mode\n    model.eval()\n    with torch.inference_mode():\n      # 6. Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])\n      transformed_image = image_transform(img).unsqueeze(dim=0)\n\n      # 7. Make a prediction on image with an extra dimension and send it to the target device\n      target_image_pred = model(transformed_image.to(device))\n\n    # 8. Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)\n    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)\n\n    # 9. Convert prediction probabilities -> prediction labels\n    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)\n\n    # 10. Plot image with predicted label and probability \n    plt.figure()\n    plt.imshow(img)\n    plt.title(f\"Pred: {class_names[target_image_pred_label]} | Prob: {target_image_pred_probs.max()*100:.2f}%\")\n    plt.axis(False);","metadata":{"execution":{"iopub.status.busy":"2024-07-09T06:49:57.657188Z","iopub.execute_input":"2024-07-09T06:49:57.657586Z","iopub.status.idle":"2024-07-09T06:49:57.668179Z","shell.execute_reply.started":"2024-07-09T06:49:57.657546Z","shell.execute_reply":"2024-07-09T06:49:57.667119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_paths = list(Path(test_dir).rglob(\"*.*\"))\nrandom_image_path = random.choice(img_paths)\npred_and_plot_image(model=pretrained_vit,\n                   image_path=random_image_path,\n                   class_names=class_names)","metadata":{"execution":{"iopub.status.busy":"2024-07-09T06:50:00.309814Z","iopub.execute_input":"2024-07-09T06:50:00.310184Z","iopub.status.idle":"2024-07-09T06:50:00.715177Z","shell.execute_reply.started":"2024-07-09T06:50:00.310154Z","shell.execute_reply":"2024-07-09T06:50:00.714289Z"},"trusted":true},"execution_count":null,"outputs":[]}]}