guides/logit_checking/k600-ssv2-logit-matching-torch-vs-keras-cv-backbone.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b71a4667",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2024-03-31T16:27:53.343125Z",
     "iopub.status.busy": "2024-03-31T16:27:53.341887Z",
     "iopub.status.idle": "2024-03-31T16:28:15.976258Z",
     "shell.execute_reply": "2024-03-31T16:28:15.975233Z"
    },
    "papermill": {
     "duration": 22.646472,
     "end_time": "2024-03-31T16:28:15.978789",
     "exception": false,
     "start_time": "2024-03-31T16:27:53.332317",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install einops -q\n",
    "import logging\n",
    "from functools import partial\n",
    "import gc\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.checkpoint as checkpoint\n",
    "import numpy as np\n",
    "from timm.models.layers import  trunc_normal_ \n",
    "\n",
    "from functools import reduce, lru_cache\n",
    "from operator import mul\n",
    "from einops import rearrange\n",
    "import logging"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d5bd06",
   "metadata": {
    "papermill": {
     "duration": 0.006646,
     "end_time": "2024-03-31T16:28:15.992867",
     "exception": false,
     "start_time": "2024-03-31T16:28:15.986221",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Video Swin Model [PyTorch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "000c7886",
   "metadata": {
    "_kg_hide-input": true,
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:16.009192Z",
     "iopub.status.busy": "2024-03-31T16:28:16.008839Z",
     "iopub.status.idle": "2024-03-31T16:28:16.095208Z",
     "shell.execute_reply": "2024-03-31T16:28:16.093560Z"
    },
    "jupyter": {
     "source_hidden": true
    },
    "papermill": {
     "duration": 0.098049,
     "end_time": "2024-03-31T16:28:16.098022",
     "exception": false,
     "start_time": "2024-03-31T16:28:15.999973",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):\n",
    "    \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n",
    "    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,\n",
    "    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...\n",
    "    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for\n",
    "    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use\n",
    "    'survival rate' as the argument.\n",
    "    \"\"\"\n",
    "    if drop_prob == 0. or not training:\n",
    "        return x\n",
    "    keep_prob = 1 - drop_prob\n",
    "    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets\n",
    "    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)\n",
    "    if keep_prob > 0.0 and scale_by_keep:\n",
    "        random_tensor.div_(keep_prob)\n",
    "    return x * random_tensor\n",
    "\n",
    "class DropPath(nn.Module):\n",
    "    \"\"\"Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).\n",
    "    \"\"\"\n",
    "    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):\n",
    "        super(DropPath, self).__init__()\n",
    "        self.drop_prob = drop_prob\n",
    "        self.scale_by_keep = scale_by_keep\n",
    "\n",
    "    def forward(self, x):\n",
    "        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)\n",
    "\n",
    "    def extra_repr(self):\n",
    "        return f'drop_prob={round(self.drop_prob,3):0.3f}'\n",
    "    \n",
    "    \n",
    "def get_root_logger(log_file=None, log_level=logging.INFO):\n",
    "    \"\"\"Use ``get_logger`` method in mmcv to get the root logger.\n",
    "    The logger will be initialized if it has not been initialized. By default a\n",
    "    StreamHandler will be added. If ``log_file`` is specified, a FileHandler\n",
    "    will also be added. The name of the root logger is the top-level package\n",
    "    name, e.g., \"mmaction\".\n",
    "    Args:\n",
    "        log_file (str | None): The log filename. If specified, a FileHandler\n",
    "            will be added to the root logger.\n",
    "        log_level (int): The root logger level. Note that only the process of\n",
    "            rank 0 is affected, while other processes will set the level to\n",
    "            \"Error\" and be silent most of the time.\n",
    "    Returns:\n",
    "        :obj:`logging.Logger`: The root logger.\n",
    "    \"\"\"\n",
    "    return get_logger(__name__.split('.')[0], log_file, log_level)\n",
    "\n",
    "\n",
    "class Mlp(nn.Module):\n",
    "    \"\"\" Multilayer perceptron.\"\"\"\n",
    "\n",
    "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n",
    "        super().__init__()\n",
    "        out_features = out_features or in_features\n",
    "        hidden_features = hidden_features or in_features\n",
    "        self.fc1 = nn.Linear(in_features, hidden_features)\n",
    "        self.act = act_layer()\n",
    "        self.fc2 = nn.Linear(hidden_features, out_features)\n",
    "        self.drop = nn.Dropout(drop)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.act(x)\n",
    "        x = self.drop(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.drop(x)\n",
    "        return x\n",
    "    \n",
    "    \n",
    "def window_partition(x, window_size):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        x: (B, D, H, W, C)\n",
    "        window_size (tuple[int]): window size\n",
    "\n",
    "    Returns:\n",
    "        windows: (B*num_windows, window_size*window_size, C)\n",
    "    \"\"\"\n",
    "    B, D, H, W, C = x.shape\n",
    "    x = x.view(\n",
    "        B, \n",
    "        D // window_size[0], window_size[0], \n",
    "        H // window_size[1], window_size[1], \n",
    "        W // window_size[2], window_size[2], \n",
    "        C\n",
    "    )\n",
    "    windows = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, reduce(mul, window_size), C)\n",
    "    return windows\n",
    "\n",
    "\n",
    "def window_reverse(windows, window_size, B, D, H, W):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        windows: (B*num_windows, window_size, window_size, C)\n",
    "        window_size (tuple[int]): Window size\n",
    "        H (int): Height of image\n",
    "        W (int): Width of image\n",
    "\n",
    "    Returns:\n",
    "        x: (B, D, H, W, C)\n",
    "    \"\"\"\n",
    "    x = windows.view(\n",
    "        B, \n",
    "        D // window_size[0], \n",
    "        H // window_size[1], \n",
    "        W // window_size[2], \n",
    "        window_size[0], \n",
    "        window_size[1], \n",
    "        window_size[2], \n",
    "        -1\n",
    "    )\n",
    "    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1)\n",
    "    return x\n",
    "\n",
    "\n",
    "def get_window_size(x_size, window_size, shift_size=None):\n",
    "    use_window_size = list(window_size)\n",
    "    if shift_size is not None:\n",
    "        use_shift_size = list(shift_size)\n",
    "    for i in range(len(x_size)):\n",
    "        if x_size[i] <= window_size[i]:\n",
    "            use_window_size[i] = x_size[i]\n",
    "            if shift_size is not None:\n",
    "                use_shift_size[i] = 0\n",
    "    if shift_size is None:\n",
    "        return tuple(use_window_size)\n",
    "    else:\n",
    "        return tuple(use_window_size), tuple(use_shift_size)\n",
    "    \n",
    "    \n",
    "class WindowAttention3D(nn.Module):\n",
    "    \"\"\" Window based multi-head self attention (W-MSA) module with relative position bias.\n",
    "    It supports both of shifted and non-shifted window.\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        window_size (tuple[int]): The temporal length, height and width of the window.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set\n",
    "        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n",
    "        proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dim, window_size, num_heads, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.window_size = window_size  # Wd, Wh, Ww\n",
    "        self.num_heads = num_heads\n",
    "        head_dim = dim // num_heads\n",
    "        self.scale = qk_scale or head_dim ** -0.5\n",
    "\n",
    "        # define a parameter table of relative position bias\n",
    "        self.relative_position_bias_table = nn.Parameter(\n",
    "            torch.zeros(\n",
    "                (2 * window_size[0] - 1) * (2 * window_size[1] - 1) * (2 * window_size[2] - 1), num_heads)\n",
    "        )  # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH\n",
    "\n",
    "        # get pair-wise relative position index for each token inside the window\n",
    "        coords_d = torch.arange(self.window_size[0])\n",
    "        coords_h = torch.arange(self.window_size[1])\n",
    "        coords_w = torch.arange(self.window_size[2])\n",
    "        coords = torch.stack(torch.meshgrid(coords_d, coords_h, coords_w))  # 3, Wd, Wh, Ww\n",
    "        coords_flatten = torch.flatten(coords, 1)  # 3, Wd*Wh*Ww\n",
    "        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 3, Wd*Wh*Ww, Wd*Wh*Ww\n",
    "        \n",
    "        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wd*Wh*Ww, Wd*Wh*Ww, 3\n",
    "        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0\n",
    "        relative_coords[:, :, 1] += self.window_size[1] - 1\n",
    "        relative_coords[:, :, 2] += self.window_size[2] - 1\n",
    "        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)\n",
    "        relative_coords[:, :, 1] *= (2 * self.window_size[2] - 1)\n",
    "        relative_position_index = relative_coords.sum(-1)  # Wd*Wh*Ww, Wd*Wh*Ww\n",
    "\n",
    "        self.register_buffer(\"relative_position_index\", relative_position_index)\n",
    "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
    "        self.attn_drop = nn.Dropout(attn_drop)\n",
    "        self.proj = nn.Linear(dim, dim)\n",
    "        self.proj_drop = nn.Dropout(proj_drop)\n",
    "\n",
    "        trunc_normal_(self.relative_position_bias_table, std=.02)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        \"\"\" Forward function.\n",
    "        Args:\n",
    "            x: input features with shape of (num_windows*B, N, C)\n",
    "            mask: (0/-inf) mask with shape of (num_windows, N, N) or None\n",
    "        \"\"\"\n",
    "        B_, N, C = x.shape\n",
    "        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n",
    "        q, k, v = qkv[0], qkv[1], qkv[2]  # B_, nH, N, C\n",
    "        q = q * self.scale\n",
    "        attn = q @ k.transpose(-2, -1)\n",
    "\n",
    "        relative_position_bias = self.relative_position_bias_table[\n",
    "            self.relative_position_index[:N, :N].reshape(-1)\n",
    "        ].reshape(N, N, -1)  # Wd*Wh*Ww,Wd*Wh*Ww,nH\n",
    "        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww\n",
    "        attn = attn + relative_position_bias.unsqueeze(0) # B_, nH, N, N\n",
    "        \n",
    "        if mask is not None:\n",
    "            nW = mask.shape[0]\n",
    "            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n",
    "            attn = attn.view(-1, self.num_heads, N, N)\n",
    "            attn = self.softmax(attn)\n",
    "        else:\n",
    "            attn = self.softmax(attn)\n",
    "\n",
    "        attn = self.attn_drop(attn)\n",
    "        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)\n",
    "        x = self.proj(x)\n",
    "        x = self.proj_drop(x)\n",
    "    \n",
    "        return x\n",
    "    \n",
    "    \n",
    "class SwinTransformerBlock3D(nn.Module):\n",
    "    \"\"\" Swin Transformer Block.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        window_size (tuple[int]): Window size.\n",
    "        shift_size (tuple[int]): Shift size for SW-MSA.\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n",
    "        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop (float, optional): Dropout rate. Default: 0.0\n",
    "        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n",
    "        drop_path (float, optional): Stochastic depth rate. Default: 0.0\n",
    "        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n",
    "        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dim, num_heads, window_size=(2,7,7), shift_size=(0,0,0),\n",
    "                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,\n",
    "                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_checkpoint=False):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.num_heads = num_heads\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = shift_size\n",
    "        self.mlp_ratio = mlp_ratio\n",
    "        self.use_checkpoint=use_checkpoint\n",
    "\n",
    "        assert 0 <= self.shift_size[0] < self.window_size[0], \"shift_size must in 0-window_size\"\n",
    "        assert 0 <= self.shift_size[1] < self.window_size[1], \"shift_size must in 0-window_size\"\n",
    "        assert 0 <= self.shift_size[2] < self.window_size[2], \"shift_size must in 0-window_size\"\n",
    "\n",
    "        self.norm1 = norm_layer(dim)\n",
    "        self.attn = WindowAttention3D(\n",
    "            dim, window_size=self.window_size, num_heads=num_heads,\n",
    "            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)\n",
    "\n",
    "        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n",
    "\n",
    "    def forward_part1(self, x, mask_matrix):\n",
    "        B, D, H, W, C = x.shape\n",
    "        window_size, shift_size = get_window_size((D, H, W), self.window_size, self.shift_size)\n",
    "\n",
    "        x = self.norm1(x)\n",
    "        # pad feature maps to multiples of window size\n",
    "        pad_l = pad_t = pad_d0 = 0\n",
    "        pad_d1 = (window_size[0] - D % window_size[0]) % window_size[0]\n",
    "        pad_b = (window_size[1] - H % window_size[1]) % window_size[1]\n",
    "        pad_r = (window_size[2] - W % window_size[2]) % window_size[2]\n",
    "        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b, pad_d0, pad_d1))\n",
    "        _, Dp, Hp, Wp, _ = x.shape\n",
    "        # cyclic shift\n",
    "        if any(i > 0 for i in shift_size):\n",
    "            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1], -shift_size[2]), dims=(1, 2, 3))\n",
    "            attn_mask = mask_matrix\n",
    "        else:\n",
    "            shifted_x = x\n",
    "            attn_mask = None\n",
    "        # partition windows\n",
    "        x_windows = window_partition(shifted_x, window_size)  # B*nW, Wd*Wh*Ww, C\n",
    "        # W-MSA/SW-MSA\n",
    "        attn_windows = self.attn(x_windows, mask=attn_mask)  # B*nW, Wd*Wh*Ww, C\n",
    "        # merge windows\n",
    "        attn_windows = attn_windows.view(-1, *(window_size+(C,)))\n",
    "        shifted_x = window_reverse(attn_windows, window_size, B, Dp, Hp, Wp)  # B D' H' W' C\n",
    "        # reverse cyclic shift\n",
    "        if any(i > 0 for i in shift_size):\n",
    "            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1], shift_size[2]), dims=(1, 2, 3))\n",
    "        else:\n",
    "            x = shifted_x\n",
    "\n",
    "        if pad_d1 >0 or pad_r > 0 or pad_b > 0:\n",
    "            x = x[:, :D, :H, :W, :].contiguous()\n",
    "        return x\n",
    "\n",
    "    def forward_part2(self, x):\n",
    "        return self.drop_path(self.mlp(self.norm2(x)))\n",
    "\n",
    "    def forward(self, x, mask_matrix):\n",
    "        \"\"\" Forward function.\n",
    "\n",
    "        Args:\n",
    "            x: Input feature, tensor size (B, D, H, W, C).\n",
    "            mask_matrix: Attention mask for cyclic shift.\n",
    "        \"\"\"\n",
    "        \n",
    "        shortcut = x\n",
    "        if self.use_checkpoint:\n",
    "            x = checkpoint.checkpoint(self.forward_part1, x, mask_matrix)\n",
    "        else:\n",
    "            x = self.forward_part1(x, mask_matrix)\n",
    "\n",
    "        x = shortcut + self.drop_path(x)\n",
    "\n",
    "        if self.use_checkpoint:\n",
    "            x = x + checkpoint.checkpoint(self.forward_part2, x)\n",
    "        else:\n",
    "            x = x + self.forward_part2(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    \n",
    "class PatchMerging(nn.Module):\n",
    "    \"\"\" Patch Merging Layer\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n",
    "    \"\"\"\n",
    "    def __init__(self, dim, norm_layer=nn.LayerNorm):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n",
    "        self.norm = norm_layer(4 * dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\" Forward function.\n",
    "\n",
    "        Args:\n",
    "            x: Input feature, tensor size (B, D, H, W, C).\n",
    "        \"\"\"\n",
    "        B, D, H, W, C = x.shape\n",
    "\n",
    "        # padding\n",
    "        pad_input = (H % 2 == 1) or (W % 2 == 1)\n",
    "        if pad_input:\n",
    "            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))\n",
    "\n",
    "        x0 = x[:, :, 0::2, 0::2, :]  # B D H/2 W/2 C\n",
    "        x1 = x[:, :, 1::2, 0::2, :]  # B D H/2 W/2 C\n",
    "        x2 = x[:, :, 0::2, 1::2, :]  # B D H/2 W/2 C\n",
    "        x3 = x[:, :, 1::2, 1::2, :]  # B D H/2 W/2 C\n",
    "        x = torch.cat([x0, x1, x2, x3], -1)  # B D H/2 W/2 4*C\n",
    "\n",
    "        x = self.norm(x)\n",
    "        x = self.reduction(x)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    \n",
    "def compute_mask(D, H, W, window_size, shift_size, device):\n",
    "    img_mask = torch.zeros((1, D, H, W, 1), device=device)  # 1 Dp Hp Wp 1\n",
    "    cnt = 0\n",
    "    for d in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0],None):\n",
    "        for h in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1],None):\n",
    "            for w in slice(-window_size[2]), slice(-window_size[2], -shift_size[2]), slice(-shift_size[2],None):\n",
    "                img_mask[:, d, h, w, :] = cnt\n",
    "                cnt += 1\n",
    "    mask_windows = window_partition(img_mask, window_size)  # nW, ws[0]*ws[1]*ws[2], 1\n",
    "    mask_windows = mask_windows.squeeze(-1)  # nW, ws[0]*ws[1]*ws[2]\n",
    "    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n",
    "    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n",
    "    return attn_mask\n",
    "\n",
    "class BasicLayer(nn.Module):\n",
    "    \"\"\" A basic Swin Transformer layer for one stage.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of feature channels\n",
    "        depth (int): Depths of this stage.\n",
    "        num_heads (int): Number of attention head.\n",
    "        window_size (tuple[int]): Local window size. Default: (1,7,7).\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.\n",
    "        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop (float, optional): Dropout rate. Default: 0.0\n",
    "        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n",
    "        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n",
    "        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 dim,\n",
    "                 depth,\n",
    "                 num_heads,\n",
    "                 window_size=(1,7,7),\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=False,\n",
    "                 qk_scale=None,\n",
    "                 drop=0.,\n",
    "                 attn_drop=0.,\n",
    "                 drop_path=0.,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 downsample=None,\n",
    "                 use_checkpoint=False):\n",
    "        super().__init__()\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = tuple(i // 2 for i in window_size)\n",
    "        self.depth = depth\n",
    "        self.use_checkpoint = use_checkpoint\n",
    "\n",
    "        # build blocks\n",
    "        self.blocks = nn.ModuleList([\n",
    "            SwinTransformerBlock3D(\n",
    "                dim=dim,\n",
    "                num_heads=num_heads,\n",
    "                window_size=window_size,\n",
    "                shift_size=(0,0,0) if (i % 2 == 0) else self.shift_size,\n",
    "                mlp_ratio=mlp_ratio,\n",
    "                qkv_bias=qkv_bias,\n",
    "                qk_scale=qk_scale,\n",
    "                drop=drop,\n",
    "                attn_drop=attn_drop,\n",
    "                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n",
    "                norm_layer=norm_layer,\n",
    "                use_checkpoint=use_checkpoint,\n",
    "            )\n",
    "            for i in range(depth)])\n",
    "        \n",
    "        self.downsample = downsample\n",
    "        if self.downsample is not None:\n",
    "            self.downsample = downsample(dim=dim, norm_layer=norm_layer)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\" Forward function.\n",
    "\n",
    "        Args:\n",
    "            x: Input feature, tensor size (B, C, D, H, W).\n",
    "        \"\"\"\n",
    "        # calculate attention mask for SW-MSA\n",
    "        B, C, D, H, W = x.shape\n",
    "        window_size, shift_size = get_window_size((D,H,W), self.window_size, self.shift_size)\n",
    "        x = rearrange(x, 'b c d h w -> b d h w c')\n",
    "        Dp = int(np.ceil(D / window_size[0])) * window_size[0]\n",
    "        Hp = int(np.ceil(H / window_size[1])) * window_size[1]\n",
    "        Wp = int(np.ceil(W / window_size[2])) * window_size[2]\n",
    "        attn_mask = compute_mask(Dp, Hp, Wp, window_size, shift_size, x.device)\n",
    "        \n",
    "\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x, attn_mask)\n",
    "        x = x.view(B, D, H, W, -1)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            x = self.downsample(x)\n",
    " \n",
    "        x = rearrange(x, 'b d h w c -> b c d h w')\n",
    "        return x\n",
    "    \n",
    "class PatchEmbed3D(nn.Module):\n",
    "    \"\"\" Video to Patch Embedding.\n",
    "\n",
    "    Args:\n",
    "        patch_size (int): Patch token size. Default: (2,4,4).\n",
    "        in_chans (int): Number of input video channels. Default: 3.\n",
    "        embed_dim (int): Number of linear projection output channels. Default: 96.\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: None\n",
    "    \"\"\"\n",
    "    def __init__(self, patch_size=(2,4,4), in_chans=3, embed_dim=96, norm_layer=None):\n",
    "        super().__init__()\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        self.in_chans = in_chans\n",
    "        self.embed_dim = embed_dim\n",
    "\n",
    "        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)\n",
    "        if norm_layer is not None:\n",
    "            self.norm = norm_layer(embed_dim)\n",
    "        else:\n",
    "            self.norm = None\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Forward function.\"\"\"\n",
    "        # padding\n",
    "        _, _, D, H, W = x.size()\n",
    "        if W % self.patch_size[2] != 0:\n",
    "            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))\n",
    "        if H % self.patch_size[1] != 0:\n",
    "            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))\n",
    "        if D % self.patch_size[0] != 0:\n",
    "            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))\n",
    "  \n",
    "        x = self.proj(x)  # B C D Wh Ww\n",
    "\n",
    "        if self.norm is not None:\n",
    "            D, Wh, Ww = x.size(2), x.size(3), x.size(4)\n",
    "            x = x.flatten(2).transpose(1, 2)\n",
    "            x = self.norm(x)\n",
    "            x = x.transpose(1, 2).view(-1, self.embed_dim, D, Wh, Ww)\n",
    "\n",
    "        return x\n",
    "    \n",
    "    \n",
    "class SwinTransformer3D(nn.Module):\n",
    "    \"\"\" Swin Transformer backbone.\n",
    "        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -\n",
    "          https://arxiv.org/pdf/2103.14030\n",
    "\n",
    "    Args:\n",
    "        patch_size (int | tuple(int)): Patch size. Default: (4,4,4).\n",
    "        in_chans (int): Number of input image channels. Default: 3.\n",
    "        embed_dim (int): Number of linear projection output channels. Default: 96.\n",
    "        depths (tuple[int]): Depths of each Swin Transformer stage.\n",
    "        num_heads (tuple[int]): Number of attention head of each stage.\n",
    "        window_size (int): Window size. Default: 7.\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.\n",
    "        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: Truee\n",
    "        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop_rate (float): Dropout rate.\n",
    "        attn_drop_rate (float): Attention dropout rate. Default: 0.\n",
    "        drop_path_rate (float): Stochastic depth rate. Default: 0.2.\n",
    "        norm_layer: Normalization layer. Default: nn.LayerNorm.\n",
    "        patch_norm (bool): If True, add normalization after patch embedding. Default: False.\n",
    "        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).\n",
    "            -1 means not freezing any parameters.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 pretrained=None,\n",
    "                 pretrained2d=True,\n",
    "                 patch_size=(4,4,4),\n",
    "                 in_chans=3,\n",
    "                 embed_dim=96,\n",
    "                 depths=[2, 2, 6, 2],\n",
    "                 num_heads=[3, 6, 12, 24],\n",
    "                 window_size=(2,7,7),\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop_rate=0.,\n",
    "                 attn_drop_rate=0.,\n",
    "                 drop_path_rate=0.2,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 patch_norm=False,\n",
    "                 frozen_stages=-1,\n",
    "                 use_checkpoint=False,\n",
    "                 \n",
    "                 # class head\n",
    "                 spatial_type='avg',\n",
    "                 in_channels=768,\n",
    "                 num_classes=400,\n",
    "                 dropout_ratio=0.5 # to do check: no dropout layer in weight state\n",
    "                ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.pretrained = pretrained\n",
    "        self.pretrained2d = pretrained2d\n",
    "        self.num_layers = len(depths)\n",
    "        self.embed_dim = embed_dim\n",
    "        self.patch_norm = patch_norm\n",
    "        self.frozen_stages = frozen_stages\n",
    "        self.window_size = window_size\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        # split image into non-overlapping patches\n",
    "        self.patch_embed = PatchEmbed3D(\n",
    "            patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,\n",
    "            norm_layer=norm_layer if self.patch_norm else None)\n",
    "\n",
    "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
    "\n",
    "        # stochastic depth\n",
    "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule\n",
    "\n",
    "        # build layers\n",
    "        self.layers = nn.ModuleList()\n",
    "        for i_layer in range(self.num_layers):\n",
    "            layer = BasicLayer(\n",
    "                dim=int(embed_dim * 2**i_layer),\n",
    "                depth=depths[i_layer],\n",
    "                num_heads=num_heads[i_layer],\n",
    "                window_size=window_size,\n",
    "                mlp_ratio=mlp_ratio,\n",
    "                qkv_bias=qkv_bias,\n",
    "                qk_scale=qk_scale,\n",
    "                drop=drop_rate,\n",
    "                attn_drop=attn_drop_rate,\n",
    "                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],\n",
    "                norm_layer=norm_layer,\n",
    "                downsample=PatchMerging if i_layer<self.num_layers-1 else None,\n",
    "                use_checkpoint=use_checkpoint)\n",
    "            self.layers.append(layer)\n",
    "\n",
    "        self.num_features = int(embed_dim * 2**(self.num_layers-1))\n",
    "\n",
    "        # add a norm layer for each output\n",
    "        self.norm = norm_layer(self.num_features)\n",
    "        self._freeze_stages()\n",
    "\n",
    "    def _freeze_stages(self):\n",
    "        if self.frozen_stages >= 0:\n",
    "            self.patch_embed.eval()\n",
    "            for param in self.patch_embed.parameters():\n",
    "                param.requires_grad = False\n",
    "\n",
    "        if self.frozen_stages >= 1:\n",
    "            self.pos_drop.eval()\n",
    "            for i in range(0, self.frozen_stages):\n",
    "                m = self.layers[i]\n",
    "                m.eval()\n",
    "                for param in m.parameters():\n",
    "                    param.requires_grad = False\n",
    "\n",
    "    def inflate_weights(self, logger):\n",
    "        \"\"\"Inflate the swin2d parameters to swin3d.\n",
    "\n",
    "        The differences between swin3d and swin2d mainly lie in an extra\n",
    "        axis. To utilize the pretrained parameters in 2d model,\n",
    "        the weight of swin2d models should be inflated to fit in the shapes of\n",
    "        the 3d counterpart.\n",
    "\n",
    "        Args:\n",
    "            logger (logging.Logger): The logger used to print\n",
    "                debugging infomation.\n",
    "        \"\"\"\n",
    "        checkpoint = torch.load(self.pretrained, map_location='cpu')\n",
    "        state_dict = checkpoint['model']\n",
    "\n",
    "        # delete relative_position_index since we always re-init it\n",
    "        relative_position_index_keys = [k for k in state_dict.keys() if \"relative_position_index\" in k]\n",
    "        for k in relative_position_index_keys:\n",
    "            del state_dict[k]\n",
    "\n",
    "        # delete attn_mask since we always re-init it\n",
    "        attn_mask_keys = [k for k in state_dict.keys() if \"attn_mask\" in k]\n",
    "        for k in attn_mask_keys:\n",
    "            del state_dict[k]\n",
    "\n",
    "        state_dict['patch_embed.proj.weight'] = state_dict['patch_embed.proj.weight'].unsqueeze(2).repeat(1,1,self.patch_size[0],1,1) / self.patch_size[0]\n",
    "\n",
    "        # bicubic interpolate relative_position_bias_table if not match\n",
    "        relative_position_bias_table_keys = [k for k in state_dict.keys() if \"relative_position_bias_table\" in k]\n",
    "        for k in relative_position_bias_table_keys:\n",
    "            relative_position_bias_table_pretrained = state_dict[k]\n",
    "            relative_position_bias_table_current = self.state_dict()[k]\n",
    "            L1, nH1 = relative_position_bias_table_pretrained.size()\n",
    "            L2, nH2 = relative_position_bias_table_current.size()\n",
    "            L2 = (2*self.window_size[1]-1) * (2*self.window_size[2]-1)\n",
    "            wd = self.window_size[0]\n",
    "            if nH1 != nH2:\n",
    "                logger.warning(f\"Error in loading {k}, passing\")\n",
    "            else:\n",
    "                if L1 != L2:\n",
    "                    S1 = int(L1 ** 0.5)\n",
    "                    relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(\n",
    "                        relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(2*self.window_size[1]-1, 2*self.window_size[2]-1),\n",
    "                        mode='bicubic')\n",
    "                    relative_position_bias_table_pretrained = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)\n",
    "            state_dict[k] = relative_position_bias_table_pretrained.repeat(2*wd-1,1)\n",
    "\n",
    "        msg = self.load_state_dict(state_dict, strict=False)\n",
    "        logger.info(msg)\n",
    "        logger.info(f\"=> loaded successfully '{self.pretrained}'\")\n",
    "        del checkpoint\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "    def init_weights(self, pretrained=None):\n",
    "        \"\"\"Initialize the weights in backbone.\n",
    "\n",
    "        Args:\n",
    "            pretrained (str, optional): Path to pre-trained weights.\n",
    "                Defaults to None.\n",
    "        \"\"\"\n",
    "        def _init_weights(m):\n",
    "            if isinstance(m, nn.Linear):\n",
    "                trunc_normal_(m.weight, std=.02)\n",
    "                if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                    nn.init.constant_(m.bias, 0)\n",
    "            elif isinstance(m, nn.LayerNorm):\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "                nn.init.constant_(m.weight, 1.0)\n",
    "\n",
    "        if pretrained:\n",
    "            self.pretrained = pretrained\n",
    "        if isinstance(self.pretrained, str):\n",
    "            self.apply(_init_weights)\n",
    "            logger = get_root_logger()\n",
    "            logger.info(f'load model from: {self.pretrained}')\n",
    "\n",
    "            if self.pretrained2d:\n",
    "                # Inflate 2D model into 3D model.\n",
    "                self.inflate_weights(logger)\n",
    "            else:\n",
    "                # Directly load 3D model.\n",
    "                load_checkpoint(self, self.pretrained, strict=False, logger=logger)\n",
    "        elif self.pretrained is None:\n",
    "            self.apply(_init_weights)\n",
    "        else:\n",
    "            raise TypeError('pretrained must be a str or None')\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Forward function.\"\"\"\n",
    "\n",
    "        x = self.patch_embed(x)\n",
    "        x = self.pos_drop(x)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x.contiguous())\n",
    " \n",
    "        x = rearrange(x, 'n c d h w -> n d h w c')\n",
    "        x = self.norm(x)\n",
    "        x = rearrange(x, 'n d h w c -> n c d h w')\n",
    "        return x\n",
    "    \n",
    "\n",
    "    def train(self, mode=True):\n",
    "        \"\"\"Convert the model into training mode while keep layers freezed.\"\"\"\n",
    "        super(SwinTransformer3D, self).train(mode)\n",
    "        self._freeze_stages()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c98f81e7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:16.115247Z",
     "iopub.status.busy": "2024-03-31T16:28:16.114872Z",
     "iopub.status.idle": "2024-03-31T16:28:16.122078Z",
     "shell.execute_reply": "2024-03-31T16:28:16.120287Z"
    },
    "papermill": {
     "duration": 0.018356,
     "end_time": "2024-03-31T16:28:16.124734",
     "exception": false,
     "start_time": "2024-03-31T16:28:16.106378",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def video_swin_base(window_size=(8,7,7), **kwargs):\n",
    "    model = SwinTransformer3D(\n",
    "        patch_size=(2,4,4),\n",
    "        embed_dim=128,\n",
    "        depths=[2, 2, 18, 2],\n",
    "        num_heads=[4, 8, 16, 32],\n",
    "        window_size=window_size,\n",
    "        mlp_ratio=4.,\n",
    "        qkv_bias=True,\n",
    "        qk_scale=None,\n",
    "        drop_rate=0.,\n",
    "        attn_drop_rate=0.,\n",
    "        drop_path_rate=0.2,\n",
    "        norm_layer=nn.LayerNorm,\n",
    "        patch_norm=True,\n",
    "        in_channels=1024,\n",
    "        **kwargs\n",
    "    )\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f2ba54",
   "metadata": {
    "papermill": {
     "duration": 0.006702,
     "end_time": "2024-03-31T16:28:16.138791",
     "exception": false,
     "start_time": "2024-03-31T16:28:16.132089",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Video Swin K600 PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d5b7dbf5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:16.155776Z",
     "iopub.status.busy": "2024-03-31T16:28:16.155312Z",
     "iopub.status.idle": "2024-03-31T16:28:17.291208Z",
     "shell.execute_reply": "2024-03-31T16:28:17.289281Z"
    },
    "papermill": {
     "duration": 1.147943,
     "end_time": "2024-03-31T16:28:17.293909",
     "exception": false,
     "start_time": "2024-03-31T16:28:16.145966",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3526.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "model_pt = video_swin_base(\n",
    "    window_size=(8,7,7), num_classes=600\n",
    ")\n",
    "model_pt.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83bb8b13",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:17.310229Z",
     "iopub.status.busy": "2024-03-31T16:28:17.309869Z",
     "iopub.status.idle": "2024-03-31T16:28:19.168922Z",
     "shell.execute_reply": "2024-03-31T16:28:19.167768Z"
    },
    "papermill": {
     "duration": 1.870111,
     "end_time": "2024-03-31T16:28:19.171289",
     "exception": false,
     "start_time": "2024-03-31T16:28:17.301178",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-31 16:28:17--  https://github.com/SwinTransformer/storage/releases/download/v1.0.4/swin_base_patch244_window877_kinetics600_22k.pth\r\n",
      "Resolving github.com (github.com)... 140.82.113.4\r\n",
      "Connecting to github.com (github.com)|140.82.113.4|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 302 Found\r\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/357198522/099f2980-d55e-11eb-8848-6616f5f65526?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162817Z&X-Amz-Expires=300&X-Amz-Signature=1ecb02f3ad86594d9680a52ff630b0668a6c2a170f2902add41aac346e8a286b&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=357198522&response-content-disposition=attachment%3B%20filename%3Dswin_base_patch244_window877_kinetics600_22k.pth&response-content-type=application%2Foctet-stream [following]\r\n",
      "--2024-03-31 16:28:17--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/357198522/099f2980-d55e-11eb-8848-6616f5f65526?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162817Z&X-Amz-Expires=300&X-Amz-Signature=1ecb02f3ad86594d9680a52ff630b0668a6c2a170f2902add41aac346e8a286b&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=357198522&response-content-disposition=attachment%3B%20filename%3Dswin_base_patch244_window877_kinetics600_22k.pth&response-content-type=application%2Foctet-stream\r\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.111.133, ...\r\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.110.133|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 200 OK\r\n",
      "Length: 382579368 (365M) [application/octet-stream]\r\n",
      "Saving to: 'checkpoint.pt'\r\n",
      "\r\n",
      "checkpoint.pt       100%[===================>] 364.86M   303MB/s    in 1.2s    \r\n",
      "\r\n",
      "2024-03-31 16:28:19 (303 MB/s) - 'checkpoint.pt' saved [382579368/382579368]\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "base_url = \"https://github.com/SwinTransformer/storage/releases/download/v1.0.4/\"\n",
    "checkpoints_pt = f\"{base_url}swin_base_patch244_window877_kinetics600_22k.pth\"\n",
    "!wget {checkpoints_pt} -O checkpoint.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ace994fc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:19.190480Z",
     "iopub.status.busy": "2024-03-31T16:28:19.190060Z",
     "iopub.status.idle": "2024-03-31T16:28:19.389671Z",
     "shell.execute_reply": "2024-03-31T16:28:19.388429Z"
    },
    "papermill": {
     "duration": 0.21121,
     "end_time": "2024-03-31T16:28:19.391988",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.180778",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "state_dict = torch.load(\n",
    "    'checkpoint.pt', map_location=\"cpu\"\n",
    ")\n",
    "state_dict = state_dict['state_dict']\n",
    "state_dict = {k.replace('backbone.', ''): v for k, v in state_dict.items()}\n",
    "state_dict = {k.replace('cls_head.', ''): v for k, v in state_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5a81cc80",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:19.425761Z",
     "iopub.status.busy": "2024-03-31T16:28:19.425410Z",
     "iopub.status.idle": "2024-03-31T16:28:19.503955Z",
     "shell.execute_reply": "2024-03-31T16:28:19.502975Z"
    },
    "papermill": {
     "duration": 0.097974,
     "end_time": "2024-03-31T16:28:19.505852",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.407878",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_IncompatibleKeys(missing_keys=[], unexpected_keys=['fc_cls.weight', 'fc_cls.bias'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_pt.load_state_dict(state_dict, strict=False) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a282a412",
   "metadata": {
    "papermill": {
     "duration": 0.007572,
     "end_time": "2024-03-31T16:28:19.521247",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.513675",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "_IncompatibleKeys(missing_keys=[], unexpected_keys=['fc_cls.weight', 'fc_cls.bias'])\n",
    "\n",
    "\n",
    "Expected, we removed final layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db2e374",
   "metadata": {
    "papermill": {
     "duration": 0.007589,
     "end_time": "2024-03-31T16:28:19.536897",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.529308",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d671cf7a",
   "metadata": {
    "papermill": {
     "duration": 0.007869,
     "end_time": "2024-03-31T16:28:19.552746",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.544877",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Video Swin K600 [Keras CV]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "97a792d6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:19.570195Z",
     "iopub.status.busy": "2024-03-31T16:28:19.569809Z",
     "iopub.status.idle": "2024-03-31T16:28:19.574972Z",
     "shell.execute_reply": "2024-03-31T16:28:19.573575Z"
    },
    "papermill": {
     "duration": 0.016661,
     "end_time": "2024-03-31T16:28:19.577193",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.560532",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"KERAS_BACKEND\"] = \"torch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8097a3bb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:19.594948Z",
     "iopub.status.busy": "2024-03-31T16:28:19.594579Z",
     "iopub.status.idle": "2024-03-31T16:28:47.664966Z",
     "shell.execute_reply": "2024-03-31T16:28:47.663253Z"
    },
    "papermill": {
     "duration": 28.081995,
     "end_time": "2024-03-31T16:28:47.667322",
     "exception": false,
     "start_time": "2024-03-31T16:28:19.585327",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'keras-cv'...\r\n",
      "remote: Enumerating objects: 13766, done.\u001b[K\r\n",
      "remote: Counting objects: 100% (1902/1902), done.\u001b[K\r\n",
      "remote: Compressing objects: 100% (764/764), done.\u001b[K\r\n",
      "remote: Total 13766 (delta 1325), reused 1611 (delta 1122), pack-reused 11864\u001b[K\r\n",
      "Receiving objects: 100% (13766/13766), 25.66 MiB | 26.81 MiB/s, done.\r\n",
      "Resolving deltas: 100% (9760/9760), done.\r\n",
      "/kaggle/working/keras-cv\n"
     ]
    }
   ],
   "source": [
    "!git clone --branch video_swin https://github.com/innat/keras-cv.git\n",
    "%cd keras-cv\n",
    "!pip install -q -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9cd988de",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:28:47.688615Z",
     "iopub.status.busy": "2024-03-31T16:28:47.688157Z",
     "iopub.status.idle": "2024-03-31T16:29:06.875822Z",
     "shell.execute_reply": "2024-03-31T16:29:06.874231Z"
    },
    "papermill": {
     "duration": 19.200845,
     "end_time": "2024-03-31T16:29:06.877810",
     "exception": false,
     "start_time": "2024-03-31T16:28:47.676965",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-31 16:28:52.233542: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-03-31 16:28:52.233664: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-03-31 16:28:52.410813: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'3.0.5'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import keras\n",
    "from keras import ops\n",
    "from keras_cv.models import VideoSwinBackbone\n",
    "from keras_cv.models import VideoClassifier\n",
    "keras.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "663ffd78",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:06.899316Z",
     "iopub.status.busy": "2024-03-31T16:29:06.898673Z",
     "iopub.status.idle": "2024-03-31T16:29:08.582025Z",
     "shell.execute_reply": "2024-03-31T16:29:08.580526Z"
    },
    "papermill": {
     "duration": 1.696839,
     "end_time": "2024-03-31T16:29:08.584655",
     "exception": false,
     "start_time": "2024-03-31T16:29:06.887816",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-31 16:29:07--  https://github.com/innat/VideoSwin/releases/download/v2.0/videoswin_base_kinetics600_imagenet22k.weights.h5\r\n",
      "Resolving github.com (github.com)... 140.82.112.4\r\n",
      "Connecting to github.com (github.com)|140.82.112.4|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 302 Found\r\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/697696973/dde3749b-9dae-47e3-8b07-5ef1ef982a48?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162907Z&X-Amz-Expires=300&X-Amz-Signature=96eef6d9c80d98ae9c19cb6985f5c8511e36a16bc6e4d0f93f403d0a5cb5654c&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=697696973&response-content-disposition=attachment%3B%20filename%3Dvideoswin_base_kinetics600_imagenet22k.weights.h5&response-content-type=application%2Foctet-stream [following]\r\n",
      "--2024-03-31 16:29:07--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/697696973/dde3749b-9dae-47e3-8b07-5ef1ef982a48?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162907Z&X-Amz-Expires=300&X-Amz-Signature=96eef6d9c80d98ae9c19cb6985f5c8511e36a16bc6e4d0f93f403d0a5cb5654c&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=697696973&response-content-disposition=attachment%3B%20filename%3Dvideoswin_base_kinetics600_imagenet22k.weights.h5&response-content-type=application%2Foctet-stream\r\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.111.133, ...\r\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.110.133|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 200 OK\r\n",
      "Length: 351381896 (335M) [application/octet-stream]\r\n",
      "Saving to: 'videoswin_base_kinetics600_imagenet22k.weights.h5'\r\n",
      "\r\n",
      "videoswin_base_kine 100%[===================>] 335.10M   305MB/s    in 1.1s    \r\n",
      "\r\n",
      "2024-03-31 16:29:08 (305 MB/s) - 'videoswin_base_kinetics600_imagenet22k.weights.h5' saved [351381896/351381896]\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "!wget https://github.com/innat/VideoSwin/releases/download/v2.0/videoswin_base_kinetics600_imagenet22k.weights.h5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "db6e9815",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:08.607646Z",
     "iopub.status.busy": "2024-03-31T16:29:08.607201Z",
     "iopub.status.idle": "2024-03-31T16:29:08.614148Z",
     "shell.execute_reply": "2024-03-31T16:29:08.612947Z"
    },
    "papermill": {
     "duration": 0.021075,
     "end_time": "2024-03-31T16:29:08.616312",
     "exception": false,
     "start_time": "2024-03-31T16:29:08.595237",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def vswin_base():\n",
    "    backbone=VideoSwinBackbone(\n",
    "        input_shape=(32, 224, 224, 3), \n",
    "        embed_dim=128,\n",
    "        depths=[2, 2, 18, 2],\n",
    "        num_heads=[4, 8, 16, 32],\n",
    "        include_rescaling=False, \n",
    "    )\n",
    "    backbone.load_weights(\n",
    "        'videoswin_base_kinetics600_imagenet22k.weights.h5'\n",
    "    )\n",
    "    return backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5300dedb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:08.640284Z",
     "iopub.status.busy": "2024-03-31T16:29:08.639902Z",
     "iopub.status.idle": "2024-03-31T16:29:10.903494Z",
     "shell.execute_reply": "2024-03-31T16:29:10.901498Z"
    },
    "papermill": {
     "duration": 2.279472,
     "end_time": "2024-03-31T16:29:10.906342",
     "exception": false,
     "start_time": "2024-03-31T16:29:08.626870",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_ks = vswin_base()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "15599c39",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:10.928739Z",
     "iopub.status.busy": "2024-03-31T16:29:10.928268Z",
     "iopub.status.idle": "2024-03-31T16:29:10.936850Z",
     "shell.execute_reply": "2024-03-31T16:29:10.935945Z"
    },
    "papermill": {
     "duration": 0.021784,
     "end_time": "2024-03-31T16:29:10.938542",
     "exception": false,
     "start_time": "2024-03-31T16:29:10.916758",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch, number of params (M): 87.64\n",
      "Keras, number of params (M): 87.64\n"
     ]
    }
   ],
   "source": [
    "n_parameters = sum(p.numel() for p in model_pt.parameters() if p.requires_grad)\n",
    "print(\"PyTorch, number of params (M): %.2f\" % (n_parameters / 1.0e6))\n",
    "n_parameters = model_ks.count_params()\n",
    "print(\"Keras, number of params (M): %.2f\" % (n_parameters / 1.0e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5910a317",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:10.960763Z",
     "iopub.status.busy": "2024-03-31T16:29:10.960337Z",
     "iopub.status.idle": "2024-03-31T16:29:11.091219Z",
     "shell.execute_reply": "2024-03-31T16:29:11.089796Z"
    },
    "papermill": {
     "duration": 0.145028,
     "end_time": "2024-03-31T16:29:11.094030",
     "exception": false,
     "start_time": "2024-03-31T16:29:10.949002",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 224, 224, 3]) torch.Size([1, 3, 32, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "common_input = np.random.normal(0, 1, (1, 32, 224, 224, 3)).astype('float32')\n",
    "keras_input = ops.array(common_input)\n",
    "torch_input = torch.from_numpy(common_input.transpose(0, 4, 1, 2, 3))\n",
    "print(keras_input.shape, torch_input.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eacab339",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:11.122133Z",
     "iopub.status.busy": "2024-03-31T16:29:11.120966Z",
     "iopub.status.idle": "2024-03-31T16:29:11.129797Z",
     "shell.execute_reply": "2024-03-31T16:29:11.128454Z"
    },
    "papermill": {
     "duration": 0.026546,
     "end_time": "2024-03-31T16:29:11.132183",
     "exception": false,
     "start_time": "2024-03-31T16:29:11.105637",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def logit_checking(keras_model, torch_model):\n",
    "    # forward pass\n",
    "    keras_predict = keras_model(keras_input)\n",
    "    torch_predict = torch_model(torch_input).permute(0,2,3,4,1)\n",
    "    print(keras_predict.shape, torch_predict.shape)\n",
    "    np.testing.assert_allclose(\n",
    "        keras_predict.detach().numpy(),\n",
    "        torch_predict.detach().numpy(),\n",
    "        1e-4, 1e-4\n",
    "    )\n",
    "    del keras_model \n",
    "    del torch_model\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "12b00ce4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:11.155937Z",
     "iopub.status.busy": "2024-03-31T16:29:11.155564Z",
     "iopub.status.idle": "2024-03-31T16:29:31.486409Z",
     "shell.execute_reply": "2024-03-31T16:29:31.484870Z"
    },
    "papermill": {
     "duration": 20.346445,
     "end_time": "2024-03-31T16:29:31.489813",
     "exception": false,
     "start_time": "2024-03-31T16:29:11.143368",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 16, 7, 7, 1024]) torch.Size([1, 16, 7, 7, 1024])\n"
     ]
    }
   ],
   "source": [
    "logit_checking(\n",
    "    model_ks, model_pt\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39156fbb",
   "metadata": {
    "papermill": {
     "duration": 0.010275,
     "end_time": "2024-03-31T16:29:31.511215",
     "exception": false,
     "start_time": "2024-03-31T16:29:31.500940",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "221fce02",
   "metadata": {
    "papermill": {
     "duration": 0.010258,
     "end_time": "2024-03-31T16:29:31.532207",
     "exception": false,
     "start_time": "2024-03-31T16:29:31.521949",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Something Something V2 PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6714c913",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:31.555235Z",
     "iopub.status.busy": "2024-03-31T16:29:31.554885Z",
     "iopub.status.idle": "2024-03-31T16:29:32.602168Z",
     "shell.execute_reply": "2024-03-31T16:29:32.600775Z"
    },
    "papermill": {
     "duration": 1.062215,
     "end_time": "2024-03-31T16:29:32.604968",
     "exception": false,
     "start_time": "2024-03-31T16:29:31.542753",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_pt = video_swin_base(\n",
    "    window_size=(16,7,7), num_classes=174\n",
    ")\n",
    "model_pt.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "13c68885",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:32.632138Z",
     "iopub.status.busy": "2024-03-31T16:29:32.631743Z",
     "iopub.status.idle": "2024-03-31T16:29:35.072409Z",
     "shell.execute_reply": "2024-03-31T16:29:35.070658Z"
    },
    "papermill": {
     "duration": 2.456138,
     "end_time": "2024-03-31T16:29:35.074968",
     "exception": false,
     "start_time": "2024-03-31T16:29:32.618830",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-31 16:29:32--  https://github.com/SwinTransformer/storage/releases/download/v1.0.4/swin_base_patch244_window1677_sthv2.pth\r\n",
      "Resolving github.com (github.com)... 140.82.114.4\r\n",
      "Connecting to github.com (github.com)|140.82.114.4|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 302 Found\r\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/357198522/20458080-d55e-11eb-9021-4730e624e0ea?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162933Z&X-Amz-Expires=300&X-Amz-Signature=978c7c4b6a25743f9500f0ec50fc31c41ef41856d255c0a3dad8261388391329&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=357198522&response-content-disposition=attachment%3B%20filename%3Dswin_base_patch244_window1677_sthv2.pth&response-content-type=application%2Foctet-stream [following]\r\n",
      "--2024-03-31 16:29:33--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/357198522/20458080-d55e-11eb-9021-4730e624e0ea?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162933Z&X-Amz-Expires=300&X-Amz-Signature=978c7c4b6a25743f9500f0ec50fc31c41ef41856d255c0a3dad8261388391329&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=357198522&response-content-disposition=attachment%3B%20filename%3Dswin_base_patch244_window1677_sthv2.pth&response-content-type=application%2Foctet-stream\r\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.111.133, ...\r\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.110.133|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 200 OK\r\n",
      "Length: 473410081 (451M) [application/octet-stream]\r\n",
      "Saving to: 'checkpoint.pt'\r\n",
      "\r\n",
      "checkpoint.pt       100%[===================>] 451.48M   262MB/s    in 1.7s    \r\n",
      "\r\n",
      "2024-03-31 16:29:34 (262 MB/s) - 'checkpoint.pt' saved [473410081/473410081]\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "base_url = \"https://github.com/SwinTransformer/storage/releases/download/v1.0.4/\"\n",
    "checkpoints_pt = f\"{base_url}swin_base_patch244_window1677_sthv2.pth\"\n",
    "!wget {checkpoints_pt} -O checkpoint.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a923aa94",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:35.099503Z",
     "iopub.status.busy": "2024-03-31T16:29:35.099110Z",
     "iopub.status.idle": "2024-03-31T16:29:35.285266Z",
     "shell.execute_reply": "2024-03-31T16:29:35.283861Z"
    },
    "papermill": {
     "duration": 0.201555,
     "end_time": "2024-03-31T16:29:35.287934",
     "exception": false,
     "start_time": "2024-03-31T16:29:35.086379",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "state_dict = torch.load(\n",
    "    'checkpoint.pt', map_location=\"cpu\"\n",
    ")\n",
    "state_dict = state_dict['state_dict']\n",
    "state_dict = {k.replace('backbone.', ''): v for k, v in state_dict.items()}\n",
    "state_dict = {k.replace('cls_head.', ''): v for k, v in state_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "fe4eec0e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:35.313216Z",
     "iopub.status.busy": "2024-03-31T16:29:35.312863Z",
     "iopub.status.idle": "2024-03-31T16:29:35.412953Z",
     "shell.execute_reply": "2024-03-31T16:29:35.411952Z"
    },
    "papermill": {
     "duration": 0.115572,
     "end_time": "2024-03-31T16:29:35.415473",
     "exception": false,
     "start_time": "2024-03-31T16:29:35.299901",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "_IncompatibleKeys(missing_keys=[], unexpected_keys=['fc_cls.weight', 'fc_cls.bias'])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_pt.load_state_dict(state_dict, strict=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46daa259",
   "metadata": {
    "papermill": {
     "duration": 0.010798,
     "end_time": "2024-03-31T16:29:35.438752",
     "exception": false,
     "start_time": "2024-03-31T16:29:35.427954",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Something Somethinb V2 KerasCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "65566f83",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:35.462922Z",
     "iopub.status.busy": "2024-03-31T16:29:35.462545Z",
     "iopub.status.idle": "2024-03-31T16:29:37.242928Z",
     "shell.execute_reply": "2024-03-31T16:29:37.241555Z"
    },
    "papermill": {
     "duration": 1.795716,
     "end_time": "2024-03-31T16:29:37.245565",
     "exception": false,
     "start_time": "2024-03-31T16:29:35.449849",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-31 16:29:35--  https://github.com/innat/VideoSwin/releases/download/v2.0/videoswin_base_something_something_v2.weights.h5\r\n",
      "Resolving github.com (github.com)... 140.82.114.4\r\n",
      "Connecting to github.com (github.com)|140.82.114.4|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 302 Found\r\n",
      "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/697696973/9faf8cf2-0308-4a38-9022-0d227ece6073?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162935Z&X-Amz-Expires=300&X-Amz-Signature=d90abf874f38fe7b07b9464b483f46bb54da77d8dde6256b1efb369af77b1896&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=697696973&response-content-disposition=attachment%3B%20filename%3Dvideoswin_base_something_something_v2.weights.h5&response-content-type=application%2Foctet-stream [following]\r\n",
      "--2024-03-31 16:29:35--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/697696973/9faf8cf2-0308-4a38-9022-0d227ece6073?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240331%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240331T162935Z&X-Amz-Expires=300&X-Amz-Signature=d90abf874f38fe7b07b9464b483f46bb54da77d8dde6256b1efb369af77b1896&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=697696973&response-content-disposition=attachment%3B%20filename%3Dvideoswin_base_something_something_v2.weights.h5&response-content-type=application%2Foctet-stream\r\n",
      "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.110.133, 185.199.108.133, 185.199.111.133, ...\r\n",
      "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.110.133|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 200 OK\r\n",
      "Length: 355448712 (339M) [application/octet-stream]\r\n",
      "Saving to: 'videoswin_base_something_something_v2.weights.h5'\r\n",
      "\r\n",
      "videoswin_base_some 100%[===================>] 338.98M   313MB/s    in 1.1s    \r\n",
      "\r\n",
      "2024-03-31 16:29:37 (313 MB/s) - 'videoswin_base_something_something_v2.weights.h5' saved [355448712/355448712]\r\n",
      "\r\n"
     ]
    }
   ],
   "source": [
    "!wget https://github.com/innat/VideoSwin/releases/download/v2.0/videoswin_base_something_something_v2.weights.h5\n",
    "\n",
    "def vswin_base():\n",
    "    backbone=VideoSwinBackbone(\n",
    "        input_shape=(32, 224, 224, 3), \n",
    "        embed_dim=128,\n",
    "        depths=[2, 2, 18, 2],\n",
    "        num_heads=[4, 8, 16, 32],\n",
    "        window_size=[16, 7, 7],\n",
    "        include_rescaling=False, \n",
    "    )\n",
    "    backbone.load_weights(\n",
    "        'videoswin_base_something_something_v2.weights.h5'\n",
    "    )\n",
    "    return backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b74004a9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:37.273090Z",
     "iopub.status.busy": "2024-03-31T16:29:37.272409Z",
     "iopub.status.idle": "2024-03-31T16:29:39.469377Z",
     "shell.execute_reply": "2024-03-31T16:29:39.467493Z"
    },
    "papermill": {
     "duration": 2.213317,
     "end_time": "2024-03-31T16:29:39.472153",
     "exception": false,
     "start_time": "2024-03-31T16:29:37.258836",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_ks = vswin_base()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "836a50fe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-31T16:29:39.499550Z",
     "iopub.status.busy": "2024-03-31T16:29:39.499121Z",
     "iopub.status.idle": "2024-03-31T16:30:20.045949Z",
     "shell.execute_reply": "2024-03-31T16:30:20.043774Z"
    },
    "papermill": {
     "duration": 40.564324,
     "end_time": "2024-03-31T16:30:20.049374",
     "exception": false,
     "start_time": "2024-03-31T16:29:39.485050",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 16, 7, 7, 1024]) torch.Size([1, 16, 7, 7, 1024])\n"
     ]
    }
   ],
   "source": [
    "logit_checking(\n",
    "    model_ks, model_pt\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e72ae4bd",
   "metadata": {
    "papermill": {
     "duration": 0.012143,
     "end_time": "2024-03-31T16:30:20.074137",
     "exception": false,
     "start_time": "2024-03-31T16:30:20.061994",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "none",
   "dataSources": [
    {
     "modelInstanceId": 17533,
     "sourceId": 21184,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30673,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "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.10.13"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 153.194598,
   "end_time": "2024-03-31T16:30:23.699815",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-03-31T16:27:50.505217",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}