ViViT_UCF101.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.8.1"
    },
    "colab": {
      "name": "ViViT_UCF101.ipynb",
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q_b_ifobbjU7"
      },
      "source": [
        "#Code for ViViT model [[paper]](https://arxiv.org/abs/2103.15691)\n"
      ],
      "id": "q_b_ifobbjU7"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rkt-0enhb0qG"
      },
      "source": [
        "## Imports and Global declarations"
      ],
      "id": "rkt-0enhb0qG"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "angry-effectiveness"
      },
      "source": [
        "# Imports\n",
        "\n",
        "import torch\n",
        "from torch import nn, einsum\n",
        "from torch.nn import functional as F\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "import torchvision as tv\n",
        "from torch.utils.data import random_split, DataLoader,Dataset\n",
        "\n",
        "import os\n",
        "import time\n",
        "import random\n",
        "import math\n",
        "\n",
        "# !pip install einops icecream\n",
        "import decord\n",
        "import numpy as np\n",
        "from einops import rearrange, repeat\n",
        "from einops.layers.torch import Rearrange\n",
        "from PIL import Image\n",
        "from tqdm.notebook import tqdm\n",
        "from icecream import ic\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline"
      ],
      "id": "angry-effectiveness",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8ifOyxNa5Cnp"
      },
      "source": [
        "# set device\n",
        "device ='cuda:0' if torch.cuda.is_available() else 'cpu'\n",
        "device"
      ],
      "id": "8ifOyxNa5Cnp",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C9gk3l2-rSt1"
      },
      "source": [
        "# Instantiate tensorboard writer\n",
        "tb_writer = SummaryWriter()"
      ],
      "id": "C9gk3l2-rSt1",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDjl36ascCbH"
      },
      "source": [
        "## DataLoader for UCF101 dataset "
      ],
      "id": "yDjl36ascCbH"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rational-neutral"
      },
      "source": [
        "# dataset params\n",
        "frames_per_clip = 16\n",
        "dataset_dir=\"../../datasets/UCF-101\""
      ],
      "id": "rational-neutral",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "documentary-python"
      },
      "source": [
        "# Dataset Class\n",
        "class UCFDataset(torch.utils.data.Dataset):\n",
        "    \"\"\"\n",
        "    Dataset Class for reading UCF101 dataset  \n",
        "    \n",
        "    Args:\n",
        "        dataset_dir: (str) - root directory of dataset\n",
        "        subset: (str) - train or test subset\n",
        "        video_list_file: (str) - file name containing list of video names \n",
        "        frames_per_clip: (int) - number of frames to be read in every video clip [default:16]\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dataset_dir, subset, video_list_file, frames_per_clip=16):\n",
        "        super().__init__()\n",
        "        self.dataset_dir = dataset_dir\n",
        "        self.subset=subset\n",
        "        self.video_list_file = video_list_file\n",
        "\n",
        "        with open(f'{dataset_dir}\\\\{video_list_file}') as video_names_file:\n",
        "            if self.subset==\"train\":\n",
        "                self.video_list,self.labels = zip(*(files[:-1].split() for files in video_names_file.readlines()))\n",
        "            else:\n",
        "                self.video_list = [files[:-1] for files in video_names_file.readlines()]\n",
        "                with open(f'{dataset_dir}\\\\classInd.txt') as classIndices:\n",
        "                    values,keys=zip(*(files[:-1].split() for files in classIndices.readlines()))\n",
        "                    self.indices = dict( (k,v) for k,v in zip(keys,values))\n",
        "\n",
        "        self.frames_per_clip = frames_per_clip\n",
        "\n",
        "        self.transform = tv.transforms.Compose([\n",
        "          tv.transforms.Resize(256),\n",
        "          tv.transforms.CenterCrop(224),\n",
        "          tv.transforms.ToTensor(),\n",
        "          tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
        "        ])\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.video_list)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        videoname = f'video_data\\\\{self.video_list[idx]}'\n",
        "        vid = decord.VideoReader(f'{self.dataset_dir}\\\\{videoname}', ctx=decord.cpu(0))\n",
        "        nframes = len(vid)\n",
        "\n",
        "        # if number of frames of video is less than frames_per_clip, repeat the frames\n",
        "        if nframes <= self.frames_per_clip:\n",
        "            idxs = np.arange(0, self.frames_per_clip).astype(np.int32)\n",
        "            idxs[nframes:] %= nframes\n",
        "\n",
        "        # else if frames_per_clip is greater, sample uniformly seperated frames\n",
        "        else:\n",
        "            idxs = np.linspace(0, nframes-1, self.frames_per_clip)\n",
        "            idxs = np.round(idxs).astype(np.int32)\n",
        "\n",
        "        imgs = []\n",
        "        for k in idxs:\n",
        "            frame = Image.fromarray(vid[k].asnumpy())\n",
        "            frame = self.transform(frame)\n",
        "            imgs.append(frame)\n",
        "        imgs = torch.stack(imgs)\n",
        "\n",
        "        # if its train subset, return both the frames and the label \n",
        "        if self.subset==\"train\":\n",
        "            label = int(self.labels[idx]) - 1    \n",
        "        # else, for test subset, read the label index\n",
        "        else:\n",
        "            label=int(classIndices[videoname.split('/')[0]])\n",
        "        return imgs,label"
      ],
      "id": "documentary-python",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6zXGDy7QCpHy"
      },
      "source": [
        "#Instantiate and create train-val-test split\n",
        "\n",
        "train_val_data = UCFDataset( dataset_dir = dataset_dir, subset=\"train\", video_list_file=\"trainlist01.txt\",frames_per_clip=frames_per_clip)\n",
        "\n",
        "train_len=int(0.85*len(train_val_data))\n",
        "train_val_split = [ train_len, len(train_val_data) - train_len ] \n",
        "\n",
        "train_data , val_data = random_split(train_val_data,train_val_split)\n",
        "test_data = UCFDataset( dataset_dir = dataset_dir, subset=\"test\", video_list_file=\"testlist01.txt\" ,frames_per_clip=frames_per_clip)\n",
        "\n",
        "print(f\"Train samples: {len(train_data)}\")\n",
        "print(f\"Validation samples: {len(val_data)}\")\n",
        "print(f\"Test samples: {len(test_data)}\")"
      ],
      "id": "6zXGDy7QCpHy",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "rotary-ladder"
      },
      "source": [
        "#example frame from a random video to check if dataloader's working properly\n",
        "plt.imshow(((train_data[35][0][4].permute(1,2,0)*255).int()).numpy())"
      ],
      "id": "rotary-ladder",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "adapted-aviation"
      },
      "source": [
        "# data loading params\n",
        "batch_size = 2\n",
        "test_batch_size = 1\n",
        "num_workers = 2\n",
        "pin_memory = True\n",
        "num_classes=101"
      ],
      "id": "adapted-aviation",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "accepting-payment"
      },
      "source": [
        "# Dataloaders\n",
        "train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
        "val_loader = DataLoader(val_data, batch_size=test_batch_size)\n",
        "test_loader = DataLoader(test_data, batch_size=test_batch_size)"
      ],
      "id": "accepting-payment",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YqviNYNCLiW4"
      },
      "source": [
        "## Model definition"
      ],
      "id": "YqviNYNCLiW4"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "comparable-crack"
      },
      "source": [
        "class MLP(nn.Module):\n",
        "    \"\"\"\n",
        "    Builds a simple feed forward network\n",
        "    \n",
        "    Args:\n",
        "        dim: (int) - inner dimension of embeddings \n",
        "        inner_dim: (int) - dimension of transformer head  \n",
        "    \n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim, inner_dim):\n",
        "        super().__init__()\n",
        "        # mlp with GELU activation function\n",
        "        self.mlp = nn.Sequential(\n",
        "            nn.Linear(dim, inner_dim),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(inner_dim, inner_dim),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(inner_dim, dim),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.mlp(x)"
      ],
      "id": "comparable-crack",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "outer-stack"
      },
      "source": [
        "class Attention(nn.Module):\n",
        "    \"\"\"\n",
        "    Builds a Self Attention Module, however, could be modified to make other attention modules too.\n",
        "    \n",
        "    Args:\n",
        "        dim: (int) - inner dimension of embeddings[default:192] \n",
        "        heads: (int) - number of attention heads[default:8] \n",
        "        dim_head: (int) - dimension of transformer head [default:64] \n",
        "    \n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim = 192, heads = 8, dim_head = 64):\n",
        "        super().__init__()\n",
        "        inner_dim = dim_head *  heads\n",
        "\n",
        "        self.heads = heads\n",
        "        self.scale = dim_head ** -0.5\n",
        "\n",
        "        # nn.Linear from 192 to (8*64)*3\n",
        "        self.make_qkv = nn.Linear(dim, inner_dim * 3, bias = False)\n",
        "\n",
        "        # Linear projection to required output dimension\n",
        "        self.get_output = nn.Sequential(nn.Linear(inner_dim, dim))\n",
        "        \n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        b, n, _ = x.shape   # b=batch_size , n=197 \n",
        "        h = self.heads      # h=8\n",
        "\n",
        "        # nn.Linear from 192 to 256*3 & then split it across q,k & v each with last dimension as 256\n",
        "        qkv = self.make_qkv(x).chunk(3, dim = -1)\n",
        "        \n",
        "        # reshaping to get the right q,k,v dimensions having 8 attn_heads(h)\n",
        "        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)\n",
        "\n",
        "        # dot product of q & k after transposing k followed by a softmax layer\n",
        "        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale     # q.kT / sqrt(d)\n",
        "        attn = dots.softmax(dim=-1)\n",
        "\n",
        "        # dot product of attention layer with v\n",
        "        output = einsum('b h i j, b h j d -> b h i d', attn, v)\n",
        "\n",
        "        # Final reshaping & nn.Linear to combine all attention head outputs to get final out.\n",
        "        output = rearrange(output, 'b h n d -> b n (h d)')\n",
        "        \n",
        "        output =  self.get_output(output)    \n",
        "        # output shape = ( b, n, dim (=192) )\n",
        "\n",
        "        return output  "
      ],
      "id": "outer-stack",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "variable-harvest"
      },
      "source": [
        "class Transformer(nn.Module):\n",
        "    \"\"\"\n",
        "    Builds a Transformer Model. Could be used as a Spatial or Temporal transformer.\n",
        "    \n",
        "    Args:\n",
        "        dim: (int) - inner dimension of embeddings \n",
        "        depth: (int) - depth of the transformer \n",
        "        heads: (int) - number of attention heads [default:8] \n",
        "        dim_head: (int) - dimension of transformer head [default:64] \n",
        "        mlp_dim: (int) - scaling dimension for attention [default:768] \n",
        "    \n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, dim, depth, heads=8, dim_head=64, mlp_dim=768):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.model_layers = nn.ModuleList([])\n",
        "        for i in range(depth):\n",
        "            self.model_layers.append(nn.ModuleList([\n",
        "                nn.LayerNorm(dim),\n",
        "                Attention(dim, heads, dim_head),\n",
        "                nn.LayerNorm(dim),\n",
        "                MLP(dim, mlp_dim)\n",
        "            ]))\n",
        "\n",
        "        self.layer_norm = nn.LayerNorm(dim)\n",
        "    \n",
        "    def forward(self, x):\n",
        "\n",
        "        for layer_norm1, attention, layer_norm2, ff_net in self.model_layers:\n",
        "            x = attention(layer_norm1(x)) + x\n",
        "            x = ff_net(layer_norm2(x)) + x\n",
        "\n",
        "        return self.layer_norm(x)"
      ],
      "id": "variable-harvest",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E-YBjPM2YN0U"
      },
      "source": [
        "Referred from [this](https://github.com/rishikksh20/ViViT-pytorc) unofficial implementation\n"
      ],
      "id": "E-YBjPM2YN0U"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "afraid-train"
      },
      "source": [
        "class ViViT(nn.Module):\n",
        "    \"\"\"\n",
        "    Builds a ViViT Model for video action recognition\n",
        "    The line-by-line comments are for default values of arguments  \n",
        "    \n",
        "    Args:\n",
        "        image_size: (int) - size of input image\n",
        "        patch_size: (int) - size of each patch of the image \n",
        "        num_classes: (int) - number of classes in the dataset\n",
        "        frames_per_clip: (int) - number of frames in every video clip. [default:16] \n",
        "        dim: (int) - inner dimension of embeddings[default:192] \n",
        "        depth: (int) - depth of the transformer[default:4] \n",
        "        heads: (int) - number of attention heads for the transformer[default:8] \n",
        "        pooling: (str) - type of pooling[default:'mean'] \n",
        "        in_channels: (int) - number of input channels for each frame [default:3] \n",
        "        dim_head: (int) - dimension of transformer head [default:64] \n",
        "        scale_dim: (int) - scaling dimension for attention [default:4] \n",
        "    \n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, image_size, patch_size, num_classes, frames_per_clip=16, dim = 192, depth = 4, heads = 8, pooling = 'mean', in_channels = 3, dim_head = 64, scale_dim = 4, ):\n",
        "        \n",
        "        super().__init__()\n",
        "\n",
        "        num_patches = (image_size // patch_size) ** 2   # => 196 for 224x224 images\n",
        "        patch_dim = in_channels * patch_size ** 2      # => 3*16*16\n",
        "\n",
        "        self.get_patch_emb = nn.Sequential(\n",
        "            # input h = 14, w=14, c=3, p1=16, p2=16\n",
        "            # reshape from (batch_size, frames, channels, 224, 224) to  (batch_size, frames, 14*14, 16*16*3 )\n",
        "            Rearrange('b t c (h p1) (w p2) -> b t (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),  \n",
        "\n",
        "            # fully connected from 16*16*3 to 192 for every patch in (batch_size, frames) \n",
        "            nn.Linear(patch_dim, dim),\n",
        "        )\n",
        "\n",
        "        # position embeddings of shape: (1, frames_per_clip = 16, num_patches + 1 = 197, 192)\n",
        "        self.pos_embedding = nn.Parameter(torch.randn(1, frames_per_clip, num_patches + 1, dim))\n",
        "\n",
        "        # space (i.e. for each image) tokens of shape: (1, 1, 192). The 192 is the tokens obtained in \"get_patch_emb\" \n",
        "        self.spatial_token = nn.Parameter(torch.randn(1, 1, dim))\n",
        "        \n",
        "        # spatial transformer ViT\n",
        "        self.spatial_transformer = Transformer(dim, depth, heads, dim_head, dim*scale_dim)\n",
        "\n",
        "        # time dimention tokens of shape: (1, 1, 192). \n",
        "        self.temporal_token = nn.Parameter(torch.randn(1, 1, dim))\n",
        "        \n",
        "        # temporal transformer which takes in spacetransformer's output tokens as the input. \n",
        "        self.temporal_transformer = Transformer(dim, depth, heads, dim_head, dim*scale_dim)\n",
        "\n",
        "        # pooling type, could be \"mean\" or \"cls\"\n",
        "        self.pooling = pooling\n",
        "\n",
        "        # mlp head for final classification\n",
        "        self.classifier_head = nn.Sequential(\n",
        "            nn.LayerNorm(dim),\n",
        "            nn.Linear(dim, num_classes),\n",
        "            nn.Softmax(dim=1)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "\n",
        "        # get patch embeddings\n",
        "        x = self.get_patch_emb(x)\n",
        "\n",
        "        # b = batch_size , t = frames , n = number of patch embeddings= 14*14 , e = embedding size\n",
        "        b, t, n, e = x.shape     # x.shape = (b, t, 196, 192) \n",
        "\n",
        "        # prepare cls_token for space transformers\n",
        "        spatial_cls_tokens = repeat(self.spatial_token, '() n d -> b t n d', b = b, t=t)\n",
        "\n",
        "        # concatenate cls_token to the patch embedding\n",
        "        x = torch.cat((spatial_cls_tokens, x), dim=2)     # => x shape = ( b, t, 197 ,192)\n",
        "\n",
        "        # add position embedding info \n",
        "        x += self.pos_embedding[:, :, :(n + 1)]\n",
        "\n",
        "        # club together the b & t dimension \n",
        "        x = rearrange(x, 'b t n d -> (b t) n d')\n",
        "\n",
        "        # pass through spatial transformer\n",
        "        x = self.spatial_transformer(x)\n",
        "\n",
        "        # declub b & t dimensions\n",
        "        x = rearrange(x[:, 0], '(b t) ... -> b t ...', b=b)\n",
        "\n",
        "        # prepare cls_token for temporal transformers & concatenate cls_token to the patch embedding\n",
        "        temporal_cls_tokens = repeat(self.temporal_token, '() n d -> b n d', b=b)\n",
        "        x = torch.cat((temporal_cls_tokens, x), dim=1)\n",
        "\n",
        "        # pass through spatial transformer\n",
        "        x = self.temporal_transformer(x)\n",
        "        \n",
        "        # if pooling is mean, then use mean of all 197 as the output, else use output corresponding to cls token as the final x\n",
        "        if self.pooling == 'mean':\n",
        "            x = x.mean(dim = 1) #( b, n, dim (=192) )\n",
        "        else:\n",
        "             x[:, 0] #( b, n, dim (=192) )\n",
        "\n",
        "        # pass through MLP classification layer\n",
        "        return self.classifier_head(x)"
      ],
      "id": "afraid-train",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "economic-limitation"
      },
      "source": [
        "# instantiate model and plot on tensorboard\n",
        "model = ViViT(image_size=224, patch_size=16, num_classes=num_classes, frames_per_clip=frames_per_clip)\n",
        "frames, _ = next(iter(train_loader))\n",
        "tb_writer.add_graph(model, frames)\n",
        "model.to(device)"
      ],
      "id": "economic-limitation",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HiyjllqxL9bV"
      },
      "source": [
        "## Training utils "
      ],
      "id": "HiyjllqxL9bV"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rY9aVH4pZM4D"
      },
      "source": [
        "# training hyper-params\n",
        "lr=0.0001\n",
        "epochs = 20 "
      ],
      "id": "rY9aVH4pZM4D",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v3JiODtrd0j-"
      },
      "source": [
        "# define the loss and optimizers\n",
        "loss_criterion = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(),lr)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.95)"
      ],
      "id": "v3JiODtrd0j-",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "least-visitor"
      },
      "source": [
        "# training step for every epoch\n",
        "def train_step(loader,epoch,):\n",
        "    \n",
        "    model.train()\n",
        "    total_epoch_loss=0\n",
        "    \n",
        "    for batch_id, (video_data,labels) in enumerate(loader):\n",
        "\n",
        "        # video_data,labels = video_data.to(device), labels.to(device)\n",
        "        video_data,labels = video_data.to(device), labels.to(device)\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        prediction = model(video_data)\n",
        "\n",
        "        loss = loss_criterion(prediction,labels)\n",
        "        total_epoch_loss += loss.item()\n",
        "\n",
        "        loss.backward()\n",
        "        \n",
        "        optimizer.step()\n",
        "\n",
        "        tb_writer.add_scalar(\"Train/Loss\",loss.item(),((len(loader))*(epoch-1))+batch_id)\n",
        " \n",
        "        print(f\"\\n[Epoch]: {epoch} Train Loss: {loss.item()}\")\n",
        "\n",
        "    return total_epoch_loss"
      ],
      "id": "least-visitor",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "elegant-tulsa"
      },
      "source": [
        "# validation step for every epoch\n",
        "def val_step(loader,epoch=None):\n",
        "\n",
        "    model.eval()\n",
        "    total_loss=0\n",
        "    corrects=0\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        for batch_id, (video_data,labels) in enumerate(loader):\n",
        "\n",
        "            video_data,labels = (video_data).to(device), labels.to(device)\n",
        "\n",
        "            prediction = model(video_data)\n",
        "            \n",
        "            loss = loss_criterion(prediction,labels)\n",
        "            total_loss += loss.item()\n",
        "            corrects+= (torch.argmax(prediction,dim=1)==labels).sum()\n",
        "    \n",
        "    accuracy = corrects/(len(loader)*batch_size)\n",
        "    \n",
        "    print(f\"\\n[Epoch]: {epoch} , Accuracy: {accuracy}, Valid Loss: {loss.item()}\")\n",
        "    tb_writer.add_scalar(\"Validation/Loss\",loss.item(),epoch)\n",
        "    tb_writer.add_scalar(\"Validation/Accuracy\",accuracy,epoch)\n",
        "\n",
        "    return accuracy"
      ],
      "id": "elegant-tulsa",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MHl6bllSMGJX"
      },
      "source": [
        "## Train Loop"
      ],
      "id": "MHl6bllSMGJX"
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "worst-bedroom"
      },
      "source": [
        "# Driving train test loop\n",
        "for epoch in tqdm(range(1,epochs+1)):\n",
        "    train_step(train_loader, epoch)\n",
        "    val_step(val_loader, epoch)\n",
        "    scheduler.step()\n",
        "    torch.save(model,\"ViViT_model.pt\")"
      ],
      "id": "worst-bedroom",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4ptG3LVMKXX"
      },
      "source": [
        "## Testing"
      ],
      "id": "r4ptG3LVMKXX"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GHHJtUrmkjb1"
      },
      "source": [
        "# test the trained model  \n",
        "def test_model(loader):\n",
        "\n",
        "    model.eval()\n",
        "    corrects=0\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        for batch_id, (input_data) in enumerate(loader):\n",
        "\n",
        "            input_data = (input_data).to(device), labels.to(device)\n",
        "\n",
        "            prediction = model(input_data)\n",
        "            loss = loss_criterion(prediction,labels)\n",
        "            total_loss += loss.item()\n",
        "            corrects+= (torch.argmax(prediction,dim=1)==labels).sum()\n",
        "    \n",
        "    accuracy = corrects/(len(loader)*batch_size)\n",
        "    print(f\"Accuracy: {accuracy}, Test Loss: {loss.item()}\")\n",
        "\n",
        "    return accuracy"
      ],
      "id": "GHHJtUrmkjb1",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "proprietary-shell"
      },
      "source": [
        "test_model(model, test_loader)"
      ],
      "id": "proprietary-shell",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BAsvl78-MO_D"
      },
      "source": [
        "## Visualize results on tensorboard"
      ],
      "id": "BAsvl78-MO_D"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "increasing-delaware"
      },
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir=runs/ "
      ],
      "id": "increasing-delaware",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cqqvXqp4G2uI"
      },
      "source": [
        ""
      ],
      "id": "cqqvXqp4G2uI",
      "execution_count": null,
      "outputs": []
    }
  ]
}