georgia.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "import shutil\n",
    "import csv\n",
    "import json\n",
    "import zipfile\n",
    "from PIL import Image\n",
    "from random import random\n",
    "import torch\n",
    "from torch.nn.functional import one_hot\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torch.utils.data import DataLoader, Dataset    \n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.transforms import Compose\n",
    "import torchvision.datasets as datasets\n",
    "from torchsummary import summary\n",
    "from numba import cuda\n",
    "\n",
    "cuda.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetGen:\n",
    "    def __init__(self, zip_path, split_ratio, transforms) -> None:\n",
    "        self.zip_path = zip_path\n",
    "        self.split_ratio = split_ratio\n",
    "        self.transforms = transforms\n",
    "        \n",
    "    def create_label_encs(self):\n",
    "        label2id, id2label = {}, {}\n",
    "        with zipfile.ZipFile(self.zip_path, 'r') as zip_file:\n",
    "            for name in zip_file.namelist():\n",
    "                if '/' in name:\n",
    "                    label = name.split('/')[1]\n",
    "                    if label not in label2id:\n",
    "                        label2id[label] = len(label2id)\n",
    "                        id2label[len(id2label)] = label\n",
    "        self.label2id = label2id\n",
    "        self.id2label = id2label\n",
    "        return label2id, id2label\n",
    "\n",
    "    def generate(self, **loader_params):\n",
    "        label2id, id2label = self.create_label_encs()\n",
    "        with zipfile.ZipFile(self.zip_path, 'r') as zip_file:\n",
    "            new_path = pathlib.Path(self.zip_path).parent / zip_file.namelist()[0].split('/')[0]\n",
    "            zip_file.extractall('datasets')\n",
    "        self.split(new_path)\n",
    "        train_dataset = datasets.ImageFolder(\n",
    "            new_path.parent / 'train', transform=self.transforms['train'])\n",
    "        train_loader = DataLoader(train_dataset, **loader_params)\n",
    "        valid_dataset = datasets.ImageFolder(\n",
    "            new_path.parent / 'valid', transform=self.transforms['valid'])\n",
    "        valid_loader = DataLoader(valid_dataset, **loader_params)\n",
    "        test_dataset = datasets.ImageFolder(\n",
    "            new_path.parent / 'test', transform=self.transforms['test'])\n",
    "        test_loader = DataLoader(test_dataset, **loader_params)\n",
    "        self.delete_path(new_path)\n",
    "        return {\n",
    "            'train_loader': train_loader,\n",
    "            'valid_loader': valid_loader,\n",
    "            'test_loader': test_loader,\n",
    "            'label2id': label2id,\n",
    "            'id2label': id2label,\n",
    "            'train_dataset': train_dataset,\n",
    "            'valid_dataset': valid_dataset,\n",
    "            'test_dataset': test_dataset\n",
    "        }\n",
    "\n",
    "    def split(self, path: pathlib.Path):\n",
    "        train_path = (path.parent / 'train')\n",
    "        valid_path = (path.parent / 'valid')\n",
    "        test_path = (path.parent / 'test')\n",
    "\n",
    "        for label in path.iterdir():\n",
    "            (train_path / label.name).mkdir(parents=True)\n",
    "            (valid_path / label.name).mkdir(parents=True)\n",
    "            (test_path / label.name).mkdir(parents=True)\n",
    "            for img in label.iterdir():\n",
    "                rand = random()\n",
    "                if rand < self.split_ratio[0]:\n",
    "                    img.rename(train_path / label.name / img.name)\n",
    "                elif rand < self.split_ratio[0] + self.split_ratio[1]:\n",
    "                    img.rename(valid_path / label.name / img.name)\n",
    "                else:\n",
    "                    img.rename(test_path / label.name / img.name)\n",
    "\n",
    "\n",
    "    def delete_path(self, path):\n",
    "        shutil.rmtree(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformations = {\n",
    "    'train': Compose([\n",
    "        transforms.RandomResizedCrop(224),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'valid': Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ]),\n",
    "    'test': Compose([\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_gen = DatasetGen('datasets/geoguessr-55countries.zip', (0.8, 0.1, 0.1), transformations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "rets = ds_gen.generate(batch_size=32, shuffle=True, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    '''\n",
    "    Convolutional block responsible for getting the hidden state of the image\n",
    "    '''\n",
    "    def __init__(self, input_channel, out_channels: list[int], filter_sizes: list[int], strides: list[int], batch_norm: list[bool], pooling_args: list | None = None, device: torch.device = torch.device('cuda')) -> None:\n",
    "        super().__init__()\n",
    "        self.layers = nn.Sequential()\n",
    "        aux = input_channel\n",
    "        for i, (out_channel, filter_size, stride, norm) in enumerate(zip(out_channels, filter_sizes, strides, batch_norm)):\n",
    "            self.layers.add_module(f'conv_{i}', nn.Conv2d(aux, out_channel, filter_size, stride, padding='same', device=device))\n",
    "            aux = out_channel\n",
    "            if norm:\n",
    "                self.layers.add_module(f'batch_norm_{i}', nn.BatchNorm2d(out_channel, device=device))\n",
    "            self.layers.add_module(f'relu_{i}', nn.ReLU())\n",
    "        if pooling_args is not None:\n",
    "            self.layers.add_module('max_pool', nn.MaxPool2d(*pooling_args))\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InceptionBlock(nn.Module):\n",
    "    def __init__(self, in_channel, out_channels, device=torch.device('cuda')) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        self.branch1 = nn.Sequential(\n",
    "            nn.Conv2d(in_channel, out_channels[0][0], 1, stride=1, padding='same'),\n",
    "            nn.ReLU()\n",
    "        ).to(device)\n",
    "        \n",
    "        self.branch2 = nn.Sequential(\n",
    "            nn.Conv2d(in_channel, out_channels[1][0], 1, stride=1, padding='same'),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(out_channels[1][0], out_channels[1][1], 3, stride=1, padding='same'),\n",
    "            nn.ReLU()\n",
    "        ).to(device)\n",
    "\n",
    "        self.branch3 = nn.Sequential(\n",
    "            nn.Conv2d(in_channel, out_channels[2][0], 1, stride=1, padding='same'),\n",
    "            nn.ReLU(),\n",
    "            nn.Conv2d(out_channels[2][0], out_channels[2][1], 5, stride=1, padding='same'),\n",
    "            nn.ReLU()    \n",
    "        ).to(device)\n",
    "\n",
    "        self.branch4 = nn.Sequential(\n",
    "            nn.MaxPool2d(3, stride=1, padding=1),\n",
    "            nn.Conv2d(in_channel, out_channels[3][0], 1, stride=1, padding='same'),\n",
    "            nn.ReLU()\n",
    "        ).to(device)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x1 = self.branch1(x)\n",
    "        x2 = self.branch2(x)\n",
    "        x3 = self.branch3(x)\n",
    "        x4 = self.branch4(x)\n",
    "\n",
    "        out = torch.cat([x1, x2, x3, x4], 1)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel1(nn.Module): #? Custom model made from scratch\n",
    "    def __init__(self, n_labels, device=torch.device('cuda')) -> None:\n",
    "        super().__init__()\n",
    "        self.n_labels = n_labels \n",
    "        self.layers = nn.Sequential(\n",
    "            EncoderBlock(3, [64, 64], [3, 3], [1, 1], [False, False], [2, 2]),\n",
    "            EncoderBlock(64, [128, 128], [3, 3], [1, 1], [False, True], [2, 2]),\n",
    "            EncoderBlock(128, [256, 256, 256], [3, 3, 3], [1, 1, 1], [False, False, True], [2, 2]),\n",
    "            EncoderBlock(256, [256, 512, 1024], [3, 3, 3], [\n",
    "                         1, 1, 1], [False, False, True], [3, 3]),\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(82944, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, n_labels),\n",
    "            nn.Softmax(dim=1)\n",
    "        ).to(device)\n",
    "\n",
    "    def forward(self, x: torch.Tensor): \n",
    "        x = self.layers(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VGG(nn.Module):\n",
    "    def __init__(self, n_labels) -> None:\n",
    "        super().__init__()\n",
    "        self.n_labels = n_labels\n",
    "        self.model = torchvision.models.vgg19(pretrained=True)\n",
    "        self.model.training = False\n",
    "        self.model.classifier[6] = nn.Linear(4096, 1024)\n",
    "        self.linear = nn.Sequential(\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(1024, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, n_labels)\n",
    "        )\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        x = self.linear(x)\n",
    "        return self.softmax(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Callback:\n",
    "    def __init__(self) -> None:\n",
    "        self.epoch_counter = 0\n",
    "        self.batch_counter = 0\n",
    "\n",
    "    def on_train_begin(self): \n",
    "        self.batch_counter = 0\n",
    "\n",
    "    def on_train_end(self):\n",
    "        return None\n",
    "\n",
    "    def on_epoch_begin(self): \n",
    "        return None\n",
    "    \n",
    "    def on_epoch_end(self, val_loss, model, *args): \n",
    "        self.epoch_counter += 1\n",
    "    \n",
    "    def on_batch_begin(self):\n",
    "        return None\n",
    "    \n",
    "    def on_batch_end(self): \n",
    "        self.batch_counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelCheckpoint(Callback):\n",
    "    def __init__(self, path) -> None:\n",
    "        super().__init__()\n",
    "        self.path = pathlib.Path(path)\n",
    "        self.best_loss = float('inf')\n",
    "\n",
    "    def on_epoch_end(self, val_loss, model, *args):\n",
    "        super().on_epoch_end(val_loss, model, *args)\n",
    "        if val_loss < self.best_loss:\n",
    "            self.best_loss = val_loss\n",
    "            torch.save(model.state_dict(), self.path)\n",
    "\n",
    "class EarlyStopping(Callback):\n",
    "    def __init__(self, patience, restore_best_weights=False, checkpoint: str|bool=False) -> None:\n",
    "        super().__init__()\n",
    "        self.patience = patience\n",
    "        self.restore_best_weights = restore_best_weights\n",
    "        self.checkpoint = checkpoint\n",
    "\n",
    "        if self.checkpoint:\n",
    "            self.model_checkpoint = ModelCheckpoint(self.checkpoint)\n",
    "        if self.restore_best_weights:\n",
    "            self.best_weights = None\n",
    "\n",
    "        self.counter = 0\n",
    "        self.best_loss = float('inf')\n",
    "\n",
    "    def on_epoch_end(self, val_loss, model, *args):\n",
    "        super().on_epoch_end(val_loss, model, *args)\n",
    "        if val_loss < self.best_loss:\n",
    "            self.best_loss = val_loss\n",
    "            self.counter = 0\n",
    "        else:\n",
    "            self.counter += 1\n",
    "            if self.counter >= self.patience:\n",
    "                print(f\"Early stopping on epoch {self.epoch_counter} with val_loss: {val_loss}\")\n",
    "                return True\n",
    "            \n",
    "class LearningRateScheduler(Callback):\n",
    "    def __init__(self, lr_scheduler) -> None:\n",
    "        super().__init__()\n",
    "        self.lr_scheduler = lr_scheduler\n",
    "\n",
    "    def on_epoch_end(self, val_loss, model, *args):\n",
    "        super().on_epoch_end(val_loss, model)\n",
    "        self.lr_scheduler(self.epoch_counter, *args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader: DataLoader, val_loader: DataLoader, loss_func, optimizer, epochs, callbacks: None | list[Callback] = None, device=torch.device('cuda')):\n",
    "    for epoch in range(epochs):\n",
    "        if callbacks is not None and any(callback.on_epoch_begin() for callback in callbacks):\n",
    "            break\n",
    "        model.train()\n",
    "        if callbacks is not None and any(callback.on_train_begin() for callback in callbacks):\n",
    "            break\n",
    "        for images, labels in train_loader:\n",
    "            images, labels = images.to(device), labels.to(device)\n",
    "\n",
    "            labels = one_hot(labels, model.n_labels).to(torch.float32)\n",
    "            if callbacks is not None and any(callback.on_batch_begin() for callback in callbacks):\n",
    "                break\n",
    "            optimizer.zero_grad()\n",
    "            output = model(images)\n",
    "\n",
    "            loss = loss_func(output, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if callbacks is not None and any(callback.on_batch_end() for callback in callbacks):\n",
    "                break\n",
    "        if callbacks is not None and any(callback.on_train_end() for callback in callbacks):\n",
    "            break\n",
    "        \n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            val_loss = 0\n",
    "            val_samples = 0\n",
    "            for images, labels in val_loader:\n",
    "                images, labels = images.to(device), labels.to(device)\n",
    "                labels = one_hot(labels, model.n_labels).to(torch.float32)\n",
    "                output = model(images)\n",
    "                loss = loss_func(output, labels)\n",
    "                val_loss += loss.item()\n",
    "                val_samples += 1\n",
    "            avg_val_loss = val_loss / val_samples\n",
    "            print(f\"Epoch: {epoch+1}, Validation Loss: {avg_val_loss}\")\n",
    "        if callbacks is not None and any(callback.on_epoch_end(avg_val_loss, model) for callback in callbacks):\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CustomModel1(len(rets['label2id']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 224, 224]           1,792\n",
      "              ReLU-2         [-1, 64, 224, 224]               0\n",
      "            Conv2d-3         [-1, 64, 224, 224]          36,928\n",
      "              ReLU-4         [-1, 64, 224, 224]               0\n",
      "         MaxPool2d-5         [-1, 64, 112, 112]               0\n",
      "      EncoderBlock-6         [-1, 64, 112, 112]               0\n",
      "            Conv2d-7        [-1, 128, 112, 112]          73,856\n",
      "              ReLU-8        [-1, 128, 112, 112]               0\n",
      "            Conv2d-9        [-1, 128, 112, 112]         147,584\n",
      "      BatchNorm2d-10        [-1, 128, 112, 112]             256\n",
      "             ReLU-11        [-1, 128, 112, 112]               0\n",
      "        MaxPool2d-12          [-1, 128, 56, 56]               0\n",
      "     EncoderBlock-13          [-1, 128, 56, 56]               0\n",
      "           Conv2d-14          [-1, 256, 56, 56]         295,168\n",
      "             ReLU-15          [-1, 256, 56, 56]               0\n",
      "           Conv2d-16          [-1, 256, 56, 56]         590,080\n",
      "             ReLU-17          [-1, 256, 56, 56]               0\n",
      "           Conv2d-18          [-1, 256, 56, 56]         590,080\n",
      "      BatchNorm2d-19          [-1, 256, 56, 56]             512\n",
      "             ReLU-20          [-1, 256, 56, 56]               0\n",
      "        MaxPool2d-21          [-1, 256, 28, 28]               0\n",
      "     EncoderBlock-22          [-1, 256, 28, 28]               0\n",
      "           Conv2d-23          [-1, 256, 28, 28]         590,080\n",
      "             ReLU-24          [-1, 256, 28, 28]               0\n",
      "           Conv2d-25          [-1, 512, 28, 28]       1,180,160\n",
      "             ReLU-26          [-1, 512, 28, 28]               0\n",
      "           Conv2d-27         [-1, 1024, 28, 28]       4,719,616\n",
      "      BatchNorm2d-28         [-1, 1024, 28, 28]           2,048\n",
      "             ReLU-29         [-1, 1024, 28, 28]               0\n",
      "        MaxPool2d-30           [-1, 1024, 9, 9]               0\n",
      "     EncoderBlock-31           [-1, 1024, 9, 9]               0\n",
      "          Flatten-32                [-1, 82944]               0\n",
      "           Linear-33                  [-1, 512]      42,467,840\n",
      "             ReLU-34                  [-1, 512]               0\n",
      "           Linear-35                   [-1, 55]          28,215\n",
      "          Softmax-36                   [-1, 55]               0\n",
      "================================================================\n",
      "Total params: 50,724,215\n",
      "Trainable params: 50,724,215\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 253.03\n",
      "Params size (MB): 193.50\n",
      "Estimated Total Size (MB): 447.10\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "summary(model, (3, 224, 224))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1, Validation Loss: 4.01376014473164\n",
      "Epoch: 2, Validation Loss: 4.013760155281134\n",
      "Epoch: 3, Validation Loss: 4.013760157391033\n"
     ]
    }
   ],
   "source": [
    "train(model, rets['train_loader'], rets['valid_loader'], nn.CrossEntropyLoss(), torch.optim.Adam(model.parameters()), 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_442627/1924512865.py:2: FutureWarning: The input object of type 'Tensor' is an array-like implementing one of the corresponding protocols (`__array__`, `__array_interface__` or `__array_struct__`); but not a sequence (or 0-D). In the future, this object will be coerced as if it was first converted using `np.array(obj)`. To retain the old behaviour, you have to either modify the type 'Tensor', or assign to an empty array created with `np.empty(correct_shape, dtype=object)`.\n",
      "  arr = np.array(sample)\n",
      "/tmp/ipykernel_442627/1924512865.py:2: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  arr = np.array(sample)\n"
     ]
    }
   ],
   "source": [
    "sample = rets['test_dataset'][0], rets['test_dataset'][1000], rets['test_dataset'][2735]\n",
    "arr = np.array(sample)\n",
    "sample_in = arr[:, 0].tolist()\n",
    "sample_lbl = arr[:, 1].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 17, 42]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_lbl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.]], device='cuda:0')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = one_hot(torch.tensor(sample_lbl), len(\n",
    "    rets['label2id'])).to(torch.float32).cuda()\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.],\n",
       "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0.]], device='cuda:0', grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out = model(torch.stack(sample_in).to('cuda'))\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 55]) torch.Size([3, 55])\n"
     ]
    }
   ],
   "source": [
    "print(out.shape, labels.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(4.0381, device='cuda:0', grad_fn=<DivBackward1>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nn.CrossEntropyLoss()(out, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([39, 39, 39], device='cuda:0')"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.argmax(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}