{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":3649,"databundleVersionId":46718,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Prologue\n## Tackling CIFAR-10 with ViT (Vision Transformer)\n\nThe title is self-explanatory - we're gonna tackle the CIFAR-10 dataset using the Vision Transformer architecture.\n\nVision Transformer is a variant of the Transformer architecture that processed images instead of sentences. The algorithm is pretty simple - an image is divided into patches of the same shape (most commonly via convolution layers), flattened out and embedded (so it'd resemble a sentence), after which the patch embedding is passed into the model.\n\nWe'll use the ViT-Base-16 version, since it's the smallest of them all. And despite that, as you'll see later, it shows really good results. It takes quite a bit of time to train, let alone predict (considering 290,000 out of 300,000 images in the test dataset aren't used to evaluate the final score). On the topic of sizes, train dataset is shrunk to speed up the training time, and once again, despite that, it shows great results.\n\nIf you have any ideas to make the notebook better or questions about how something works in the notebook, you're very welcome to leave a comment regarding the said theme. I'm always welcome for suggestions. ","metadata":{}},{"cell_type":"markdown","source":"# Importing libraries","metadata":{}},{"cell_type":"markdown","source":"We'll need a number of libraries to complete the task at hand. ","metadata":{}},{"cell_type":"code","source":"!pip install py7zr\n\n# base\nimport numpy as np\nimport pandas as pd\n\n# visualization\nimport matplotlib.pyplot as plt\nfrom torchinfo import summary\nfrom tqdm.notebook import tqdm\n\n# machine learning\nimport torch\nimport torchvision\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader, random_split, RandomSampler, SequentialSampler, Subset\nfrom torchvision import datasets\nfrom transformers import get_linear_schedule_with_warmup\n\n# metric tracking\nimport time\nfrom datetime import timedelta\n\n# libraries for dataset unpacking / loading\nimport py7zr\nimport os\nimport PIL","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:23:45.425578Z","iopub.execute_input":"2025-08-18T17:23:45.425851Z","iopub.status.idle":"2025-08-18T17:24:01.122885Z","shell.execute_reply.started":"2025-08-18T17:23:45.425826Z","shell.execute_reply":"2025-08-18T17:24:01.121972Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We'll REALLY need a GPU. It takes a lot of time train the model, let alone predict. You'll see why later on.","metadata":{}},{"cell_type":"code","source":"if torch.cuda.is_available():\n    print(\"CUDA is available; using as the main device\")\n    device = \"cuda\"\nelse:\n    print(\"CUDA is not available; using CPU as the main device\")\n    device = \"cpu\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:01.124800Z","iopub.execute_input":"2025-08-18T17:24:01.125352Z","iopub.status.idle":"2025-08-18T17:24:01.190387Z","shell.execute_reply.started":"2025-08-18T17:24:01.125328Z","shell.execute_reply":"2025-08-18T17:24:01.189682Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Preparing data and the model\n\nInstead of unpacking the zipfile with training data (which takes a lot time to unzip and convert to datasets), we'll use `torchvision.datasets.CIFAR10`, which we can ever so conveniently load and create dataloaders with.\n\nThe model we'll use, as said in the beginning, is ViT-Base-16. It can be found in the `torchvision` library, more precisely - it's `models` module. We'll initialize it with the appropriate weights, freeze all the layers and add our own classifier instead of the original final layer.\n\nIn addition to everything above, we'll have to pass the transforms used by the architecture into the dataset. This'll help the model with processing the images.","metadata":{}},{"cell_type":"code","source":"vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT\nvit_transforms = vit_weights.transforms()\n\nvit_transforms","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:01.191162Z","iopub.execute_input":"2025-08-18T17:24:01.191432Z","iopub.status.idle":"2025-08-18T17:24:01.209260Z","shell.execute_reply.started":"2025-08-18T17:24:01.191409Z","shell.execute_reply":"2025-08-18T17:24:01.208534Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The cell below loads the main CIFAR-10 dataset. Pretty convenient, right?","metadata":{}},{"cell_type":"code","source":"labeled_dataset = datasets.CIFAR10(root=\"/kaggle/working\",\n                                   train=True,\n                                   transform=vit_transforms,\n                                   download=True)\nclasses = labeled_dataset.classes\n\nimage, label = next(iter(labeled_dataset))\nprint(f\"Image tensor shape:  {image.shape}\")\nprint(f\"Label type: {type(label)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:01.210147Z","iopub.execute_input":"2025-08-18T17:24:01.210498Z","iopub.status.idle":"2025-08-18T17:24:09.467864Z","shell.execute_reply.started":"2025-08-18T17:24:01.210469Z","shell.execute_reply":"2025-08-18T17:24:09.467083Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"As I mentioned above, we'll use a smaller version of the dataset to speed up the training. PyTorch's `torch.utils.data.Subset` will assist us in reducing the dataset size.","metadata":{}},{"cell_type":"code","source":"used_dataset_size = 10000\nindices = torch.arange(used_dataset_size)\n\nlabeled_dataset = Subset(labeled_dataset, indices)\n\nprint(f\"Used dataset size: {len(labeled_dataset)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:09.468831Z","iopub.execute_input":"2025-08-18T17:24:09.469088Z","iopub.status.idle":"2025-08-18T17:24:09.475502Z","shell.execute_reply.started":"2025-08-18T17:24:09.469061Z","shell.execute_reply":"2025-08-18T17:24:09.474850Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"We'll use `torch.utils.data.random_split` to split the main (labeled) data into datasets for training and validation.","metadata":{}},{"cell_type":"code","source":"train_dataset_size = int(0.8 * len(labeled_dataset))\nvalidation_dataset_size = len(labeled_dataset) - train_dataset_size\n\ntrain_dataset, validation_dataset = random_split(dataset=labeled_dataset,\n                                                 lengths=(train_dataset_size, validation_dataset_size))\n\nprint(f\"Train dataset length: {len(train_dataset)}\")\nprint(f\"Validation dataset length: {len(validation_dataset)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:09.476150Z","iopub.execute_input":"2025-08-18T17:24:09.476352Z","iopub.status.idle":"2025-08-18T17:24:09.491573Z","shell.execute_reply.started":"2025-08-18T17:24:09.476336Z","shell.execute_reply":"2025-08-18T17:24:09.491021Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Creating dataloaders isn't any more difficult than most of the things above. We'll also inspect an image sample from the training dataloader to ensure everything's working correctly.","metadata":{}},{"cell_type":"code","source":"batch_size = 16\ntrain_sampler = RandomSampler(train_dataset)\nvalidation_sampler = SequentialSampler(validation_dataset)\n\ntrain_dataloader = DataLoader(train_dataset,\n                              batch_size=batch_size,\n                              sampler=train_sampler)\n\nvalidation_dataloader = DataLoader(validation_dataset,\n                                   batch_size=batch_size,\n                                   sampler=validation_sampler)\n\nimage_batch, label_batch = next(iter(train_dataloader))\n\nsample_index = 3\nimage_sample = image_batch[sample_index]\nlabel_sample = label_batch[sample_index]\n\nplt.imshow(image_sample.permute(1, 2, 0))\nplt.title(classes[label_sample])\nplt.axis(False);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:09.493666Z","iopub.execute_input":"2025-08-18T17:24:09.493912Z","iopub.status.idle":"2025-08-18T17:24:09.798091Z","shell.execute_reply.started":"2025-08-18T17:24:09.493886Z","shell.execute_reply":"2025-08-18T17:24:09.797277Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Here's the model loading part.\n\nSince the models from `torchvision` have gradient tracking enabled by default, we'll have to turn it off. We don't want to re-train the model, do we?\n\nIn addition, we'll have to change the last layer to define the output dimension that'll suit us. That's normal practice when using `torchvision.models`.","metadata":{}},{"cell_type":"code","source":"model = torchvision.models.vit_b_16(weights=vit_weights).to(device)\n\nfor parameter in model.parameters():\n    parameter.requires_grad = False\n\nmodel.heads = nn.Linear(768, len(classes)).to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:09.798902Z","iopub.execute_input":"2025-08-18T17:24:09.799162Z","iopub.status.idle":"2025-08-18T17:24:13.134177Z","shell.execute_reply.started":"2025-08-18T17:24:09.799140Z","shell.execute_reply":"2025-08-18T17:24:13.133626Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"`torchinfo.summary` is a great tool to visualize PyTorch models. And not just its layers, but also the input and output shapes across all of them. \n\nYou can also specify additional parameters to inspect more details - for example, `\"trainable\"` in the `col_names` parameter list to show whether a layer is frozen or not.","metadata":{}},{"cell_type":"code","source":"summary(model,\n        input_size=(1, 3, 224, 224),\n        col_names=[\"input_size\", \"output_size\", \"num_params\", \"trainable\"],\n        col_width=20,\n        row_settings=[\"var_names\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:13.134858Z","iopub.execute_input":"2025-08-18T17:24:13.135071Z","iopub.status.idle":"2025-08-18T17:24:13.638174Z","shell.execute_reply.started":"2025-08-18T17:24:13.135054Z","shell.execute_reply":"2025-08-18T17:24:13.637471Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training\n\nWe'll use the hyperparameters specified below to train our model. GPU usage is required if you want to achieve results as fast as possible.","metadata":{}},{"cell_type":"code","source":"epochs = 10\nnum_training_steps = len(train_dataloader) * epochs\n\nloss_function = nn.CrossEntropyLoss()\n\noptimizer = torch.optim.Adam(params=model.parameters(),\n                            lr=3e-3,\n                            betas=(0.9, 0.999),\n                            weight_decay=0.1)\n\nschedule = get_linear_schedule_with_warmup(optimizer,\n                                          num_warmup_steps=0,\n                                          num_training_steps=num_training_steps)\n\noptimizer","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:13.639009Z","iopub.execute_input":"2025-08-18T17:24:13.639786Z","iopub.status.idle":"2025-08-18T17:24:13.645961Z","shell.execute_reply.started":"2025-08-18T17:24:13.639762Z","shell.execute_reply":"2025-08-18T17:24:13.645431Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Dividing the larger task into smaller, more simple pieces is a very nice strategy. You can see it utilized in the cell below. The larger function, `fit`, which is used to train and evaluate the model for a specific number of epochs, is divided into `train_step` and `validation_step`, both of which are simplified main components of the task at hand.\n\nWith their help, the main training function only has to iterate through the given number of epochs, calling the said simplified functions and recording model metrics.\n\n`format_time` will be used to create understandable representations of how much time specific processes took.","metadata":{}},{"cell_type":"code","source":"def format_time(elapsed_time_seconds: time.time) -> str:\n    return str(timedelta(seconds=round(elapsed_time_seconds)))\n\ndef train_step(model,\n               dataloader: DataLoader,\n               loss_function: nn.Module,\n               optimizer: torch.optim.Optimizer,\n               schedule,\n               device: torch.device) -> tuple[float, float]:\n    \"\"\" Performs a single step on a batch from the given dataloader. Returns average loss and accuracy. \"\"\"\n    \n    total_loss = 0.0\n    total_accuracy = 0.0\n    \n    model.train()\n    for feature_batch, label_batch in dataloader:\n        feature_batch = feature_batch.to(device)\n        label_batch = label_batch.to(device)\n\n        logits = model(feature_batch)\n        loss = loss_function(logits, label_batch)\n        total_loss += loss.item()\n\n        predicted_classes = torch.argmax(torch.softmax(logits, dim=-1), dim=-1)\n        accuracy = (predicted_classes == label_batch).sum().item() / len(label_batch)\n        total_accuracy += accuracy\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        schedule.step()\n\n    average_loss = total_loss / len(dataloader)\n    average_accuracy = total_accuracy / len(dataloader)\n\n    return average_loss, average_accuracy\n    \ndef validation_step(model,\n                   dataloader: DataLoader,\n                   loss_function: nn.Module,\n                   device: torch.device) -> tuple[float, float]:\n    \"\"\" Evaluates the passed model on a single batch from the dataloader. Returns average loss and accuracy. \"\"\"\n    \n    total_loss = 0.0\n    total_accuracy = 0.0\n    \n    model.eval()\n    with torch.inference_mode():\n        for feature_batch, label_batch in dataloader:\n            feature_batch = feature_batch.to(device)\n            label_batch = label_batch.to(device)\n            \n            logits = model(feature_batch)\n\n            loss = loss_function(logits, label_batch)\n            total_loss += loss.item()\n\n            predicted_classes = torch.argmax(torch.softmax(logits, dim=-1), dim=-1)\n            accuracy = (predicted_classes == label_batch).sum().item() / len(label_batch)\n            total_accuracy += accuracy\n\n    average_loss = total_loss / len(dataloader)\n    average_accuracy = total_accuracy / len(dataloader)\n\n    return average_loss, average_accuracy\n\ndef fit(model,\n        train_dataloader: DataLoader,\n        validation_dataloader: DataLoader,\n        loss_function: nn.Module,\n        optimizer: torch.optim.Optimizer,\n        schedule,\n        epochs: int,\n        device: torch.device):\n    history = {\n        \"train_loss\": [],\n        \"train_accuracy\": [],\n        \"validation_loss\": [],\n        \"validation_accuracy\": []\n    }\n    \n    for epoch in tqdm(range(1, epochs + 1)):\n        epoch_start_time = time.time()\n        \n        train_loss, train_accuracy = train_step(model,\n                                                dataloader=train_dataloader,\n                                                loss_function=loss_function,\n                                                optimizer=optimizer,\n                                                schedule=schedule,\n                                                device=device)\n\n        validation_loss, validation_accuracy = validation_step(model,\n                                                               dataloader=validation_dataloader,\n                                                               loss_function=loss_function,\n                                                               device=device)\n\n        history[\"train_loss\"].append(train_loss)\n        history[\"train_accuracy\"].append(train_accuracy)\n        history[\"validation_loss\"].append(validation_loss)\n        history[\"validation_accuracy\"].append(validation_accuracy)\n        \n        #print(f\"Epoch {epoch:2d}/{epochs} | train loss: {train_loss:.4f} | train accuracy: {train_accuracy:.4f} | validation loss: {validation_loss:.4f} | validation accuracy: {validation_accuracy:.4f}\")\n        print(f\"Epoch {epoch:2d}/{epochs}\")\n        print(f\"├ Training loss: {train_loss:.4f}, accuracy: {train_accuracy:.4f}\")\n        print(f\"├ Validation loss: {validation_loss:.4f}, accuracy: {validation_accuracy:.4f}\")\n        print(f\"└ Epoch took: {format_time(time.time() - epoch_start_time)}\\n\")\n\n    return history","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:13.646715Z","iopub.execute_input":"2025-08-18T17:24:13.646946Z","iopub.status.idle":"2025-08-18T17:24:13.663765Z","shell.execute_reply.started":"2025-08-18T17:24:13.646926Z","shell.execute_reply":"2025-08-18T17:24:13.662957Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Using the code below is all it takes to train the model. That's it. After it ends, your model is ready for inference.","metadata":{}},{"cell_type":"code","source":"training_start_time = time.time()\nhistory = fit(model,\n              train_dataloader=train_dataloader,\n              validation_dataloader=validation_dataloader,\n              loss_function=loss_function,\n              optimizer=optimizer,\n              schedule=schedule,\n              epochs=epochs,\n              device=device)\n\nprint(f\"Training took: {format_time(time.time() - training_start_time)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:24:13.664381Z","iopub.execute_input":"2025-08-18T17:24:13.664611Z","iopub.status.idle":"2025-08-18T17:42:42.477681Z","shell.execute_reply.started":"2025-08-18T17:24:13.664596Z","shell.execute_reply":"2025-08-18T17:42:42.476752Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"figure, axes = plt.subplots(1, 2, figsize=(12, 5))\naxes[0].plot(history[\"train_loss\"], color=\"green\", label=\"train_loss\")\naxes[0].plot(history[\"validation_loss\"], color=\"orange\", label=\"validation_loss\")\naxes[0].set(title=\"Train and Validation losses\")\naxes[0].legend()\n\naxes[1].plot(history[\"train_accuracy\"], color=\"green\", label=\"train_accuracy\")\naxes[1].plot(history[\"validation_accuracy\"], color=\"orange\", label=\"validation_accuracy\")\naxes[1].set(title=\"Train and Validation accuracies\")\naxes[1].legend()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:46:42.228020Z","iopub.execute_input":"2025-08-18T17:46:42.228591Z","iopub.status.idle":"2025-08-18T17:46:42.557427Z","shell.execute_reply.started":"2025-08-18T17:46:42.228565Z","shell.execute_reply":"2025-08-18T17:46:42.556773Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Predicting\n\nWho would've thought the prediction would take the most time. That's pretty obvious, though! 290,000 out of 300,000 images in the testing dataset are junk, meaning they won't be used to form the final score.\n\nThe cell below will show the available input files.","metadata":{}},{"cell_type":"code","source":"!ls /kaggle/input/cifar-10","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:46:56.283439Z","iopub.execute_input":"2025-08-18T17:46:56.284069Z","iopub.status.idle":"2025-08-18T17:46:56.441261Z","shell.execute_reply.started":"2025-08-18T17:46:56.284048Z","shell.execute_reply":"2025-08-18T17:46:56.440592Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Now we know where the testing dataset is located. Let's unpack it now.","metadata":{}},{"cell_type":"code","source":"image_dir = \"/kaggle/working\"\n\nif not os.path.isdir(image_dir + \"/test\"):\n    with py7zr.SevenZipFile(f\"/kaggle/input/cifar-10/test.7z\", mode=\"r\") as zipf:\n        zipf.extractall(image_dir)\n    print(f\"Test split zipfile extracted\")\n\nelse:     \n    print(f\"Test split zipfile already extracted, skipping...\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:46:57.002625Z","iopub.execute_input":"2025-08-18T17:46:57.003248Z","iopub.status.idle":"2025-08-18T17:48:55.585915Z","shell.execute_reply.started":"2025-08-18T17:46:57.003220Z","shell.execute_reply":"2025-08-18T17:48:55.585172Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Before doing the fun stuff, we'll load the sample submission to utilize it later.","metadata":{}},{"cell_type":"code","source":"submission = pd.read_csv(\"/kaggle/input/cifar-10/sampleSubmission.csv\")\nsubmission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:48:55.587106Z","iopub.execute_input":"2025-08-18T17:48:55.587405Z","iopub.status.idle":"2025-08-18T17:48:55.733097Z","shell.execute_reply.started":"2025-08-18T17:48:55.587388Z","shell.execute_reply":"2025-08-18T17:48:55.732506Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"The custom dataset class defenition is pretty simple. The only part that's more or less difficult is the image IDs. You can get them from the sample submission file and then use them to load the corresponding images, since the image names are their IDs.","metadata":{}},{"cell_type":"code","source":"class CIFAR10TestDataset(Dataset):\n    def __init__(self,\n                 image_folder: str,\n                 image_ids: list,\n                 transforms = None):\n        self.image_folder = image_folder\n        self.image_ids = image_ids\n        self.transforms = transforms\n\n    def __len__(self):\n        return len(self.image_ids)\n    \n    def __getitem__(self, idx):\n        image_path = os.path.join(self.image_folder, self.image_ids[idx] + \".png\")\n        image = PIL.Image.open(image_path).convert(\"RGB\")\n        if self.transforms is not None:\n            image = self.transforms(image)\n            \n        return image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:48:55.733822Z","iopub.execute_input":"2025-08-18T17:48:55.734039Z","iopub.status.idle":"2025-08-18T17:48:55.739128Z","shell.execute_reply.started":"2025-08-18T17:48:55.734015Z","shell.execute_reply":"2025-08-18T17:48:55.738390Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"And that's how we'll create a dataloader for testing. Don't forget to turn the image IDs to strings so they could be appended to string paths.","metadata":{}},{"cell_type":"code","source":"test_folder = \"/kaggle/working/test\"\ntest_image_ids = submission[\"id\"].astype(str).to_list()\ntest_dataset = CIFAR10TestDataset(image_folder=test_folder,\n                                  image_ids=test_image_ids,\n                                  transforms=vit_transforms)\n\ntest_sampler = SequentialSampler(test_dataset)\ntest_dataloader = DataLoader(test_dataset,\n                             batch_size=batch_size,\n                             sampler=test_sampler)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:48:55.740750Z","iopub.execute_input":"2025-08-18T17:48:55.740975Z","iopub.status.idle":"2025-08-18T17:48:55.824397Z","shell.execute_reply.started":"2025-08-18T17:48:55.740960Z","shell.execute_reply":"2025-08-18T17:48:55.823600Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"This cell alone takes the most time to complete. \n\nThe loading bar will show the approximate prediction time. Using GPU P100, it takes around 38 minutes, whereas training via dataset of size 10,000 and batch size of 16 for 10 epochs takes around 18 minutes.\n\nThough, after it finishes, we'll get what we needed all along - a list containing predictions for every image in the testing dataset.","metadata":{}},{"cell_type":"code","source":"predicted_labels = []\n\nmodel.eval()\n\nfor image_batch in tqdm(test_dataloader):\n    image_batch = image_batch.to(device)\n\n    with torch.inference_mode():\n        logits = model(image_batch)\n\n    logits = logits.detach().cpu()\n    predicted_label_batch = torch.argmax(torch.softmax(logits, dim=-1), dim=-1)\n    predicted_labels.extend(predicted_label_batch.tolist())\n\npredicted_classes = [classes[label] for label in predicted_labels]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T17:48:55.825077Z","iopub.execute_input":"2025-08-18T17:48:55.825285Z","iopub.status.idle":"2025-08-18T18:27:25.272676Z","shell.execute_reply.started":"2025-08-18T17:48:55.825269Z","shell.execute_reply":"2025-08-18T18:27:25.271925Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"And this is it. Swap the `label` column in the sample submission dataframe we've loaded before, and you get what you wanted.\n\nDon't forget to turn off the index column creation, since Kaggle's evaluation system doesn't like unnecessary columns.","metadata":{}},{"cell_type":"code","source":"submission[\"label\"] = predicted_classes\nsubmission.to_csv(\"submission.csv\", index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-18T18:27:25.273425Z","iopub.execute_input":"2025-08-18T18:27:25.273617Z","iopub.status.idle":"2025-08-18T18:27:25.502741Z","shell.execute_reply.started":"2025-08-18T18:27:25.273602Z","shell.execute_reply":"2025-08-18T18:27:25.502030Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Epilogue\n\nThe Vision Transformer is, indeed, a very powerful architecture, capable of solving different computer vision tasks. Transformer, the architecture it's based on, has shown great results throughout the years. Many new implementations have found their use. DeepSeek's Multi-head Latent Attention is a great example. It compresses keys and values in a shared representation, which makes it MUCH faster than the original Multi-head Attention.\n\nOnce again, thanks for reading. If you have any questions or ideas that I could implement here, the comment section very warmly welcomes you.","metadata":{}}]}