GATTransformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ccca774a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Custom GNN Transformer architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af43b49",
   "metadata": {},
   "source": [
    "Authors:\n",
    "- Marina Debogović\n",
    "- Marko Njegomir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f106786b",
   "metadata": {},
   "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: rdflib in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (6.1.1)\n",
      "Requirement already satisfied: pandas in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.4.2)\n",
      "Requirement already satisfied: tqdm in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (4.64.0)\n",
      "Requirement already satisfied: networkx in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (2.8)\n",
      "Requirement already satisfied: yacs in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (0.1.8)\n",
      "Requirement already satisfied: pyparsing in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (3.0.8)\n",
      "Requirement already satisfied: PyYAML in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (6.0)\n",
      "Requirement already satisfied: googledrivedownloader in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (0.4)\n",
      "Requirement already satisfied: numpy in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.21.6)\n",
      "Requirement already satisfied: jinja2 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (3.1.1)\n",
      "Requirement already satisfied: scikit-learn in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (1.0.2)\n",
      "Requirement already satisfied: requests in c:\\users\\marko\\anaconda3\\lib\\site-packages (from torch-geometric) (2.27.1)\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: pytz>=2020.1 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from pandas->torch-geometric) (2022.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: 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: 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: 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: joblib>=0.11 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from scikit-learn->torch-geometric) (1.1.0)\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: 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: GitPython>=1.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (3.1.27)\n",
      "Requirement already satisfied: setuptools in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (62.1.0)\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: pathtools in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (0.1.2)\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: setproctitle in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (1.3.1)\n",
      "Requirement already satisfied: promise<3,>=2.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (2.3)\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: psutil>=5.0.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (5.9.0)\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: PyYAML in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (6.0)\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: protobuf<4.0dev,>=3.12.0 in c:\\users\\marko\\anaconda3\\lib\\site-packages (from wandb) (3.14.0)\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: 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: 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: 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: 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: 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": "e37cf507",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import time\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric\n",
    "from torch_geometric.nn import GATConv, GATv2Conv\n",
    "import torchvision\n",
    "from einops import rearrange\n",
    "import wandb\n",
    "from sklearn import metrics \n",
    "import os\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "de952e9f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "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, **kwargs):\n",
    "        return self.net(x)\n",
    "    \n",
    "class QuickFix(nn.Module):\n",
    "    def __init__(self, dim, heads, fn):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.heads = heads\n",
    "        self.linear = nn.Linear(dim * heads, dim)\n",
    "        self.fn = fn\n",
    "        \n",
    "    def forward(self, x, **kwargs):\n",
    "        return self.linear(self.fn(x, **kwargs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6d3d558f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Out implementation of GAT Transformer\n",
    "\n",
    "class GATTransformer(torch.nn.Module):\n",
    "    def __init__(self, dim, depth, heads, mlp_dim):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        self.depth = depth\n",
    "        for _ in range(depth):\n",
    "            self.layers.append(nn.ModuleList([\n",
    "                Residual(PreNorm(dim, QuickFix(dim, heads, GATv2Conv(in_channels=dim, out_channels=dim, heads=heads, add_self_loops=False)))),\n",
    "                Residual(PreNorm(dim, FeedForward(dim, mlp_dim)))\n",
    "            ]))\n",
    "\n",
    "\n",
    "    def forward(self, data):\n",
    "        x, edge_index = data\n",
    "\n",
    "        for attn, ff in self.layers:\n",
    "            x = attn(x=x, edge_index=edge_index)\n",
    "            x = ff(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "82198700",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "class GraphViT(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.dim = dim\n",
    "\n",
    "        self.patch_size = patch_size\n",
    "        self.num_patches = num_patches\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 = GATTransformer(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",
    "        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",
    "        \n",
    "    \n",
    "        x_shape = x.shape\n",
    "        edge_index = torch.ones((2, x.shape[1]**2*x.shape[0]), dtype=torch.long)\n",
    "        for i in range(x.shape[0]*x.shape[1]):\n",
    "            for j in range(x.shape[1]):\n",
    "                edge_index[0,i*x.shape[1] + j] = i\n",
    "                edge_index[1,i*x.shape[1] + j] = x.shape[1]*(i // x.shape[1]) + j\n",
    "        x = x.view((x.shape[0]*x.shape[1], x.shape[2]))\n",
    "        edge_index = edge_index.to(device)\n",
    "\n",
    "        \n",
    "        x = self.transformer([x, edge_index])\n",
    "        x = x.view(x_shape)\n",
    "        \n",
    "        x = self.to_cls_token(x[:,0])\n",
    "        x = x.to(device)\n",
    "        \n",
    "        out = self.mlp_head(x)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b346c02f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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_193719-3tdcyodv</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/gat-test/runs/3tdcyodv\" target=\"_blank\">volcanic-glitter-13</a></strong> to <a href=\"https://wandb.ai/njmarko/gat-test\" 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"
    }
   ],
   "source": [
    "# torch.manual_seed(42)\n",
    "wandb.init(entity = 'njmarko', project = 'gat-test')\n",
    "\n",
    "DOWNLOAD_PATH = '/data/mnist'\n",
    "BATCH_SIZE_TRAIN = 200\n",
    "BATCH_SIZE_VAL = 200\n",
    "BATCH_SIZE_TEST = 200\n",
    "\n",
    "## CIFAR 10 transforms\n",
    "# mean = [0.4914, 0.4822, 0.4465]\n",
    "# std = [0.2023, 0.1994, 0.2010]\n",
    "\n",
    "# transform_mnist = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "#                                torchvision.transforms.Normalize(mean, std)])\n",
    "\n",
    "## MNIST transforms\n",
    "transform_mnist = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "                                                  torchvision.transforms.Normalize((0.1307,), (0.3081,))])\n",
    "\n",
    "train_set = torchvision.datasets.MNIST(DOWNLOAD_PATH, train=True, download=True,\n",
    "                                       transform=transform_mnist)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE_TRAIN, shuffle=True)\n",
    "\n",
    "test_set = torchvision.datasets.MNIST(DOWNLOAD_PATH, train=False, download=True,\n",
    "                                      transform=transform_mnist)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size=BATCH_SIZE_TEST, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "47c7a4af",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def train_epoch(model, optimizer, data_loader, loss_history, scheduler, best_model_acc):\n",
    "    total_samples = len(data_loader.dataset)\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    old_time = time.time()\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",
    "        target = target.to(device)\n",
    "        output = output.to(device)\n",
    "        \n",
    "        _, pred = torch.max(output, dim=1)\n",
    "        \n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "        running_loss += loss.item()\n",
    "        \n",
    "        correct_samples += pred.eq(target).sum()\n",
    "        target = target.cpu().detach().numpy()\n",
    "        pred = pred.cpu().detach().numpy()\n",
    "        \n",
    "        f1_score = metrics.f1_score(target, pred, average='micro')\n",
    "        \n",
    "        if i % 100 == 0:\n",
    "            new_time = time.time()\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",
    "            print(f'Execution time: {new_time - old_time}')\n",
    "            old_time = time.time()\n",
    "\n",
    "        \n",
    "    acc = 100.0 * correct_samples / total_samples\n",
    "    print(f'Accuracy: ' + '{:4.2f}'.format(acc) + '%')\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "51a6ac55",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def evaluate(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",
    "    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",
    "\n",
    "            _, pred = torch.max(output, dim=1)\n",
    "                        \n",
    "            total_loss += loss.item()\n",
    "            correct_samples += pred.eq(target).sum()\n",
    "            # LR Scheduler is called per epoch\n",
    "            scheduler.step()\n",
    "\n",
    "            \n",
    "    avg_loss = total_loss / total_samples\n",
    "    acc = 100.0 * correct_samples / total_samples\n",
    "    loss_history.append(avg_loss)\n",
    "    \n",
    "    target = target.cpu().detach().numpy()\n",
    "    pred = pred.cpu().detach().numpy()\n",
    "    f1_score = metrics.f1_score(target, pred, average='micro')\n",
    "    \n",
    "    wandb.log({\n",
    "        'test_loss': loss.item(),\n",
    "        'test_f1_score': f1_score\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(acc) + '%)\\n')\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dbe6de3",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "N_EPOCHS = 20\n",
    "\n",
    "Path(\"models/\").mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "## CIFAR 10 model hyperparams\n",
    "# model = GraphViT(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",
    "## MNIST model hyperparams\n",
    "model = GraphViT(image_size=28, patch_size=7, num_classes=10, channels=1,\n",
    "            dim=64, depth=6, heads=8, mlp_dim=128)\n",
    "optimizer = optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.05)\n",
    "early_stop_tolerance = 10e-4\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",
    "model = model.to(device)\n",
    "\n",
    "\n",
    "\n",
    "train_loss_history, test_loss_history = [], []\n",
    "best_model_acc = -1\n",
    "for epoch in range(1, N_EPOCHS + 1):\n",
    "    print('Epoch:', epoch)\n",
    "    train_epoch(model, optimizer, train_loader, train_loss_history, scheduler, best_model_acc)\n",
    "    acc = evaluate(model, test_loader, test_loss_history)\n",
    "    if acc > best_model_acc:\n",
    "      best_model_acc = acc\n",
    "      torch.save(model.state_dict(), os.path.join('models', f'train-4-{epoch}-acc{acc}'))\n",
    "    if len(train_loss_history) > 2 and np.isclose(train_loss_history[-2], train_loss_history[-1], atol=early_stop_tolerance):\n",
    "      print(f\"Early stop with tolerance {early_stop_tolerance} for losses {train_loss_history[-2]} and {train_loss_history[-1]}\")\n",
    "\n",
    "print('Execution time:', '{:5.2f}'.format(time.time() - start_time), 'seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94c7b6b4",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}