Video Inference with Train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j540ho0jlEq1"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "Videos are sequences of images. Let's assume you have an image\n",
    "representation model (CNN, ViT, etc.) and a sequence model\n",
    "(RNN, LSTM, etc.) at hand. We ask you to tweak the model for video\n",
    "classification. The simplest approach would be to apply the image\n",
    "model to individual frames, use the sequence model to learn\n",
    "sequences of image features, then apply a classification head on\n",
    "the learned sequence representation.\n",
    "(https://keras.io/examples/vision/video_transformers/).\n",
    "\n",
    "In this example, we minimally implement\n",
    "[ViViT: A Video Vision Transformer](https://arxiv.org/abs/2103.15691)\n",
    "by Arnab et al., a **pure Transformer-based** model\n",
    "for video classification. The authors propose a novel embedding scheme\n",
    "and a number of Transformer variants to model video clips. We implement\n",
    "the embedding scheme and one of the variants of the Transformer\n",
    "architecture, for simplicity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_yuPg2KNlEq3"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "fPyUO9chlEq4"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-09-09 05:26:46.958413: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /glob/development-tools/versions/oneapi/2022.2/oneapi/vpl/2022.1.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/tbb/2021.6.0/env/../lib/intel64/gcc4.8:/glob/development-tools/versions/oneapi/2022.2/oneapi/rkcommon/1.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray_studio/0.10.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray/2.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/openvkl/1.2.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/oidn/1.4.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//libfabric/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib/release:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mkl/2022.1.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/itac/2021.6.0/slib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ispc/1.17.0/lib/lib64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ippcp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/embree/3.13.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dnnl/2022.1.0/cpu_dpcpp_gpu_dpcpp/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/gdb/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/libipt/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/dep/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dal/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/x64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/oclfpga/host/linux64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/compiler/lib/intel64_lin:/glob/development-tools/versions/oneapi/2022.2/oneapi/ccl/2021.6.0/lib/cpu_gpu_dpcpp:/glob/development-tools/versions/oneapi/2022.2/oneapi/vpl/2022.1.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/tbb/2021.6.0/env/../lib/intel64/gcc4.8:/glob/development-tools/versions/oneapi/2022.2/oneapi/rkcommon/1.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray_studio/0.10.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray/2.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/openvkl/1.2.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/oidn/1.4.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//libfabric/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib/release:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mkl/2022.1.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/itac/2021.6.0/slib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ispc/1.17.0/lib/lib64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ippcp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/embree/3.13.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dnnl/2022.1.0/cpu_dpcpp_gpu_dpcpp/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/gdb/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/libipt/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/dep/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dal/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/x64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/oclfpga/host/linux64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/compiler/lib/intel64_lin:/glob/development-tools/versions/oneapi/2022.2/oneapi/ccl/2021.6.0/lib/cpu_gpu_dpcpp\n",
      "2022-09-09 05:26:46.958523: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import io\n",
    "import imageio\n",
    "import medmnist\n",
    "import ipywidgets\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "# Setting seed for reproducibility\n",
    "SEED = 42\n",
    "keras.utils.set_random_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SSby1Z3xlEq5"
   },
   "source": [
    "## Hyperparameters\n",
    "\n",
    "The hyperparameters are chosen via hyperparameter\n",
    "search. You can learn more about the process in the \"conclusion\" section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "dy3Pfo1YlEq6"
   },
   "outputs": [],
   "source": [
    "# DATA\n",
    "DATASET_NAME = \"organmnist3d\"\n",
    "BATCH_SIZE = 32\n",
    "AUTO = tf.data.AUTOTUNE\n",
    "INPUT_SHAPE = (28, 28, 28, 1)\n",
    "NUM_CLASSES = 11\n",
    "\n",
    "# OPTIMIZER\n",
    "LEARNING_RATE = 1e-4\n",
    "WEIGHT_DECAY = 1e-5\n",
    "\n",
    "# TRAINING\n",
    "EPOCHS = 60\n",
    "\n",
    "# TUBELET EMBEDDING\n",
    "PATCH_SIZE = (8, 8, 8)\n",
    "NUM_PATCHES = (INPUT_SHAPE[0] // PATCH_SIZE[0]) ** 2\n",
    "\n",
    "# ViViT ARCHITECTURE\n",
    "LAYER_NORM_EPS = 1e-6\n",
    "PROJECTION_DIM = 128\n",
    "NUM_HEADS = 8\n",
    "NUM_LAYERS = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LEUs6r5HlEq6"
   },
   "source": [
    "## Dataset\n",
    "\n",
    "For our example we use the\n",
    "[MedMNIST v2: A Large-Scale Lightweight Benchmark for 2D and 3D Biomedical Image Classification](https://medmnist.com/)\n",
    "dataset. The videos are lightweight and easy to train on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "TdW7QTsclEq7"
   },
   "outputs": [],
   "source": [
    "\n",
    "def download_and_prepare_dataset(data_info: dict):\n",
    "    \"\"\"Utility function to download the dataset.\n",
    "\n",
    "    Arguments:\n",
    "        data_info (dict): Dataset metadata.\n",
    "    \"\"\"\n",
    "    data_path = keras.utils.get_file(origin=data_info[\"url\"], md5_hash=data_info[\"MD5\"])\n",
    "\n",
    "    with np.load(data_path) as data:\n",
    "        # Get videos\n",
    "        train_videos = data[\"train_images\"]\n",
    "        valid_videos = data[\"val_images\"]\n",
    "        test_videos = data[\"test_images\"]\n",
    "\n",
    "        # Get labels\n",
    "        train_labels = data[\"train_labels\"].flatten()\n",
    "        valid_labels = data[\"val_labels\"].flatten()\n",
    "        test_labels = data[\"test_labels\"].flatten()\n",
    "\n",
    "    return (\n",
    "        (train_videos, train_labels),\n",
    "        (valid_videos, valid_labels),\n",
    "        (test_videos, test_labels),\n",
    "    )\n",
    "\n",
    "\n",
    "# Get the metadata of the dataset\n",
    "info = medmnist.INFO[DATASET_NAME]\n",
    "\n",
    "# Get the dataset\n",
    "prepared_dataset = download_and_prepare_dataset(info)\n",
    "(train_videos, train_labels) = prepared_dataset[0]\n",
    "(valid_videos, valid_labels) = prepared_dataset[1]\n",
    "(test_videos, test_labels) = prepared_dataset[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tiJqUB5-lEq9"
   },
   "source": [
    "### `tf.data` pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "3RyKfv1clEq9"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-09-09 05:27:13.391307: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /glob/development-tools/versions/oneapi/2022.2/oneapi/vpl/2022.1.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/tbb/2021.6.0/env/../lib/intel64/gcc4.8:/glob/development-tools/versions/oneapi/2022.2/oneapi/rkcommon/1.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray_studio/0.10.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray/2.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/openvkl/1.2.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/oidn/1.4.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//libfabric/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib/release:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mkl/2022.1.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/itac/2021.6.0/slib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ispc/1.17.0/lib/lib64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ippcp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/embree/3.13.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dnnl/2022.1.0/cpu_dpcpp_gpu_dpcpp/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/gdb/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/libipt/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/dep/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dal/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/x64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/oclfpga/host/linux64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/compiler/lib/intel64_lin:/glob/development-tools/versions/oneapi/2022.2/oneapi/ccl/2021.6.0/lib/cpu_gpu_dpcpp:/glob/development-tools/versions/oneapi/2022.2/oneapi/vpl/2022.1.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/tbb/2021.6.0/env/../lib/intel64/gcc4.8:/glob/development-tools/versions/oneapi/2022.2/oneapi/rkcommon/1.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray_studio/0.10.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ospray/2.9.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/openvkl/1.2.0/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/oidn/1.4.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//libfabric/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib/release:/glob/development-tools/versions/oneapi/2022.2/oneapi/mpi/2021.6.0//lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/mkl/2022.1.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/itac/2021.6.0/slib:/glob/development-tools/versions/oneapi/2022.2/oneapi/ispc/1.17.0/lib/lib64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ippcp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/ipp/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/embree/3.13.3/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dnnl/2022.1.0/cpu_dpcpp_gpu_dpcpp/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/gdb/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/libipt/intel64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/debugger/2021.6.0/dep/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/dal/2021.6.0/lib/intel64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/x64:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/lib/oclfpga/host/linux64/lib:/glob/development-tools/versions/oneapi/2022.2/oneapi/compiler/2022.1.0/linux/compiler/lib/intel64_lin:/glob/development-tools/versions/oneapi/2022.2/oneapi/ccl/2021.6.0/lib/cpu_gpu_dpcpp\n",
      "2022-09-09 05:27:13.391402: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)\n",
      "2022-09-09 05:27:13.391456: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (s001-n007): /proc/driver/nvidia/version does not exist\n",
      "2022-09-09 05:27:13.438531: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "@tf.function\n",
    "def preprocess(frames: tf.Tensor, label: tf.Tensor):\n",
    "    \"\"\"Preprocess the frames tensors and parse the labels.\"\"\"\n",
    "    # Preprocess images\n",
    "    frames = tf.image.convert_image_dtype(\n",
    "        frames[\n",
    "            ..., tf.newaxis\n",
    "        ],  # The new axis is to help for further processing with Conv3D layers\n",
    "        tf.float32,\n",
    "    )\n",
    "    # Parse label\n",
    "    label = tf.cast(label, tf.float32)\n",
    "    return frames, label\n",
    "\n",
    "\n",
    "def prepare_dataloader(\n",
    "    videos: np.ndarray,\n",
    "    labels: np.ndarray,\n",
    "    loader_type: str = \"train\",\n",
    "    batch_size: int = BATCH_SIZE,\n",
    "):\n",
    "    \"\"\"Utility function to prepare the dataloader.\"\"\"\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((videos, labels))\n",
    "\n",
    "    if loader_type == \"train\":\n",
    "        dataset = dataset.shuffle(BATCH_SIZE * 2)\n",
    "\n",
    "    dataloader = (\n",
    "        dataset.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)\n",
    "        .batch(batch_size)\n",
    "        .prefetch(tf.data.AUTOTUNE)\n",
    "    )\n",
    "    return dataloader\n",
    "\n",
    "\n",
    "trainloader = prepare_dataloader(train_videos, train_labels, \"train\")\n",
    "validloader = prepare_dataloader(valid_videos, valid_labels, \"valid\")\n",
    "testloader = prepare_dataloader(test_videos, test_labels, \"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JatUOaWKlEq-"
   },
   "source": [
    "## Tubelet Embedding\n",
    "\n",
    "In ViTs, an image is divided into patches, which are then spatially\n",
    "flattened, a process known as tokenization. For a video, one can\n",
    "repeat this process for individual frames. **Uniform frame sampling**\n",
    "as suggested by the authors is a tokenization scheme in which we\n",
    "sample frames from the video clip and perform simple ViT tokenization.\n",
    "\n",
    "| ![uniform frame sampling](https://i.imgur.com/aaPyLPX.png) |\n",
    "| :--: |\n",
    "| Uniform Frame Sampling [Source](https://arxiv.org/abs/2103.15691) |\n",
    "\n",
    "**Tubelet Embedding** is different in terms of capturing temporal\n",
    "information from the video.\n",
    "First, we extract volumes from the video -- these volumes contain\n",
    "patches of the frame and the temporal information as well. The volumes\n",
    "are then flattened to build video tokens.\n",
    "\n",
    "| ![tubelet embedding](https://i.imgur.com/9G7QTfV.png) |\n",
    "| :--: |\n",
    "| Tubelet Embedding [Source](https://arxiv.org/abs/2103.15691) |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "OO1D7QAmlEq_"
   },
   "outputs": [],
   "source": [
    "class TubeletEmbedding(layers.Layer):\n",
    "    def __init__(self, embed_dim, patch_size, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.projection = layers.Conv3D(\n",
    "            filters=embed_dim,\n",
    "            kernel_size=patch_size,\n",
    "            strides=patch_size,\n",
    "            padding=\"VALID\",\n",
    "        )\n",
    "        self.flatten = layers.Reshape(target_shape=(-1, embed_dim))\n",
    "\n",
    "    def call(self, videos):\n",
    "        projected_patches = self.projection(videos)\n",
    "        flattened_patches = self.flatten(projected_patches)\n",
    "        return flattened_patches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yTfSBwqSlEq_"
   },
   "source": [
    "## Positional Embedding\n",
    "\n",
    "This layer adds positional information to the encoded video tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "GtItEC2rlErA"
   },
   "outputs": [],
   "source": [
    "class PositionalEncoder(layers.Layer):\n",
    "    def __init__(self, embed_dim, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.embed_dim = embed_dim\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        _, num_tokens, _ = input_shape\n",
    "        self.position_embedding = layers.Embedding(\n",
    "            input_dim=num_tokens, output_dim=self.embed_dim\n",
    "        )\n",
    "        self.positions = tf.range(start=0, limit=num_tokens, delta=1)\n",
    "\n",
    "    def call(self, encoded_tokens):\n",
    "        # Encode the positions and add it to the encoded tokens\n",
    "        encoded_positions = self.position_embedding(self.positions)\n",
    "        encoded_tokens = encoded_tokens + encoded_positions\n",
    "        return encoded_tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VfvTMxJ4lErA"
   },
   "source": [
    "## Video Vision Transformer\n",
    "\n",
    "The authors suggest 4 variants of Vision Transformer:\n",
    "\n",
    "- Spatio-temporal attention\n",
    "- Factorized encoder\n",
    "- Factorized self-attention\n",
    "- Factorized dot-product attention\n",
    "\n",
    "In this example, we will implement the **Spatio-temporal attention**\n",
    "model for simplicity. The following code snippet is heavily inspired from\n",
    "[Image classification with Vision Transformer](https://keras.io/examples/vision/image_classification_with_vision_transformer/).\n",
    "One can also refer to the\n",
    "[official repository of ViViT](https://github.com/google-research/scenic/tree/main/scenic/projects/vivit)\n",
    "which contains all the variants, implemented in JAX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "1IRS0w3KlErB"
   },
   "outputs": [],
   "source": [
    "def create_vivit_classifier(\n",
    "    tubelet_embedder,\n",
    "    positional_encoder,\n",
    "    input_shape=INPUT_SHAPE,\n",
    "    transformer_layers=NUM_LAYERS,\n",
    "    num_heads=NUM_HEADS,\n",
    "    embed_dim=PROJECTION_DIM,\n",
    "    layer_norm_eps=LAYER_NORM_EPS,\n",
    "    num_classes=NUM_CLASSES,\n",
    "):\n",
    "    # Get the input layer\n",
    "    inputs = layers.Input(shape=input_shape)\n",
    "    # Create patches.\n",
    "    patches = tubelet_embedder(inputs)\n",
    "    # Encode patches.\n",
    "    encoded_patches = positional_encoder(patches)\n",
    "\n",
    "    # Create multiple layers of the Transformer block.\n",
    "    for _ in range(transformer_layers):\n",
    "        # Layer normalization and MHSA\n",
    "        x1 = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
    "        attention_output = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=embed_dim // num_heads, dropout=0.1\n",
    "        )(x1, x1)\n",
    "\n",
    "        # Skip connection\n",
    "        x2 = layers.Add()([attention_output, encoded_patches])\n",
    "\n",
    "        # Layer Normalization and MLP\n",
    "        x3 = layers.LayerNormalization(epsilon=1e-6)(x2)\n",
    "        x3 = keras.Sequential(\n",
    "            [\n",
    "                layers.Dense(units=embed_dim * 4, activation=tf.nn.gelu),\n",
    "                layers.Dense(units=embed_dim, activation=tf.nn.gelu),\n",
    "            ]\n",
    "        )(x3)\n",
    "\n",
    "        # Skip connection\n",
    "        encoded_patches = layers.Add()([x3, x2])\n",
    "\n",
    "    # Layer normalization and Global average pooling.\n",
    "    representation = layers.LayerNormalization(epsilon=layer_norm_eps)(encoded_patches)\n",
    "    representation = layers.GlobalAvgPool1D()(representation)\n",
    "\n",
    "    # Classify outputs.\n",
    "    outputs = layers.Dense(units=num_classes, activation=\"softmax\")(representation)\n",
    "\n",
    "    # Create the Keras model.\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iZANaiirlErB"
   },
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "7eCVyK26lErB"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/60\n",
      "31/31 [==============================] - 14s 203ms/step - loss: 2.4100 - accuracy: 0.1307 - top-5-accuracy: 0.6111 - val_loss: 2.2180 - val_accuracy: 0.2422 - val_top-5-accuracy: 0.7081\n",
      "Epoch 2/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 2.0878 - accuracy: 0.2058 - top-5-accuracy: 0.7531 - val_loss: 1.9416 - val_accuracy: 0.2298 - val_top-5-accuracy: 0.7391\n",
      "Epoch 3/60\n",
      "31/31 [==============================] - 5s 172ms/step - loss: 1.9232 - accuracy: 0.2500 - top-5-accuracy: 0.8282 - val_loss: 1.7017 - val_accuracy: 0.3168 - val_top-5-accuracy: 0.8696\n",
      "Epoch 4/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 1.7617 - accuracy: 0.2932 - top-5-accuracy: 0.8786 - val_loss: 1.5232 - val_accuracy: 0.3416 - val_top-5-accuracy: 0.9379\n",
      "Epoch 5/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 1.6091 - accuracy: 0.3467 - top-5-accuracy: 0.9043 - val_loss: 1.4992 - val_accuracy: 0.3168 - val_top-5-accuracy: 0.9503\n",
      "Epoch 6/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 1.4432 - accuracy: 0.4393 - top-5-accuracy: 0.9239 - val_loss: 1.2170 - val_accuracy: 0.4720 - val_top-5-accuracy: 0.9565\n",
      "Epoch 7/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 1.3383 - accuracy: 0.4805 - top-5-accuracy: 0.9414 - val_loss: 1.1929 - val_accuracy: 0.4969 - val_top-5-accuracy: 0.9752\n",
      "Epoch 8/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 1.3102 - accuracy: 0.4722 - top-5-accuracy: 0.9465 - val_loss: 1.0678 - val_accuracy: 0.5652 - val_top-5-accuracy: 0.9752\n",
      "Epoch 9/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 1.2364 - accuracy: 0.5185 - top-5-accuracy: 0.9496 - val_loss: 0.9949 - val_accuracy: 0.5714 - val_top-5-accuracy: 0.9814\n",
      "Epoch 10/60\n",
      "31/31 [==============================] - 5s 166ms/step - loss: 1.2058 - accuracy: 0.5278 - top-5-accuracy: 0.9527 - val_loss: 0.8689 - val_accuracy: 0.6832 - val_top-5-accuracy: 0.9876\n",
      "Epoch 11/60\n",
      "31/31 [==============================] - 5s 167ms/step - loss: 1.0700 - accuracy: 0.5813 - top-5-accuracy: 0.9619 - val_loss: 0.7665 - val_accuracy: 0.7391 - val_top-5-accuracy: 0.9938\n",
      "Epoch 12/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.9877 - accuracy: 0.6183 - top-5-accuracy: 0.9691 - val_loss: 0.7561 - val_accuracy: 0.6708 - val_top-5-accuracy: 1.0000\n",
      "Epoch 13/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.9471 - accuracy: 0.6440 - top-5-accuracy: 0.9784 - val_loss: 0.8702 - val_accuracy: 0.6832 - val_top-5-accuracy: 0.9752\n",
      "Epoch 14/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.9140 - accuracy: 0.6348 - top-5-accuracy: 0.9815 - val_loss: 0.7462 - val_accuracy: 0.7329 - val_top-5-accuracy: 0.9938\n",
      "Epoch 15/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.8060 - accuracy: 0.6903 - top-5-accuracy: 0.9846 - val_loss: 0.6559 - val_accuracy: 0.7453 - val_top-5-accuracy: 1.0000\n",
      "Epoch 16/60\n",
      "31/31 [==============================] - 5s 164ms/step - loss: 0.7669 - accuracy: 0.7315 - top-5-accuracy: 0.9856 - val_loss: 0.4809 - val_accuracy: 0.8634 - val_top-5-accuracy: 0.9938\n",
      "Epoch 17/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.6272 - accuracy: 0.7840 - top-5-accuracy: 0.9897 - val_loss: 0.5235 - val_accuracy: 0.8199 - val_top-5-accuracy: 1.0000\n",
      "Epoch 18/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.5904 - accuracy: 0.8076 - top-5-accuracy: 0.9938 - val_loss: 0.5101 - val_accuracy: 0.8385 - val_top-5-accuracy: 1.0000\n",
      "Epoch 19/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.5864 - accuracy: 0.7778 - top-5-accuracy: 0.9959 - val_loss: 0.5114 - val_accuracy: 0.8261 - val_top-5-accuracy: 1.0000\n",
      "Epoch 20/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.6273 - accuracy: 0.7819 - top-5-accuracy: 0.9846 - val_loss: 0.5055 - val_accuracy: 0.8509 - val_top-5-accuracy: 1.0000\n",
      "Epoch 21/60\n",
      "31/31 [==============================] - 5s 167ms/step - loss: 0.4806 - accuracy: 0.8374 - top-5-accuracy: 0.9918 - val_loss: 0.3943 - val_accuracy: 0.8944 - val_top-5-accuracy: 0.9938\n",
      "Epoch 22/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.3801 - accuracy: 0.8704 - top-5-accuracy: 0.9959 - val_loss: 0.4661 - val_accuracy: 0.8634 - val_top-5-accuracy: 1.0000\n",
      "Epoch 23/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.3685 - accuracy: 0.8693 - top-5-accuracy: 0.9979 - val_loss: 0.3805 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 24/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.3833 - accuracy: 0.8611 - top-5-accuracy: 0.9990 - val_loss: 0.3609 - val_accuracy: 0.8634 - val_top-5-accuracy: 0.9938\n",
      "Epoch 25/60\n",
      "31/31 [==============================] - 5s 166ms/step - loss: 0.3050 - accuracy: 0.8981 - top-5-accuracy: 0.9990 - val_loss: 0.4151 - val_accuracy: 0.9006 - val_top-5-accuracy: 0.9938\n",
      "Epoch 26/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.2785 - accuracy: 0.9023 - top-5-accuracy: 0.9979 - val_loss: 0.3222 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 27/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.2503 - accuracy: 0.9105 - top-5-accuracy: 0.9990 - val_loss: 0.2723 - val_accuracy: 0.9255 - val_top-5-accuracy: 1.0000\n",
      "Epoch 28/60\n",
      "31/31 [==============================] - 5s 167ms/step - loss: 0.2370 - accuracy: 0.9228 - top-5-accuracy: 1.0000 - val_loss: 0.2787 - val_accuracy: 0.9130 - val_top-5-accuracy: 1.0000\n",
      "Epoch 29/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.2280 - accuracy: 0.9249 - top-5-accuracy: 1.0000 - val_loss: 0.2850 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 30/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.1884 - accuracy: 0.9331 - top-5-accuracy: 1.0000 - val_loss: 0.2703 - val_accuracy: 0.9193 - val_top-5-accuracy: 1.0000\n",
      "Epoch 31/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.1568 - accuracy: 0.9568 - top-5-accuracy: 1.0000 - val_loss: 0.2673 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 32/60\n",
      "31/31 [==============================] - 5s 167ms/step - loss: 0.1641 - accuracy: 0.9506 - top-5-accuracy: 1.0000 - val_loss: 0.3681 - val_accuracy: 0.8758 - val_top-5-accuracy: 1.0000\n",
      "Epoch 33/60\n",
      "31/31 [==============================] - 5s 165ms/step - loss: 0.1389 - accuracy: 0.9496 - top-5-accuracy: 1.0000 - val_loss: 0.3349 - val_accuracy: 0.9130 - val_top-5-accuracy: 1.0000\n",
      "Epoch 34/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.0995 - accuracy: 0.9743 - top-5-accuracy: 1.0000 - val_loss: 0.3538 - val_accuracy: 0.8944 - val_top-5-accuracy: 1.0000\n",
      "Epoch 35/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.1143 - accuracy: 0.9599 - top-5-accuracy: 1.0000 - val_loss: 0.4529 - val_accuracy: 0.8820 - val_top-5-accuracy: 1.0000\n",
      "Epoch 36/60\n",
      "31/31 [==============================] - 5s 171ms/step - loss: 0.0776 - accuracy: 0.9784 - top-5-accuracy: 1.0000 - val_loss: 0.3750 - val_accuracy: 0.9068 - val_top-5-accuracy: 1.0000\n",
      "Epoch 37/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0601 - accuracy: 0.9825 - top-5-accuracy: 1.0000 - val_loss: 0.3436 - val_accuracy: 0.8820 - val_top-5-accuracy: 1.0000\n",
      "Epoch 38/60\n",
      "31/31 [==============================] - 5s 171ms/step - loss: 0.0493 - accuracy: 0.9907 - top-5-accuracy: 1.0000 - val_loss: 0.2529 - val_accuracy: 0.9193 - val_top-5-accuracy: 1.0000\n",
      "Epoch 39/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0372 - accuracy: 0.9938 - top-5-accuracy: 1.0000 - val_loss: 0.2900 - val_accuracy: 0.9379 - val_top-5-accuracy: 1.0000\n",
      "Epoch 40/60\n",
      "31/31 [==============================] - 5s 166ms/step - loss: 0.0337 - accuracy: 0.9928 - top-5-accuracy: 1.0000 - val_loss: 0.4346 - val_accuracy: 0.9130 - val_top-5-accuracy: 0.9938\n",
      "Epoch 41/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0972 - accuracy: 0.9640 - top-5-accuracy: 1.0000 - val_loss: 0.3314 - val_accuracy: 0.9255 - val_top-5-accuracy: 1.0000\n",
      "Epoch 42/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0844 - accuracy: 0.9753 - top-5-accuracy: 1.0000 - val_loss: 0.4804 - val_accuracy: 0.8696 - val_top-5-accuracy: 1.0000\n",
      "Epoch 43/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0616 - accuracy: 0.9866 - top-5-accuracy: 1.0000 - val_loss: 0.2963 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 44/60\n",
      "31/31 [==============================] - 5s 165ms/step - loss: 0.0479 - accuracy: 0.9856 - top-5-accuracy: 1.0000 - val_loss: 0.2857 - val_accuracy: 0.8820 - val_top-5-accuracy: 1.0000\n",
      "Epoch 45/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0427 - accuracy: 0.9866 - top-5-accuracy: 1.0000 - val_loss: 0.3745 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 46/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0566 - accuracy: 0.9835 - top-5-accuracy: 1.0000 - val_loss: 0.3385 - val_accuracy: 0.8882 - val_top-5-accuracy: 1.0000\n",
      "Epoch 47/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0834 - accuracy: 0.9702 - top-5-accuracy: 1.0000 - val_loss: 0.2583 - val_accuracy: 0.9255 - val_top-5-accuracy: 0.9938\n",
      "Epoch 48/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0707 - accuracy: 0.9794 - top-5-accuracy: 1.0000 - val_loss: 0.3085 - val_accuracy: 0.9130 - val_top-5-accuracy: 1.0000\n",
      "Epoch 49/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0297 - accuracy: 0.9949 - top-5-accuracy: 1.0000 - val_loss: 0.2739 - val_accuracy: 0.9193 - val_top-5-accuracy: 1.0000\n",
      "Epoch 50/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0168 - accuracy: 0.9949 - top-5-accuracy: 1.0000 - val_loss: 0.2748 - val_accuracy: 0.9193 - val_top-5-accuracy: 1.0000\n",
      "Epoch 51/60\n",
      "31/31 [==============================] - 5s 164ms/step - loss: 0.0126 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2656 - val_accuracy: 0.9255 - val_top-5-accuracy: 1.0000\n",
      "Epoch 52/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.0181 - accuracy: 0.9959 - top-5-accuracy: 1.0000 - val_loss: 0.2471 - val_accuracy: 0.9503 - val_top-5-accuracy: 1.0000\n",
      "Epoch 53/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0425 - accuracy: 0.9887 - top-5-accuracy: 1.0000 - val_loss: 0.3769 - val_accuracy: 0.9006 - val_top-5-accuracy: 1.0000\n",
      "Epoch 54/60\n",
      "31/31 [==============================] - 5s 170ms/step - loss: 0.0183 - accuracy: 0.9979 - top-5-accuracy: 1.0000 - val_loss: 0.3074 - val_accuracy: 0.9130 - val_top-5-accuracy: 1.0000\n",
      "Epoch 55/60\n",
      "31/31 [==============================] - 5s 168ms/step - loss: 0.0052 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2340 - val_accuracy: 0.9441 - val_top-5-accuracy: 1.0000\n",
      "Epoch 56/60\n",
      "31/31 [==============================] - 5s 166ms/step - loss: 0.0035 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2452 - val_accuracy: 0.9379 - val_top-5-accuracy: 1.0000\n",
      "Epoch 57/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0024 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2057 - val_accuracy: 0.9503 - val_top-5-accuracy: 1.0000\n",
      "Epoch 58/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0018 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2103 - val_accuracy: 0.9379 - val_top-5-accuracy: 1.0000\n",
      "Epoch 59/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0015 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2167 - val_accuracy: 0.9317 - val_top-5-accuracy: 1.0000\n",
      "Epoch 60/60\n",
      "31/31 [==============================] - 5s 169ms/step - loss: 0.0015 - accuracy: 1.0000 - top-5-accuracy: 1.0000 - val_loss: 0.2181 - val_accuracy: 0.9379 - val_top-5-accuracy: 1.0000\n",
      "20/20 [==============================] - 1s 70ms/step - loss: 0.9767 - accuracy: 0.7902 - top-5-accuracy: 0.9754\n",
      "Test accuracy: 79.02%\n",
      "Test top 5 accuracy: 97.54%\n"
     ]
    }
   ],
   "source": [
    "def run_experiment():\n",
    "    # Initialize model\n",
    "    model = create_vivit_classifier(\n",
    "        tubelet_embedder=TubeletEmbedding(\n",
    "            embed_dim=PROJECTION_DIM, patch_size=PATCH_SIZE\n",
    "        ),\n",
    "        positional_encoder=PositionalEncoder(embed_dim=PROJECTION_DIM),\n",
    "    )\n",
    "\n",
    "    # Compile the model with the optimizer, loss function\n",
    "    # and the metrics.\n",
    "    optimizer = keras.optimizers.Adam(learning_rate=LEARNING_RATE)\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        loss=\"sparse_categorical_crossentropy\",\n",
    "        metrics=[\n",
    "            keras.metrics.SparseCategoricalAccuracy(name=\"accuracy\"),\n",
    "            keras.metrics.SparseTopKCategoricalAccuracy(5, name=\"top-5-accuracy\"),\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    # Train the model.\n",
    "    _ = model.fit(trainloader, epochs=EPOCHS, validation_data=validloader)\n",
    "\n",
    "    _, accuracy, top_5_accuracy = model.evaluate(testloader)\n",
    "    print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")\n",
    "    print(f\"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%\")\n",
    "\n",
    "    return model\n",
    "\n",
    "model = run_experiment()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I4KgP8-ClErE"
   },
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "WixAMFNWlErF"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 1s 989ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 22ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 24ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n",
      "1/1 [==============================] - 0s 23ms/step\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd8e118689aa4f6285ef1d4919f14aa6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "GridBox(children=(VBox(children=(HTML(value=\"'T: pancreas | P: pancreas'\"), Box(children=(Image(value=b'GIF89a…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "NUM_SAMPLES_VIZ = 25\n",
    "testsamples, labels = next(iter(testloader))\n",
    "testsamples, labels = testsamples[:NUM_SAMPLES_VIZ], labels[:NUM_SAMPLES_VIZ]\n",
    "\n",
    "ground_truths = []\n",
    "preds = []\n",
    "videos = []\n",
    "\n",
    "for i, (testsample, label) in enumerate(zip(testsamples, labels)):\n",
    "    # Generate gif\n",
    "    with io.BytesIO() as gif:\n",
    "        imageio.mimsave(gif, (testsample.numpy() * 255).astype(\"float32\"), \"GIF\", fps=5)\n",
    "        videos.append(gif.getvalue())\n",
    "\n",
    "    # Get model prediction\n",
    "    output = model.predict(tf.expand_dims(testsample, axis=0))[0]\n",
    "    pred = np.argmax(output, axis=0)\n",
    "\n",
    "    ground_truths.append(label.numpy().astype(\"int\"))\n",
    "    preds.append(pred)\n",
    "\n",
    "def make_box_for_grid(image_widget, fit):\n",
    "    \"\"\"Make a VBox to hold caption/image for demonstrating option_fit values.\n",
    "    Source: https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20Styling.html\n",
    "    \"\"\"\n",
    "    # Make the caption\n",
    "    if fit is not None:\n",
    "        fit_str = \"'{}'\".format(fit)\n",
    "    else:\n",
    "        fit_str = str(fit)\n",
    "\n",
    "    h = ipywidgets.HTML(value=\"\" + str(fit_str) + \"\")\n",
    "\n",
    "    # Make the green box with the image widget inside it\n",
    "    boxb = ipywidgets.widgets.Box()\n",
    "    boxb.children = [image_widget]\n",
    "\n",
    "    # Compose into a vertical box\n",
    "    vb = ipywidgets.widgets.VBox()\n",
    "    vb.layout.align_items = \"center\"\n",
    "    vb.children = [h, boxb]\n",
    "    return vb\n",
    "\n",
    "boxes = []\n",
    "for i in range(NUM_SAMPLES_VIZ):\n",
    "    ib = ipywidgets.widgets.Image(value=videos[i], width=100, height=100)\n",
    "    true_class = info[\"label\"][str(ground_truths[i])]\n",
    "    pred_class = info[\"label\"][str(preds[i])]\n",
    "    caption = f\"T: {true_class} | P: {pred_class}\"\n",
    "\n",
    "    boxes.append(make_box_for_grid(ib, caption))\n",
    "\n",
    "ipywidgets.widgets.GridBox(\n",
    "    boxes, layout=ipywidgets.widgets.Layout(grid_template_columns=\"repeat(5, 200px)\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F_chA1UnlErF"
   },
   "source": [
    "## Final thoughts\n",
    "\n",
    "With a vanilla implementation, we achieve ~79-80% Top-1 accuracy on the\n",
    "test dataset.\n",
    "\n",
    "The hyperparameters used in this tutorial were finalized by running a\n",
    "hyperparameter search using\n",
    "[W&B Sweeps](https://docs.wandb.ai/guides/sweeps).\n",
    "You can find out our sweeps result\n",
    "[here](https://wandb.ai/minimal-implementations/vivit/sweeps/66fp0lhz)\n",
    "and our quick analysis of the results\n",
    "[here](https://wandb.ai/minimal-implementations/vivit/reports/Hyperparameter-Tuning-Analysis--VmlldzoxNDEwNzcx)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "vivit",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Tensorflow (AI kit)",
   "language": "python",
   "name": "c009-intel_distribution_of_python_3_oneapi-beta05-tf"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}