Transformer - CIFAR10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "23e83631",
   "metadata": {
    "id": "23e83631"
   },
   "source": [
    "# Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319062fd",
   "metadata": {
    "id": "319062fd"
   },
   "source": [
    "Authors:\n",
    "- Marina Debogović\n",
    "- Marko Njegomir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "092182cd",
   "metadata": {
    "id": "092182cd"
   },
   "source": [
    "## VIT transformer implemented in pytorch that is used as a baseline for comparison\n",
    "\n",
    "[Code for regular implementation of VIT transformer in pytorch](https://towardsdatascience.com/a-demonstration-of-using-vision-transformers-in-pytorch-mnist-handwritten-digit-recognition-407eafbc15b0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "Jh-Almm1kPu3",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jh-Almm1kPu3",
    "outputId": "b4ae8a27-ccd4-4d79-ae7e-304e663e75fb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://download.pytorch.org/whl/cu116\n",
      "Requirement already satisfied: torch in c:\\users\\marko\\anaconda3\\lib\\site-packages (1.10.1)\n",
      "Requirement already satisfied: torchvision in c:\\users\\marko\\anaconda3\\lib\\site-packages (0.11.2)\n",
      "Requirement already satisfied: torchaudio in c:\\users\\marko\\anaconda3\\lib\\site-packages (0.10.1)\n",
      "Requirement already satisfied: typing_extensions in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch) (4.2.0)\n",
      "Requirement already satisfied: numpy in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torchvision) (1.21.6)\n",
      "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torchvision) (9.1.0)\n",
      "Looking in links: https://data.pyg.org/whl/torch-1.12.0+cu116.html\n",
      "Requirement already satisfied: torch-scatter in c:\\users\\marko\\anaconda3\\lib\\site-packages (2.0.9)\n",
      "Requirement already satisfied: torch-sparse in c:\\users\\marko\\anaconda3\\lib\\site-packages (0.6.12)\n",
      "Requirement already satisfied: torch-cluster in c:\\users\\marko\\anaconda3\\lib\\site-packages (1.5.9)\n",
      "Requirement already satisfied: torch-spline-conv in c:\\users\\marko\\anaconda3\\lib\\site-packages (1.2.1)\n",
      "Requirement already satisfied: torch-geometric in c:\\users\\marko\\anaconda3\\lib\\site-packages (2.0.3)\n",
      "Requirement already satisfied: scipy in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-sparse) (1.8.0)\n",
      "Requirement already satisfied: tqdm in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (4.64.0)\n",
      "Requirement already satisfied: rdflib in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (6.1.1)\n",
      "Requirement already satisfied: numpy in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.21.6)\n",
      "Requirement already satisfied: requests in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (2.27.1)\n",
      "Requirement already satisfied: pyparsing in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (3.0.8)\n",
      "Requirement already satisfied: yacs in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (0.1.8)\n",
      "Requirement already satisfied: jinja2 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (3.1.1)\n",
      "Requirement already satisfied: PyYAML in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (6.0)\n",
      "Requirement already satisfied: scikit-learn in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.0.2)\n",
      "Requirement already satisfied: pandas in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.4.2)\n",
      "Requirement already satisfied: googledrivedownloader in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (0.4)\n",
      "Requirement already satisfied: networkx in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (2.8)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from jinja2->torch-geometric) (2.1.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from pandas->torch-geometric) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from pandas->torch-geometric) (2022.1)\n",
      "Requirement already satisfied: setuptools in c:\\users\\marko\\anaconda3\\lib\\site-packages (from rdflib->torch-geometric) (62.1.0)\n",
      "Requirement already satisfied: isodate in c:\\users\\marko\\anaconda3\\lib\\site-packages (from rdflib->torch-geometric) (0.6.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests->torch-geometric) (3.3)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests->torch-geometric) (2.0.12)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests->torch-geometric) (2022.6.15)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests->torch-geometric) (1.26.11)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from scikit-learn->torch-geometric) (3.1.0)\n",
      "Requirement already satisfied: joblib>=0.11 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from scikit-learn->torch-geometric) (1.1.0)\n",
      "Requirement already satisfied: colorama in c:\\users\\marko\\anaconda3\\lib\\site-packages (from tqdm->torch-geometric) (0.4.4)\n",
      "Requirement already satisfied: six>=1.5 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from python-dateutil>=2.8.1->pandas->torch-geometric) (1.16.0)\n",
      "Requirement already satisfied: einops in c:\\users\\marko\\anaconda3\\lib\\site-packages (0.4.1)\n",
      "Requirement already satisfied: wandb in c:\\users\\marko\\anaconda3\\lib\\site-packages (0.13.1)\n",
      "Requirement already satisfied: shortuuid>=0.5.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (1.0.9)\n",
      "Requirement already satisfied: PyYAML in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (6.0)\n",
      "Requirement already satisfied: requests<3,>=2.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (2.27.1)\n",
      "Requirement already satisfied: GitPython>=1.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (3.1.27)\n",
      "Requirement already satisfied: docker-pycreds>=0.4.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (0.4.0)\n",
      "Requirement already satisfied: protobuf<4.0dev,>=3.12.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (3.14.0)\n",
      "Requirement already satisfied: six>=1.13.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (1.16.0)\n",
      "Requirement already satisfied: Click!=8.0.0,>=7.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (8.1.3)\n",
      "Requirement already satisfied: setproctitle in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (1.3.1)\n",
      "Requirement already satisfied: pathtools in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (0.1.2)\n",
      "Requirement already satisfied: psutil>=5.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (5.9.0)\n",
      "Requirement already satisfied: setuptools in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (62.1.0)\n",
      "Requirement already satisfied: promise<3,>=2.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (2.3)\n",
      "Requirement already satisfied: sentry-sdk>=1.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (1.9.2)\n",
      "Requirement already satisfied: colorama in c:\\users\\marko\\anaconda3\\lib\\site-packages (from Click!=8.0.0,>=7.0->wandb) (0.4.4)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from GitPython>=1.0.0->wandb) (4.0.9)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests<3,>=2.0.0->wandb) (3.3)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests<3,>=2.0.0->wandb) (2.0.12)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests<3,>=2.0.0->wandb) (1.26.11)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from requests<3,>=2.0.0->wandb) (2022.6.15)\n",
      "Requirement already satisfied: smmap<6,>=3.0.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from gitdb<5,>=4.0.1->GitPython>=1.0.0->wandb) (3.0.5)\n"
     ]
    }
   ],
   "source": [
    "# Pytorch installation\n",
    "!pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116\n",
    "# Pytorch Geometric installation\n",
    "!pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.12.0+cu116.html\n",
    "# Einops installation (Flexible and powerful tensor operations for readable and reliable code)\n",
    "!pip install einops\n",
    "# Weights and Biases installation\n",
    "!pip install wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4b4d80d6",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "4b4d80d6",
    "outputId": "fa8f3c4b-cb95-4cfd-d23f-ed1d2836462f"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'NVIDIA GeForce GTX 1660 Ti'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import time\n",
    "from torch import optim\n",
    "import wandb\n",
    "from sklearn import metrics \n",
    "import os\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b77c71c6",
   "metadata": {
    "id": "b77c71c6"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch import nn\n",
    "from einops import rearrange\n",
    "\n",
    "class Residual(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(x, **kwargs) + x\n",
    "\n",
    "class PreNorm(nn.Module):\n",
    "    def __init__(self, dim, fn):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(dim)\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.fn(self.norm(x), **kwargs)\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(dim, hidden_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(hidden_dim, dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    def __init__(self, dim, heads=8):\n",
    "        super().__init__()\n",
    "        self.heads = heads\n",
    "        self.scale = dim ** -0.5\n",
    "\n",
    "        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)\n",
    "        self.to_out = nn.Linear(dim, dim)\n",
    "\n",
    "    def forward(self, x, mask = None):\n",
    "        b, n, _, h = *x.shape, self.heads\n",
    "        qkv = self.to_qkv(x)\n",
    "        q, k, v = rearrange(qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=h)\n",
    "\n",
    "        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale\n",
    "        # matricno mnozenje q i k\n",
    "\n",
    "        if mask is not None:\n",
    "            mask = F.pad(mask.flatten(1), (1, 0), value = True)\n",
    "            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'\n",
    "            mask = mask[:, None, :] * mask[:, :, None]\n",
    "            dots.masked_fill_(~mask, float('-inf'))\n",
    "            del mask\n",
    "\n",
    "        attn = dots.softmax(dim=-1)\n",
    "\n",
    "        out = torch.einsum('bhij,bhjd->bhid', attn, v)\n",
    "        out = rearrange(out, 'b h n d -> b n (h d)')\n",
    "        out =  self.to_out(out)\n",
    "        return out\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "    def __init__(self, dim, depth, heads, mlp_dim):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList([])\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                Residual(PreNorm(dim, Attention(dim, heads = heads))),\n",
    "                Residual(PreNorm(dim, FeedForward(dim, mlp_dim)))\n",
    "            ]))\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x, mask=mask)\n",
    "            x = ff(x)\n",
    "        return x\n",
    "\n",
    "class ViT(nn.Module):\n",
    "    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3):\n",
    "        super().__init__()\n",
    "        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'\n",
    "        num_patches = (image_size // patch_size) ** 2\n",
    "        patch_dim = channels * patch_size ** 2\n",
    "\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))\n",
    "        self.patch_to_embedding = nn.Linear(patch_dim, dim)\n",
    "        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))\n",
    "        self.transformer = Transformer(dim, depth, heads, mlp_dim)\n",
    "\n",
    "        self.to_cls_token = nn.Identity()\n",
    "\n",
    "        self.mlp_head = nn.Sequential(\n",
    "            nn.Linear(dim, mlp_dim),\n",
    "            nn.GELU(),\n",
    "            nn.Linear(mlp_dim, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, img, mask=None):\n",
    "        if mask: mask.to(device)\n",
    "        p = self.patch_size\n",
    "\n",
    "        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)       \n",
    "        x = x.to(device)\n",
    "        x = self.patch_to_embedding(x)\n",
    "\n",
    "        cls_tokens = self.cls_token.expand(img.shape[0], -1, -1)\n",
    "        cls_tokens = cls_tokens.to(device)\n",
    "        x = torch.cat((cls_tokens, x), dim=1)\n",
    "        x += self.pos_embedding\n",
    "        x = self.transformer(x, mask)\n",
    "\n",
    "        x = self.to_cls_token(x[:, 0])\n",
    "        x = x.to(device)\n",
    "        return self.mlp_head(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "13a02805",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 173,
     "referenced_widgets": [
      "b0891c5aaee44dfc89bc7f349503067b",
      "5338af9b42d74a0296190ddfa826838e",
      "68be3b568a454d71b68800bc9f65ac5e",
      "b203b9e3b730464bbf6d96c3d792ecc4",
      "adf7d92793a64e659eaeb6f346adbd2c",
      "fe106ac6bb604381a0e9192e7507b250",
      "e74f53a63d9040f79ff2721f32fb0d1e",
      "d483e53f040e4e438f18dbc6154602aa"
     ]
    },
    "id": "13a02805",
    "outputId": "ff097202-3699-4b06-fc22-b30d6dcc78c0"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mnjmarko\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.1"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>C:\\Users\\Marko\\Desktop\\temp\\graph-transformer-psiml\\wandb\\run-20220814_194001-2u0e4o6v</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/njmarko/CIFAR10/runs/2u0e4o6v\" target=\"_blank\">daily-microwave-5</a></strong> to <a href=\"https://wandb.ai/njmarko/CIFAR10\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "# torch.manual_seed(42)\n",
    "projectName = 'CIFAR10'\n",
    "wandb.init(entity = 'njmarko', project = projectName)\n",
    "\n",
    "DOWNLOAD_PATH = '/data/cifar10'\n",
    "BATCH_SIZE_TRAIN = 250\n",
    "BATCH_SIZE_VAL = 250\n",
    "BATCH_SIZE_TEST = 250\n",
    "\n",
    "mean = [0.4914, 0.4822, 0.4465]\n",
    "std = [0.2023, 0.1994, 0.2010]\n",
    "\n",
    "transform_cifar10 = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "                                                  torchvision.transforms.transforms.Normalize(mean, std)])\n",
    "train_data = torchvision.datasets.CIFAR10(DOWNLOAD_PATH, train=True, download=True, transform=transform_cifar10)\n",
    "train_set, val_set = torch.utils.data.random_split(train_data, [45000, 5000])\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE_TRAIN, shuffle=True)\n",
    "val_loader = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE_VAL, shuffle=True)\n",
    "test_set = torchvision.datasets.CIFAR10(DOWNLOAD_PATH, train=False, download=True, transform=transform_cifar10)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=BATCH_SIZE_TEST, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bc7abb4b",
   "metadata": {
    "id": "bc7abb4b"
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, optimizer, data_loader, loss_history, scheduler):\n",
    "    total_samples = len(data_loader.dataset)\n",
    "    model.train()\n",
    "\n",
    "    correct_samples = 0\n",
    "    for i, (data, target) in enumerate(data_loader):\n",
    "        optimizer.zero_grad()\n",
    "        output = F.log_softmax(model(data), dim=1)\n",
    "\n",
    "        target = target.to(device)\n",
    "        output = output.to(device)\n",
    "        \n",
    "        _, pred = torch.max(output, dim=1)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        correct_samples += pred.eq(target).sum()\n",
    "        target = target.cpu().detach().numpy()\n",
    "        pred = pred.cpu().detach().numpy()\n",
    "      \n",
    "        \n",
    "        f1_score = metrics.f1_score(target, pred, average='micro')\n",
    "\n",
    "        if i % 100 == 0:\n",
    "            print('[' +  '{:5}'.format(i * len(data)) + '/' + '{:5}'.format(total_samples) +\n",
    "                  ' (' + '{:3.0f}'.format(100 * i / len(data_loader)) + '%)]  Loss: ' +\n",
    "                  '{:6.4f}'.format(loss.item()))\n",
    "            loss_history.append(loss.item())\n",
    "            wandb.log({\n",
    "                'train_loss': loss.item(),\n",
    "                'train_f1_score': f1_score\n",
    "            })\n",
    "            \n",
    "    acc = 100.0 * correct_samples / total_samples\n",
    "    wandb.log({\n",
    "        'train_accuracy': acc\n",
    "    })\n",
    "    print(f'Accuracy: ' + '{:4.2f}'.format(acc) + '%')        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a9828c4a",
   "metadata": {
    "id": "a9828c4a"
   },
   "outputs": [],
   "source": [
    "def validate(model, data_loader, loss_history):\n",
    "    model.eval()\n",
    "    \n",
    "    total_samples = len(data_loader.dataset)\n",
    "    correct_samples = 0\n",
    "    total_loss = 0\n",
    "    \n",
    "    global_target = np.array([])\n",
    "    global_pred = np.array([]) \n",
    "\n",
    "    with torch.no_grad():\n",
    "        for data, target in data_loader:\n",
    "            res = model(data)\n",
    "            res = res.to(device)\n",
    "            output = F.log_softmax(res, dim=1)\n",
    "            target = target.to(device)\n",
    "            output = output.to(device)\n",
    "            loss = F.nll_loss(output, target, reduction='sum')\n",
    "            _, pred = torch.max(output, dim=1)\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            correct_samples += pred.eq(target).sum()\n",
    "            \n",
    "            target = target.cpu().detach().numpy()\n",
    "            pred = pred.cpu().detach().numpy()\n",
    "            \n",
    "            global_target = np.concatenate((global_target, target))\n",
    "            global_pred = np.concatenate((global_pred, pred))\n",
    "            \n",
    "    avg_loss = total_loss / total_samples\n",
    "    acc = 100.0 * correct_samples / total_samples\n",
    "    loss_history.append(avg_loss)\n",
    "\n",
    "   \n",
    "    f1_score = metrics.f1_score(global_target, global_pred, average='micro')\n",
    "\n",
    "    print('\\nAverage test loss: ' + '{:.4f}'.format(avg_loss) +\n",
    "          '  Accuracy:' + '{:5}'.format(correct_samples) + '/' +\n",
    "          '{:5}'.format(total_samples) + ' (' +\n",
    "          '{:4.2f}'.format(100.0 * correct_samples / total_samples) + '%)\\n')\n",
    "    \n",
    "    wandb.log({\n",
    "        'val_loss': loss.item(),\n",
    "        'val_f1_score': f1_score,\n",
    "        'val_accuracy': acc\n",
    "    })\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b682f1e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, data_loader, loss_history):\n",
    "    model.eval()\n",
    "    \n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    \n",
    "    total_samples = len(data_loader.dataset)\n",
    "    correct_samples = 0\n",
    "    total_loss = 0\n",
    "    \n",
    "    global_target = np.array([])\n",
    "    global_pred = np.array([]) \n",
    "\n",
    "    with torch.no_grad():\n",
    "        for data, target in data_loader:\n",
    "            res = model(data)\n",
    "            res = res.to(device)\n",
    "            output = F.log_softmax(res, dim=1)\n",
    "            target = target.to(device)\n",
    "            output = output.to(device)\n",
    "            loss = F.nll_loss(output, target, reduction='sum')\n",
    "            _, pred = torch.max(output, dim=1)\n",
    "            \n",
    "            total_loss += loss.item()\n",
    "            correct_samples += pred.eq(target).sum()\n",
    "            \n",
    "            target = target.cpu().detach().numpy()\n",
    "            pred = pred.cpu().detach().numpy()\n",
    "            \n",
    "            global_target = np.concatenate((global_target, target))\n",
    "            global_pred = np.concatenate((global_pred, pred))\n",
    "\n",
    "    avg_loss = total_loss / total_samples\n",
    "    acc = 100.0 * correct_samples / total_samples\n",
    "    loss_history.append(avg_loss)\n",
    "\n",
    "    f1_score = metrics.f1_score(global_target, global_pred, average='micro')\n",
    "    precision = metrics.precision_score(global_target, global_pred, average='micro')\n",
    "    recall = metrics.recall_score(global_target, global_pred, average='micro')\n",
    "\n",
    "    print('\\nTest loss: ' + '{:.4f}'.format(avg_loss) +\n",
    "          '  Accuracy:' + '{:5}'.format(correct_samples) + '/' +\n",
    "          '{:5}'.format(total_samples) + ' (' +\n",
    "          '{:4.2f}'.format(acc) + '%)  Precision: ' + '{:4.2f}'.format(precision) +\n",
    "          '  Recall: ' + '{:4.2f}'.format(recall) + '\\n')\n",
    "    \n",
    "    cm = metrics.confusion_matrix(global_target, global_pred)\n",
    "    print(f'Confusion matrix:\\n {cm}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e171700b",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 676
    },
    "id": "e171700b",
    "outputId": "2af08caa-c733-46b8-8949-2adbc2fd7885"
   },
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "N_EPOCHS = 20\n",
    "best_acc = -1\n",
    "best_epoch = -1\n",
    "\n",
    "path = 'CIFAR10-classic'\n",
    "Path(path).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "start_time = time.time()\n",
    "model = ViT(image_size=32, patch_size=8, num_classes=10, channels=3,\n",
    "            dim=128, depth=8, heads=8, mlp_dim=256)\n",
    "optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.05)\n",
    "early_stop_tolerance = 10e-4\n",
    "model = model.to(device)\n",
    "scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.0001, max_lr=0.001, cycle_momentum=False, step_size_up = 1000)\n",
    "\n",
    "\n",
    "train_loss_history, val_loss_history, test_loss_history = [], [], []\n",
    "for epoch in range(1, N_EPOCHS + 1):\n",
    "    early_stopping = 0\n",
    "    print('Epoch:', epoch)\n",
    "    train_epoch(model, optimizer, train_loader, train_loss_history, scheduler)\n",
    "    acc = validate(model, val_loader, val_loss_history)\n",
    "        \n",
    "    if best_acc < acc:\n",
    "        best_acc = acc\n",
    "        torch.save(model.state_dict(), os.path.join(path, f'classic_transfromer-{projectName}-acc{acc}'))\n",
    "\n",
    "    if len(train_loss_history) > 2 and np.isclose(train_loss_history[-2], train_loss_history[-1], atol=early_stop_tolerance):\n",
    "        early_stopping += 1\n",
    "        if (early_stopping == 5):\n",
    "            print(f\"Early stop with tolerance {early_stop_tolerance} for losses {train_loss_history[-2]} and {train_loss_history[-1]}\")\n",
    "            break\n",
    "    else:\n",
    "        early_stopping = 0\n",
    "        \n",
    "# Testiranje modela\n",
    "path2 = f'CIFAR10-classic/classic_transfromer-{projectName}-acc{best_acc}'\n",
    "print(path2)\n",
    "model.load_state_dict(torch.load(path2))\n",
    "model = model.to(device)\n",
    "evaluate(model, test_loader, test_loss_history)\n",
    "\n",
    "        \n",
    "print('Execution time:', '{:5.2f}'.format(time.time() - start_time), 'seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf0315c",
   "metadata": {
    "id": "0cf0315c"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Transformer.ipynb",
   "provenance": []
  },
  "gpuClass": "standard",
  "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.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "5338af9b42d74a0296190ddfa826838e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "LabelModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "LabelModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "LabelView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_adf7d92793a64e659eaeb6f346adbd2c",
      "placeholder": "​",
      "style": "IPY_MODEL_fe106ac6bb604381a0e9192e7507b250",
      "value": "0.008 MB of 0.008 MB uploaded (0.000 MB deduped)\r"
     }
    },
    "68be3b568a454d71b68800bc9f65ac5e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_e74f53a63d9040f79ff2721f32fb0d1e",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_d483e53f040e4e438f18dbc6154602aa",
      "value": 1
     }
    },
    "adf7d92793a64e659eaeb6f346adbd2c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b0891c5aaee44dfc89bc7f349503067b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "VBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "VBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "VBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_5338af9b42d74a0296190ddfa826838e",
       "IPY_MODEL_68be3b568a454d71b68800bc9f65ac5e"
      ],
      "layout": "IPY_MODEL_b203b9e3b730464bbf6d96c3d792ecc4"
     }
    },
    "b203b9e3b730464bbf6d96c3d792ecc4": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d483e53f040e4e438f18dbc6154602aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "e74f53a63d9040f79ff2721f32fb0d1e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "fe106ac6bb604381a0e9192e7507b250": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}