vivit_test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Base settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# %env CUDA_VISIBLE_DEVICES=1\n",
    "\n",
    "import os\n",
    "import time\n",
    "from pprint import pprint\n",
    "from pathlib import Path\n",
    "from collections import OrderedDict\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import plotly.express as px\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from einops import rearrange, reduce, repeat\n",
    "from IPython.display import display\n",
    "\n",
    "import video_transformer.data_transform as T\n",
    "from video_transformer.dataset import DecordInit, load_annotation_data\n",
    "from video_transformer.transformer import PatchEmbed, TransformerContainer, ClassificationHead\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare kinetics400 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open('./data/kinetics400_val_metadata.pkl', 'rb') as f:\n",
    "    kinetics400_val_metadata = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/simc/workspace/dataset/kinetics/k400\n"
     ]
    }
   ],
   "source": [
    "# from torchvision.datasets import Kinetics\n",
    "from ood_with_vit.datasets.kinetics import MyKinetics\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "dataset_root = '~/workspace/dataset/kinetics/k400'\n",
    "dataset_root = os.path.expanduser(dataset_root)\n",
    "print(dataset_root)\n",
    "\n",
    "dataset_mean, dataset_std = (0.45, 0.45, 0.45), (0.225, 0.225, 0.225)\n",
    "val_transform = T.create_video_transform(\n",
    "    input_size=224,\n",
    "    is_training=False,\n",
    "    interpolation='bicubic',\n",
    "    mean=dataset_mean,\n",
    "    std=dataset_std,\n",
    ")\n",
    "# data_transform = T.Compose([\n",
    "#     T.Resize(scale_range=(-1, 256)),\n",
    "#     T.ThreeCrop(size=224),\n",
    "#     T.ToTensor(),\n",
    "#     T.Normalize(dataset_mean, dataset_std)\n",
    "# ])\n",
    "# data_transform.randomize_parameters()\n",
    "\n",
    "kinetics400_val_ds = MyKinetics(\n",
    "    root=dataset_root,\n",
    "    frames_per_clip=16,\n",
    "    split='val',\n",
    "    num_workers=24,\n",
    "    frame_rate=2,\n",
    "    step_between_clips=1,\n",
    "    transform=val_transform,\n",
    "    _precomputed_metadata=kinetics400_val_metadata,\n",
    ")\n",
    "\n",
    "kinetics400_val_dl = DataLoader(\n",
    "    dataset=kinetics400_val_ds,\n",
    "    batch_size=32,\n",
    "    shuffle=False,\n",
    "    num_workers=32,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 3, 224, 224]) abseiling\n"
     ]
    }
   ],
   "source": [
    "video, label = kinetics400_val_ds[0]\n",
    "print(video.shape, kinetics400_val_ds.classes[label])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "video paths: <class 'list'> 19881 /home/simc/workspace/dataset/kinetics/k400/val/abseiling/0wR5jVB-WPk_000417_000427.mp4\n",
      "video pts: <class 'list'> 19881 torch.Size([300])\n",
      "video fps: <class 'list'> 19881 29.97002997002997\n",
      "clips <class 'list'> 19881 torch.Size([5, 16])\n",
      "cumulative sizes: 88540\n"
     ]
    }
   ],
   "source": [
    "video_paths = kinetics400_val_ds.metadata['video_paths']\n",
    "video_pts = kinetics400_val_ds.metadata['video_pts']\n",
    "video_fps = kinetics400_val_ds.metadata['video_fps']\n",
    "clips = kinetics400_val_ds.video_clips.clips\n",
    "print('video paths:', type(video_paths), len(video_paths), video_paths[0])\n",
    "print('video pts:', type(video_pts), len(video_pts), video_pts[0].shape)\n",
    "print('video fps:', type(video_fps), len(video_fps), video_fps[0])\n",
    "print('clips', type(clips), len(clips), clips[0].shape)\n",
    "print('cumulative sizes:', kinetics400_val_ds.video_clips.cumulative_sizes[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare ViViT model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from video_transformer.video_transformer import ViViT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_state_dict(state_dict):\n",
    "\tfor old_key in list(state_dict.keys()):\n",
    "\t\tif old_key.startswith('model'):\n",
    "\t\t\tnew_key = old_key[6:]\n",
    "\t\t\tstate_dict[new_key] = state_dict.pop(old_key)\n",
    "\t\telse:\n",
    "\t\t\tnew_key = old_key[9:]\n",
    "\t\t\tstate_dict[new_key] = state_dict.pop(old_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_from_pretrain_(module, pretrained, init_module):\n",
    "    if torch.cuda.is_available():\n",
    "        state_dict = torch.load(pretrained)\n",
    "    else:\n",
    "        state_dict = torch.load(pretrained, map_location=torch.device('cpu'))\n",
    "    if init_module == 'transformer':\n",
    "        replace_state_dict(state_dict)\n",
    "    elif init_module == 'cls_head':\n",
    "        replace_state_dict(state_dict)\n",
    "    else:\n",
    "        raise TypeError(f'pretrained weights do not include the {init_module} module')\n",
    "    msg = module.load_state_dict(state_dict, strict=False)\n",
    "    return msg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_IncompatibleKeys(missing_keys=[], unexpected_keys=['cls_head.weight', 'cls_head.bias'])\n",
      "load model finished, the missing key of cls is:[]\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\") \n",
    "num_frames = 8\n",
    "frame_interval = 32\n",
    "num_class = 400\n",
    "arch = 'vivit' # turn to vivit for initializing vivit model\n",
    "\n",
    "pretrain_pth = './logs/vivit/vivit_model.pth'\n",
    "num_frames = num_frames * 2\n",
    "frame_interval = frame_interval // 2\n",
    "model = ViViT(\n",
    "    num_frames=num_frames,\n",
    "    img_size=224,\n",
    "    patch_size=16,\n",
    "    embed_dims=768,\n",
    "    in_channels=3,\n",
    "    attention_type='fact_encoder',\n",
    "    return_cls_token=True,\n",
    "    pretrain_pth=pretrain_pth,\n",
    "    weights_from='kinetics',\n",
    ")\n",
    "\n",
    "cls_head = ClassificationHead(num_classes=num_class, in_channels=768)\n",
    "# msg_trans = init_from_pretrain_(model, pretrain_pth, init_module='transformer')\n",
    "msg_cls = init_from_pretrain_(cls_head, pretrain_pth, init_module='cls_head')\n",
    "model.eval()\n",
    "cls_head.eval()\n",
    "model = model.to(device)\n",
    "cls_head = cls_head.to(device)\n",
    "print(f'load model finished, the missing key of cls is:{msg_cls[0]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2767/2767 [57:24<00:00,  1.24s/it]  "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.105 | Test Acc: 74.190% (65688/88540)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model.eval()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "total_test_loss, n_correct, n_total = 0, 0, 0\n",
    "with torch.no_grad():\n",
    "    start = time.time()\n",
    "    for batch_idx, (x, y) in enumerate(tqdm(kinetics400_val_dl)):\n",
    "        # print(f'batch: {batch_idx} {time.time() - start:.3f}')\n",
    "        start = time.time()\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        outputs = model(x)\n",
    "        outputs = cls_head(outputs)\n",
    "        loss = criterion(outputs, y)\n",
    "        # print(outputs.shape)\n",
    "        # print(f'model: {time.time() - start:.3f}')\n",
    "        start = time.time()\n",
    "\n",
    "        total_test_loss += loss.item()\n",
    "        _, predicted = outputs.max(1)\n",
    "        n_total += y.size(0)\n",
    "        n_correct += predicted.eq(y).sum().item()\n",
    "\n",
    "    avg_test_loss = total_test_loss / (batch_idx + 1)\n",
    "    test_accuracy = 100. * n_correct / n_total\n",
    "    print(f'Test Loss: {avg_test_loss:.3f} | Test Acc: {test_accuracy:.3f}% ({n_correct}/{n_total})')\n",
    "\n",
    "# return total_test_loss, test_accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<video controls width=\"480\" height=\"480\" src=\"./YABnJL_bDzw.mp4\">animation</video>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "video_path = './YABnJL_bDzw.mp4'\n",
    "html_str = '''\n",
    "<video controls width=\\\"480\\\" height=\\\"480\\\" src=\\\"{}\\\">animation</video>\n",
    "'''.format(video_path)\n",
    "display(HTML(html_str))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "302 22 278\n",
      "original: (16, 256, 454, 3)\n",
      "transformed: torch.Size([3, 16, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "# Prepare data preprocess\n",
    "mean, std = (0.45, 0.45, 0.45), (0.225, 0.225, 0.225)\n",
    "data_transform = T.Compose([\n",
    "        T.Resize(scale_range=(-1, 256)),\n",
    "        T.ThreeCrop(size=224),\n",
    "        T.ToTensor(),\n",
    "        T.Normalize(mean, std)\n",
    "        ])\n",
    "temporal_sample = T.TemporalRandomCrop(num_frames*frame_interval)\n",
    "\n",
    "# Sampling video frames\n",
    "video_decoder = DecordInit()\n",
    "v_reader = video_decoder(video_path)\n",
    "total_frames = len(v_reader)\n",
    "start_frame_ind, end_frame_ind = temporal_sample(total_frames)\n",
    "print(total_frames, start_frame_ind, end_frame_ind)\n",
    "if end_frame_ind-start_frame_ind < num_frames:\n",
    "    raise ValueError(f'the total frames of the video {video_path} is less than {num_frames}')\n",
    "frame_indice = np.linspace(0, end_frame_ind-start_frame_ind-1, num_frames, dtype=int)\n",
    "video = v_reader.get_batch(frame_indice).asnumpy()\n",
    "del v_reader\n",
    "\n",
    "print('original:', video.shape)\n",
    "video = torch.from_numpy(video).permute(0,3,1,2) # Video transform: T C H W\n",
    "data_transform.randomize_parameters()\n",
    "video = data_transform(video)\n",
    "video = video.to(device)\n",
    "print('transformed:', video.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 400])\n",
      "the shape of ouptut: torch.Size([400]),     and the prediction is: laughing\n"
     ]
    }
   ],
   "source": [
    "# Predict class label\n",
    "with torch.no_grad():\n",
    "    logits = model(video)\n",
    "    output = cls_head(logits)\n",
    "    print(output.shape)\n",
    "    output = output.view(3, 400).mean(0)\n",
    "    cls_pred = output.argmax().item()\n",
    "    \n",
    "print(f'the shape of ouptut: {output.shape}, \\\n",
    "    and the prediction is: {kinetics400_val_ds.classes[cls_pred]}')"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "949777d72b0d2535278d3dc13498b2535136f6dfe0678499012e853ee9abcab1"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}