isl-recognition-with-transformers.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:34.519315Z",
     "iopub.status.busy": "2023-05-06T10:34:34.518597Z",
     "iopub.status.idle": "2023-05-06T10:34:34.544954Z",
     "shell.execute_reply": "2023-05-06T10:34:34.543924Z",
     "shell.execute_reply.started": "2023-05-06T10:34:34.519276Z"
    }
   },
   "outputs": [],
   "source": [
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load\n",
    "\n",
    "#import numpy as np # linear algebra\n",
    "#import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the read-only \"../input/\" directory\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n",
    "\n",
    "#import os\n",
    "#for dirname, _, filenames in os.walk('/kaggle/input'):\n",
    "    #for filename in filenames:\n",
    "        #print(os.path.join(dirname, filename))\n",
    "\n",
    "# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n",
    "# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:34.546977Z",
     "iopub.status.busy": "2023-05-06T10:34:34.546656Z",
     "iopub.status.idle": "2023-05-06T10:34:43.557798Z",
     "shell.execute_reply": "2023-05-06T10:34:43.556703Z",
     "shell.execute_reply.started": "2023-05-06T10:34:34.546947Z"
    }
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'wandb'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[1], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mjson\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mrandom\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mmath\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mscipy\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mwandb\u001b[39;00m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mnumpy\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnp\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mpandas\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mpd\u001b[39;00m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'wandb'"
     ]
    }
   ],
   "source": [
    "import os, json, random, math, scipy, wandb\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_addons as tfa\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.notebook import tqdm\n",
    "from sklearn.model_selection import StratifiedGroupKFold \n",
    "from types import SimpleNamespace\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:43.560199Z",
     "iopub.status.busy": "2023-05-06T10:34:43.559287Z",
     "iopub.status.idle": "2023-05-06T10:34:51.532349Z",
     "shell.execute_reply": "2023-05-06T10:34:51.531219Z",
     "shell.execute_reply.started": "2023-05-06T10:34:43.560159Z"
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'wandb' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m----> 2\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mkaggle_secrets\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m UserSecretsClient\n\u001b[0;32m      3\u001b[0m     user_secrets \u001b[38;5;241m=\u001b[39m UserSecretsClient()\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'kaggle_secrets'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 7\u001b[0m\n\u001b[0;32m      5\u001b[0m     wandb\u001b[38;5;241m.\u001b[39mlogin(key\u001b[38;5;241m=\u001b[39msecret_value_0)\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m----> 7\u001b[0m     wandb\u001b[38;5;241m.\u001b[39mlogin()\n",
      "\u001b[1;31mNameError\u001b[0m: name 'wandb' is not defined"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    from kaggle_secrets import UserSecretsClient\n",
    "    user_secrets = UserSecretsClient()\n",
    "    secret_value_0 = user_secrets.get_secret(\"82c0e9e68dd26ea25025b0f01592b0ed1c13c87b\")\n",
    "    wandb.login(key=secret_value_0)\n",
    "except:\n",
    "    wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.535479Z",
     "iopub.status.busy": "2023-05-06T10:34:51.535099Z",
     "iopub.status.idle": "2023-05-06T10:34:51.542206Z",
     "shell.execute_reply": "2023-05-06T10:34:51.540823Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.535436Z"
    }
   },
   "outputs": [],
   "source": [
    "# Source: https://www.kaggle.com/competitions/asl-signs/overview/evaluation\n",
    "ROWS_PER_FRAME = 543\n",
    "\n",
    "def load_relevant_data_subset(pq_path):\n",
    "    data_columns = ['x', 'y', 'z']\n",
    "    data = pd.read_parquet(pq_path, columns=data_columns)\n",
    "    n_frames = int(len(data) / ROWS_PER_FRAME)\n",
    "    data = data.values.reshape(n_frames, ROWS_PER_FRAME, len(data_columns))\n",
    "    return data.astype(np.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config and Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.544769Z",
     "iopub.status.busy": "2023-05-06T10:34:51.543973Z",
     "iopub.status.idle": "2023-05-06T10:34:51.555241Z",
     "shell.execute_reply": "2023-05-06T10:34:51.554226Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.544728Z"
    }
   },
   "outputs": [],
   "source": [
    "cfg = SimpleNamespace()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.557352Z",
     "iopub.status.busy": "2023-05-06T10:34:51.556945Z",
     "iopub.status.idle": "2023-05-06T10:34:51.570342Z",
     "shell.execute_reply": "2023-05-06T10:34:51.569362Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.557316Z"
    }
   },
   "outputs": [],
   "source": [
    "iskaggle = os.environ.get('KAGGLE_KERNEL_RUN_TYPE', '')\n",
    "iskaggle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.572411Z",
     "iopub.status.busy": "2023-05-06T10:34:51.572117Z",
     "iopub.status.idle": "2023-05-06T10:34:51.583927Z",
     "shell.execute_reply": "2023-05-06T10:34:51.582976Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.572385Z"
    }
   },
   "outputs": [],
   "source": [
    "DATA_DIR         = Path('../data/') if not iskaggle else Path('/kaggle/input/asl-signs/')\n",
    "TRAIN_CSV_PATH   = DATA_DIR/'train.csv'\n",
    "LANDMARK_DIR     = DATA_DIR/'train_landmark_files'\n",
    "LABEL_MAP_PATH   = DATA_DIR/'sign_to_prediction_index_map.json'\n",
    "DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.586225Z",
     "iopub.status.busy": "2023-05-06T10:34:51.585727Z",
     "iopub.status.idle": "2023-05-06T10:34:51.595635Z",
     "shell.execute_reply": "2023-05-06T10:34:51.594627Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.58619Z"
    }
   },
   "outputs": [],
   "source": [
    "cfg.PREPROCESS_DATA = False\n",
    "cfg.TRAIN_MODEL = True\n",
    "cfg.N_ROWS = 543\n",
    "cfg.N_DIMS = 3\n",
    "cfg.DIM_NAMES = ['x', 'y', 'z']\n",
    "cfg.SEED = 42\n",
    "cfg.NUM_CLASSES = 250\n",
    "cfg.IS_INTERACTIVE = True\n",
    "cfg.VERBOSE = 2\n",
    "cfg.INPUT_SIZE = 32\n",
    "cfg.BATCH_ALL_SIGNS_N = 4\n",
    "cfg.BATCH_SIZE = 256\n",
    "cfg.N_EPOCHS = 50\n",
    "cfg.LR_MAX = 1e-3\n",
    "cfg.N_WARMUP_EPOCHS = 0\n",
    "cfg.WD_RATIO = 0.05\n",
    "cfg.MASK_VAL = 4237"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.600949Z",
     "iopub.status.busy": "2023-05-06T10:34:51.600671Z",
     "iopub.status.idle": "2023-05-06T10:34:51.838805Z",
     "shell.execute_reply": "2023-05-06T10:34:51.837833Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.600924Z"
    }
   },
   "outputs": [],
   "source": [
    "train = pd.read_csv(TRAIN_CSV_PATH)\n",
    "N_SAMPLES = len(train)\n",
    "N_SAMPLES"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Below, instead of updating the path on the train file, we create a symlink instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:51.850137Z",
     "iopub.status.busy": "2023-05-06T10:34:51.844503Z",
     "iopub.status.idle": "2023-05-06T10:34:52.877066Z",
     "shell.execute_reply": "2023-05-06T10:34:52.875755Z",
     "shell.execute_reply.started": "2023-05-06T10:34:51.850073Z"
    }
   },
   "outputs": [],
   "source": [
    "# Get complete file path to file\n",
    "def get_file_path(path):\n",
    "    return f'/kaggle/input/asl-signs/{path}'\n",
    "\n",
    "!ln -s {LANDMARK_DIR} ./train_landmark_files\n",
    "train['file_path'] = train['path'].values\n",
    "train.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:52.879593Z",
     "iopub.status.busy": "2023-05-06T10:34:52.879196Z",
     "iopub.status.idle": "2023-05-06T10:34:52.906816Z",
     "shell.execute_reply": "2023-05-06T10:34:52.905716Z",
     "shell.execute_reply.started": "2023-05-06T10:34:52.879546Z"
    }
   },
   "outputs": [],
   "source": [
    "with open(LABEL_MAP_PATH, 'r') as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "# Create a dictionary to map the indexes to their corresponding values\n",
    "index_to_value = {str(index): value for index, value in data.items()}\n",
    "\n",
    "# add a new column with the corresponding keys using map\n",
    "train['sign_ord'] = train['sign'].map(index_to_value)\n",
    "train.head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Kfold"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we create K-folds using StratifiedGroupKFold, the idea is to have different participants in\n",
    "train and validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:52.908963Z",
     "iopub.status.busy": "2023-05-06T10:34:52.908566Z",
     "iopub.status.idle": "2023-05-06T10:34:52.913922Z",
     "shell.execute_reply": "2023-05-06T10:34:52.912676Z",
     "shell.execute_reply.started": "2023-05-06T10:34:52.908926Z"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import StratifiedGroupKFold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:52.916557Z",
     "iopub.status.busy": "2023-05-06T10:34:52.915741Z",
     "iopub.status.idle": "2023-05-06T10:34:53.222347Z",
     "shell.execute_reply": "2023-05-06T10:34:53.221288Z",
     "shell.execute_reply.started": "2023-05-06T10:34:52.91652Z"
    }
   },
   "outputs": [],
   "source": [
    "N_PARTICIPANTS = train.participant_id.nunique()\n",
    "sgkf = StratifiedGroupKFold(n_splits=7, shuffle=True, random_state=43)\n",
    "train['fold'] = -1\n",
    "for i, (train_idx, val_idx) in enumerate(sgkf.split(train.index, train.sign, train.participant_id)):\n",
    "    train.loc[val_idx, 'fold'] = i\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.224449Z",
     "iopub.status.busy": "2023-05-06T10:34:53.224041Z",
     "iopub.status.idle": "2023-05-06T10:34:53.255805Z",
     "shell.execute_reply": "2023-05-06T10:34:53.254879Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.224407Z"
    }
   },
   "outputs": [],
   "source": [
    "# create indexes using fold `0` for now\n",
    "train_idxs = train.query(\"fold!=0\").index.values\n",
    "val_idxs = train.query(\"fold==0\").index.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.257796Z",
     "iopub.status.busy": "2023-05-06T10:34:53.257216Z",
     "iopub.status.idle": "2023-05-06T10:34:53.368087Z",
     "shell.execute_reply": "2023-05-06T10:34:53.367092Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.257742Z"
    }
   },
   "outputs": [],
   "source": [
    "data = pd.read_parquet(train.loc[0, \"file_path\"])\n",
    "data[data['frame'] == 20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process Data Tensorflow\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.370165Z",
     "iopub.status.busy": "2023-05-06T10:34:53.369533Z",
     "iopub.status.idle": "2023-05-06T10:34:53.381929Z",
     "shell.execute_reply": "2023-05-06T10:34:53.380712Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.370125Z"
    }
   },
   "outputs": [],
   "source": [
    "# landmark indices in original data\n",
    "LIPS_IDXS0 = np.array([\n",
    "        61, 185, 40, 39, 37, 0, 267, 269, 270, 409,\n",
    "        291, 146, 91, 181, 84, 17, 314, 405, 321, 375,\n",
    "        78, 191, 80, 81, 82, 13, 312, 311, 310, 415,\n",
    "        95, 88, 178, 87, 14, 317, 402, 318, 324, 308,\n",
    "    ])\n",
    "LEFT_HAND_IDXS0  = np.arange(468,489)\n",
    "RIGHT_HAND_IDXS0 = np.arange(522,543)\n",
    "POSE_IDXS0       = np.arange(502, 512)\n",
    "LANDMARK_IDXS0   = np.concatenate((LIPS_IDXS0, LEFT_HAND_IDXS0, RIGHT_HAND_IDXS0, POSE_IDXS0))\n",
    "HAND_IDXS0       = np.concatenate((LEFT_HAND_IDXS0, RIGHT_HAND_IDXS0), axis=0)\n",
    "N_COLS           = LANDMARK_IDXS0.size\n",
    "N_COLS, LANDMARK_IDXS0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.384241Z",
     "iopub.status.busy": "2023-05-06T10:34:53.383177Z",
     "iopub.status.idle": "2023-05-06T10:34:53.396951Z",
     "shell.execute_reply": "2023-05-06T10:34:53.396071Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.384203Z"
    }
   },
   "outputs": [],
   "source": [
    "# Landmark indices in processed data\n",
    "LIPS_IDXS       = np.argwhere(np.isin(LANDMARK_IDXS0, LIPS_IDXS0)).squeeze()\n",
    "LEFT_HAND_IDXS  = np.argwhere(np.isin(LANDMARK_IDXS0, LEFT_HAND_IDXS0)).squeeze()\n",
    "RIGHT_HAND_IDXS = np.argwhere(np.isin(LANDMARK_IDXS0, RIGHT_HAND_IDXS0)).squeeze()\n",
    "HAND_IDXS       = np.argwhere(np.isin(LANDMARK_IDXS0, HAND_IDXS0)).squeeze()\n",
    "POSE_IDXS       = np.argwhere(np.isin(LANDMARK_IDXS0, POSE_IDXS0)).squeeze()\n",
    "HAND_IDXS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.399228Z",
     "iopub.status.busy": "2023-05-06T10:34:53.398502Z",
     "iopub.status.idle": "2023-05-06T10:34:53.428813Z",
     "shell.execute_reply": "2023-05-06T10:34:53.427831Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.399191Z"
    }
   },
   "outputs": [],
   "source": [
    "class PreprocessLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self):\n",
    "        super(PreprocessLayer, self).__init__()\n",
    "        \n",
    "    def pad_edge(self, t, repeats, side):\n",
    "        if side == 'LEFT':\n",
    "            return tf.concat((tf.repeat(t[:1], repeats=repeats, axis=0), t), axis=0)\n",
    "        elif side == 'RIGHT':\n",
    "            return tf.concat((t, tf.repeat(t[-1:], repeats=repeats, axis=0)), axis=0)\n",
    "    \n",
    "    @tf.function(\n",
    "        input_signature=(tf.TensorSpec(shape=[None,cfg.N_ROWS,cfg.N_DIMS], dtype=tf.float32),),\n",
    "    )\n",
    "    def call(self, data0):\n",
    "        # Number of Frames in Video\n",
    "        N_FRAMES0 = tf.shape(data0)[0]\n",
    "        \n",
    "        # Keep only non-empty frames in data\n",
    "        frames_hands_nansum = tf.experimental.numpy.nanmean(tf.gather(data0, HAND_IDXS0, axis=1), axis=[1,2])\n",
    "        non_empty_frames_idxs = tf.where(frames_hands_nansum > 0)\n",
    "        non_empty_frames_idxs = tf.squeeze(non_empty_frames_idxs, axis=1)\n",
    "        data = tf.gather(data0, non_empty_frames_idxs, axis=0)\n",
    "        \n",
    "        non_empty_frames_idxs = tf.cast(non_empty_frames_idxs, tf.float32) \n",
    "        \n",
    "        # Number of non-empty frames\n",
    "        N_FRAMES = tf.shape(data)[0]\n",
    "        data = tf.gather(data, LANDMARK_IDXS0, axis=1)\n",
    "        \n",
    "        if N_FRAMES < cfg.INPUT_SIZE:\n",
    "            # Video fits in cfg.INPUT_SIZE\n",
    "            non_empty_frames_idxs = tf.pad(non_empty_frames_idxs, [[0, cfg.INPUT_SIZE-N_FRAMES]], constant_values=-1)\n",
    "            data = tf.pad(data, [[0, cfg.INPUT_SIZE-N_FRAMES], [0,0], [0,0]], constant_values=0)\n",
    "            data = tf.where(tf.math.is_nan(data), 0.0, data)\n",
    "            return data, non_empty_frames_idxs\n",
    "        else:\n",
    "            # Video needs to be downsampled to cfg.INPUT_SIZE\n",
    "            if N_FRAMES < cfg.INPUT_SIZE**2:\n",
    "                repeats = tf.math.floordiv(cfg.INPUT_SIZE * cfg.INPUT_SIZE, N_FRAMES0)\n",
    "                data = tf.repeat(data, repeats=repeats, axis=0)\n",
    "                non_empty_frames_idxs = tf.repeat(non_empty_frames_idxs, repeats=repeats, axis=0)\n",
    "\n",
    "            # Pad To Multiple Of Input Size\n",
    "            pool_size = tf.math.floordiv(len(data), cfg.INPUT_SIZE)\n",
    "            if tf.math.mod(len(data), cfg.INPUT_SIZE) > 0:\n",
    "                pool_size += 1\n",
    "            if pool_size == 1:\n",
    "                pad_size = (pool_size * cfg.INPUT_SIZE) - len(data)\n",
    "            else:\n",
    "                pad_size = (pool_size * cfg.INPUT_SIZE) % len(data)\n",
    "\n",
    "            # Pad Start/End with Start/End value\n",
    "            pad_left = tf.math.floordiv(pad_size, 2) + tf.math.floordiv(cfg.INPUT_SIZE, 2)\n",
    "            pad_right = tf.math.floordiv(pad_size, 2) + tf.math.floordiv(cfg.INPUT_SIZE, 2)\n",
    "            if tf.math.mod(pad_size, 2) > 0:\n",
    "                pad_right += 1\n",
    "\n",
    "            # Pad By Concatenating Left/Right Edge Values\n",
    "            data = self.pad_edge(data, pad_left, 'LEFT')\n",
    "            data = self.pad_edge(data, pad_right, 'RIGHT')\n",
    "\n",
    "            # Pad Non Empty Frame Indices\n",
    "            non_empty_frames_idxs = self.pad_edge(non_empty_frames_idxs, pad_left, 'LEFT')\n",
    "            non_empty_frames_idxs = self.pad_edge(non_empty_frames_idxs, pad_right, 'RIGHT')\n",
    "\n",
    "            # Reshape to Mean Pool\n",
    "            data = tf.reshape(data, [cfg.INPUT_SIZE, -1, N_COLS, cfg.N_DIMS])\n",
    "            non_empty_frames_idxs = tf.reshape(non_empty_frames_idxs, [cfg.INPUT_SIZE, -1])\n",
    "\n",
    "            # Mean Pool\n",
    "            data = tf.experimental.numpy.nanmean(data, axis=1)\n",
    "            non_empty_frames_idxs = tf.experimental.numpy.nanmean(non_empty_frames_idxs, axis=1)\n",
    "\n",
    "            # Fill NaN Values With 0\n",
    "            data = tf.where(tf.math.is_nan(data), 0.0, data)\n",
    "            \n",
    "            return data, non_empty_frames_idxs\n",
    "    \n",
    "preprocess_layer = PreprocessLayer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.431408Z",
     "iopub.status.busy": "2023-05-06T10:34:53.430601Z",
     "iopub.status.idle": "2023-05-06T10:34:53.46645Z",
     "shell.execute_reply": "2023-05-06T10:34:53.465823Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.431371Z"
    }
   },
   "outputs": [],
   "source": [
    "sample = load_relevant_data_subset(train.path[2])\n",
    "sample.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:53.468493Z",
     "iopub.status.busy": "2023-05-06T10:34:53.467541Z",
     "iopub.status.idle": "2023-05-06T10:34:57.014792Z",
     "shell.execute_reply": "2023-05-06T10:34:57.013773Z",
     "shell.execute_reply.started": "2023-05-06T10:34:53.468456Z"
    }
   },
   "outputs": [],
   "source": [
    "data, non_empty_frames_idxs = preprocess_layer(sample)\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:57.016865Z",
     "iopub.status.busy": "2023-05-06T10:34:57.01648Z",
     "iopub.status.idle": "2023-05-06T10:34:57.024162Z",
     "shell.execute_reply": "2023-05-06T10:34:57.023189Z",
     "shell.execute_reply.started": "2023-05-06T10:34:57.016827Z"
    }
   },
   "outputs": [],
   "source": [
    "LIPS_START = 0\n",
    "LEFT_HAND_START = LIPS_IDXS.size\n",
    "RIGHT_HAND_START = LEFT_HAND_START + LEFT_HAND_IDXS.size\n",
    "POSE_START = RIGHT_HAND_START + RIGHT_HAND_IDXS.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:57.026442Z",
     "iopub.status.busy": "2023-05-06T10:34:57.026029Z",
     "iopub.status.idle": "2023-05-06T10:34:57.033958Z",
     "shell.execute_reply": "2023-05-06T10:34:57.032844Z",
     "shell.execute_reply.started": "2023-05-06T10:34:57.026404Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_data(file_path):\n",
    "    data = load_relevant_data_subset(file_path)\n",
    "    data = preprocess_layer(data)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:57.036286Z",
     "iopub.status.busy": "2023-05-06T10:34:57.035472Z",
     "iopub.status.idle": "2023-05-06T10:34:57.049111Z",
     "shell.execute_reply": "2023-05-06T10:34:57.048125Z",
     "shell.execute_reply.started": "2023-05-06T10:34:57.03625Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_x_y():\n",
    "    # Create arrays to save data\n",
    "    X = np.zeros([N_SAMPLES, cfg.INPUT_SIZE, N_COLS, cfg.N_DIMS], dtype=np.float32)\n",
    "    y = np.zeros([N_SAMPLES], dtype=np.int32)\n",
    "    NON_EMPTY_FRAME_IDXS = np.full([N_SAMPLES, cfg.INPUT_SIZE], -1, dtype=np.float32)\n",
    "\n",
    "    for row_idx, (file_path, sign_ord) in enumerate(tqdm(train[['file_path', 'sign_ord']].values)):\n",
    "        if row_idx % 5000 == 0:\n",
    "            print(f'Generated {row_idx}/{N_SAMPLES}')\n",
    "\n",
    "        data, non_empty_frame_idxs = get_data(file_path)\n",
    "        X[row_idx] = data\n",
    "        y[row_idx] = sign_ord\n",
    "        NON_EMPTY_FRAME_IDXS[row_idx] = non_empty_frame_idxs\n",
    "        if np.isnan(data).sum() > 0: return data\n",
    "\n",
    "    # Save X/y\n",
    "    np.save('X.npy', X)\n",
    "    np.save('y.npy', y)\n",
    "    np.save('NON_EMPTY_FRAME_IDXS.npy', NON_EMPTY_FRAME_IDXS)\n",
    "    \n",
    "    return X, y, NON_EMPTY_FRAME_IDXS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T10:34:57.05629Z",
     "iopub.status.busy": "2023-05-06T10:34:57.056012Z",
     "iopub.status.idle": "2023-05-06T11:04:47.584161Z",
     "shell.execute_reply": "2023-05-06T11:04:47.583013Z",
     "shell.execute_reply.started": "2023-05-06T10:34:57.056255Z"
    }
   },
   "outputs": [],
   "source": [
    "if not cfg.PREPROCESS_DATA:\n",
    "    X, y, NON_EMPTY_FRAME_IDXS = get_x_y()\n",
    "else:\n",
    "    X = np.load('/kaggle/working/X.npy')\n",
    "    y = np.load('/kaggle/working/y.npy')\n",
    "    NON_EMPTY_FRAME_IDXS = np.load('/kaggle/working/NON_EMPTY_FRAME_IDXS.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:04:47.588556Z",
     "iopub.status.busy": "2023-05-06T11:04:47.586676Z",
     "iopub.status.idle": "2023-05-06T11:04:47.597548Z",
     "shell.execute_reply": "2023-05-06T11:04:47.596425Z",
     "shell.execute_reply.started": "2023-05-06T11:04:47.58847Z"
    }
   },
   "outputs": [],
   "source": [
    "X.shape, y.shape, NON_EMPTY_FRAME_IDXS.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature Statistics - Lips\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:04:47.600018Z",
     "iopub.status.busy": "2023-05-06T11:04:47.59964Z",
     "iopub.status.idle": "2023-05-06T11:04:58.331927Z",
     "shell.execute_reply": "2023-05-06T11:04:58.330848Z",
     "shell.execute_reply.started": "2023-05-06T11:04:47.599979Z"
    }
   },
   "outputs": [],
   "source": [
    "# LIPS\n",
    "LIPS_MEAN_X  = np.zeros([LIPS_IDXS.size], dtype=np.float32)\n",
    "LIPS_MEAN_Y  = np.zeros([LIPS_IDXS.size], dtype=np.float32)\n",
    "LIPS_STD_X   = np.zeros([LIPS_IDXS.size], dtype=np.float32)\n",
    "LIPS_STD_Y   = np.zeros([LIPS_IDXS.size], dtype=np.float32)\n",
    "\n",
    "for col, ll in enumerate(tqdm( np.transpose(X[:,:,LIPS_IDXS], [2,3,0,1]).reshape([LIPS_IDXS.size, cfg.N_DIMS, -1]) )):\n",
    "    for dim, l in enumerate(ll):\n",
    "        v = l[np.nonzero(l)]\n",
    "        if dim == 0: # X\n",
    "            LIPS_MEAN_X[col] = v.mean()\n",
    "            LIPS_STD_X[col] = v.std()\n",
    "        if dim == 1: # Y\n",
    "            LIPS_MEAN_Y[col] = v.mean()\n",
    "            LIPS_STD_Y[col] = v.std()\n",
    "        \n",
    "LIPS_MEAN = np.array([LIPS_MEAN_X, LIPS_MEAN_Y]).T\n",
    "LIPS_STD = np.array([LIPS_STD_X, LIPS_STD_Y]).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature Statistics - Hands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:04:58.334307Z",
     "iopub.status.busy": "2023-05-06T11:04:58.333328Z",
     "iopub.status.idle": "2023-05-06T11:05:08.178963Z",
     "shell.execute_reply": "2023-05-06T11:05:08.177845Z",
     "shell.execute_reply.started": "2023-05-06T11:04:58.334267Z"
    }
   },
   "outputs": [],
   "source": [
    "# LEFT HAND\n",
    "LEFT_HANDS_MEAN_X = np.zeros([LEFT_HAND_IDXS.size], dtype=np.float32)\n",
    "LEFT_HANDS_MEAN_Y = np.zeros([LEFT_HAND_IDXS.size], dtype=np.float32)\n",
    "LEFT_HANDS_STD_X = np.zeros([LEFT_HAND_IDXS.size], dtype=np.float32)\n",
    "LEFT_HANDS_STD_Y = np.zeros([LEFT_HAND_IDXS.size], dtype=np.float32)\n",
    "# RIGHT HAND\n",
    "RIGHT_HANDS_MEAN_X = np.zeros([RIGHT_HAND_IDXS.size], dtype=np.float32)\n",
    "RIGHT_HANDS_MEAN_Y = np.zeros([RIGHT_HAND_IDXS.size], dtype=np.float32)\n",
    "RIGHT_HANDS_STD_X = np.zeros([RIGHT_HAND_IDXS.size], dtype=np.float32)\n",
    "RIGHT_HANDS_STD_Y = np.zeros([RIGHT_HAND_IDXS.size], dtype=np.float32)\n",
    "\n",
    "for col, ll in enumerate(tqdm( np.transpose(X[:,:,HAND_IDXS], [2,3,0,1]).reshape([HAND_IDXS.size, cfg.N_DIMS, -1]) )):\n",
    "    for dim, l in enumerate(ll):\n",
    "        v = l[np.nonzero(l)]\n",
    "        if dim == 0: # X\n",
    "            if col < RIGHT_HAND_IDXS.size: # LEFT HAND\n",
    "                LEFT_HANDS_MEAN_X[col] = v.mean()\n",
    "                LEFT_HANDS_STD_X[col] = v.std()\n",
    "            else:\n",
    "                RIGHT_HANDS_MEAN_X[col - LEFT_HAND_IDXS.size] = v.mean()\n",
    "                RIGHT_HANDS_STD_X[col - LEFT_HAND_IDXS.size] = v.std()\n",
    "        if dim == 1: # Y\n",
    "            if col < RIGHT_HAND_IDXS.size: # LEFT HAND\n",
    "                LEFT_HANDS_MEAN_Y[col] = v.mean()\n",
    "                LEFT_HANDS_STD_Y[col] = v.std()\n",
    "            else: # RIGHT HAND\n",
    "                RIGHT_HANDS_MEAN_Y[col - LEFT_HAND_IDXS.size] = v.mean()\n",
    "                RIGHT_HANDS_STD_Y[col - LEFT_HAND_IDXS.size] = v.std()\n",
    "        \n",
    "LEFT_HANDS_MEAN = np.array([LEFT_HANDS_MEAN_X, LEFT_HANDS_MEAN_Y]).T\n",
    "LEFT_HANDS_STD = np.array([LEFT_HANDS_STD_X, LEFT_HANDS_STD_Y]).T\n",
    "RIGHT_HANDS_MEAN = np.array([RIGHT_HANDS_MEAN_X, RIGHT_HANDS_MEAN_Y]).T\n",
    "RIGHT_HANDS_STD = np.array([RIGHT_HANDS_STD_X, RIGHT_HANDS_STD_Y]).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature Statistics - Pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:08.181239Z",
     "iopub.status.busy": "2023-05-06T11:05:08.180576Z",
     "iopub.status.idle": "2023-05-06T11:05:10.171142Z",
     "shell.execute_reply": "2023-05-06T11:05:10.169843Z",
     "shell.execute_reply.started": "2023-05-06T11:05:08.181197Z"
    }
   },
   "outputs": [],
   "source": [
    "# POSE\n",
    "POSE_MEAN_X = np.zeros([POSE_IDXS.size], dtype=np.float32)\n",
    "POSE_MEAN_Y = np.zeros([POSE_IDXS.size], dtype=np.float32)\n",
    "POSE_STD_X = np.zeros([POSE_IDXS.size], dtype=np.float32)\n",
    "POSE_STD_Y = np.zeros([POSE_IDXS.size], dtype=np.float32)\n",
    "\n",
    "for col, ll in enumerate(tqdm( np.transpose(X[:,:,POSE_IDXS], [2,3,0,1]).reshape([POSE_IDXS.size, cfg.N_DIMS, -1]) )):\n",
    "    for dim, l in enumerate(ll):\n",
    "        v = l[np.nonzero(l)]\n",
    "        if dim == 0: # X\n",
    "            POSE_MEAN_X[col] = v.mean()\n",
    "            POSE_STD_X[col] = v.std()\n",
    "        if dim == 1: # Y\n",
    "            POSE_MEAN_Y[col] = v.mean()\n",
    "            POSE_STD_Y[col] = v.std()\n",
    "        \n",
    "POSE_MEAN = np.array([POSE_MEAN_X, POSE_MEAN_Y]).T\n",
    "POSE_STD = np.array([POSE_STD_X, POSE_STD_Y]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.17338Z",
     "iopub.status.busy": "2023-05-06T11:05:10.172986Z",
     "iopub.status.idle": "2023-05-06T11:05:10.185532Z",
     "shell.execute_reply": "2023-05-06T11:05:10.18446Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.173339Z"
    }
   },
   "outputs": [],
   "source": [
    "# Custom sampler to get a batch containing N times all signs\n",
    "def get_train_batch_all_signs(X, y, NON_EMPTY_FRAME_IDXS, n=cfg.BATCH_ALL_SIGNS_N):\n",
    "    # Arrays to store batch in\n",
    "    X_batch = np.zeros([cfg.NUM_CLASSES*n, cfg.INPUT_SIZE, N_COLS, cfg.N_DIMS], dtype=np.float32)\n",
    "    y_batch = np.arange(0, cfg.NUM_CLASSES, step=1/n, dtype=np.float32).astype(np.int64)\n",
    "    non_empty_frame_idxs_batch = np.zeros([cfg.NUM_CLASSES*n, cfg.INPUT_SIZE], dtype=np.float32)\n",
    "    \n",
    "    # Dictionary mapping ordinally encoded sign to corresponding sample indices\n",
    "    CLASS2IDXS = {}\n",
    "    for i in range(cfg.NUM_CLASSES):\n",
    "        CLASS2IDXS[i] = np.argwhere(y == i).squeeze().astype(np.int32)\n",
    "            \n",
    "    while True:\n",
    "        # Fill batch arrays\n",
    "        for i in range(cfg.NUM_CLASSES):\n",
    "            idxs = np.random.choice(CLASS2IDXS[i], n)\n",
    "            X_batch[i*n:(i+1)*n] = X[idxs]\n",
    "            non_empty_frame_idxs_batch[i*n:(i+1)*n] = NON_EMPTY_FRAME_IDXS[idxs]\n",
    "        \n",
    "        yield { 'frames': X_batch, 'non_empty_frame_idxs': non_empty_frame_idxs_batch }, y_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.187661Z",
     "iopub.status.busy": "2023-05-06T11:05:10.186969Z",
     "iopub.status.idle": "2023-05-06T11:05:10.248808Z",
     "shell.execute_reply": "2023-05-06T11:05:10.247692Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.187621Z"
    }
   },
   "outputs": [],
   "source": [
    "dummy_dataset = get_train_batch_all_signs(X, y, NON_EMPTY_FRAME_IDXS)\n",
    "X_batch, y_batch = next(dummy_dataset)\n",
    "X_batch.keys(), X_batch['frames'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.250899Z",
     "iopub.status.busy": "2023-05-06T11:05:10.250362Z",
     "iopub.status.idle": "2023-05-06T11:05:10.257644Z",
     "shell.execute_reply": "2023-05-06T11:05:10.256635Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.250856Z"
    }
   },
   "outputs": [],
   "source": [
    "# Epsilon value for layer normalisation\n",
    "LAYER_NORM_EPS = 1e-6\n",
    "\n",
    "# Dense layer units for landmarks\n",
    "LIPS_UNITS = 384\n",
    "HANDS_UNITS = 384\n",
    "POSE_UNITS = 384\n",
    "# final embedding and transformer embedding size\n",
    "UNITS = 384\n",
    "\n",
    "# Transformer\n",
    "NUM_BLOCKS = 2\n",
    "MLP_RATIO = 2\n",
    "\n",
    "# Dropout\n",
    "EMBEDDING_DROPOUT = 0.00\n",
    "MLP_DROPOUT_RATIO = 0.30\n",
    "CLASSIFIER_DROPOUT_RATIO = 0.10\n",
    "\n",
    "# Initiailizers\n",
    "INIT_HE_UNIFORM = tf.keras.initializers.he_uniform\n",
    "INIT_GLOROT_UNIFORM = tf.keras.initializers.glorot_uniform\n",
    "INIT_ZEROS = tf.keras.initializers.constant(0.0)\n",
    "# Activations\n",
    "GELU = tf.keras.activations.gelu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer\n",
    "\n",
    "Need to implement transformer from scratch as TFLite does not support the native TF implementation of MultiHeadAttention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.259958Z",
     "iopub.status.busy": "2023-05-06T11:05:10.259272Z",
     "iopub.status.idle": "2023-05-06T11:05:10.278477Z",
     "shell.execute_reply": "2023-05-06T11:05:10.277324Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.259918Z"
    }
   },
   "outputs": [],
   "source": [
    "def scaled_dot_product(q,k,v, softmax, attention_mask):\n",
    "    #calculates Q . K(transpose)\n",
    "    qkt = tf.matmul(q,k,transpose_b=True)\n",
    "    #caculates scaling factor\n",
    "    dk = tf.math.sqrt(tf.cast(q.shape[-1],dtype=tf.float32))\n",
    "    scaled_qkt = qkt/dk\n",
    "    softmax = softmax(scaled_qkt, mask=attention_mask)\n",
    "    \n",
    "    z = tf.matmul(softmax,v)\n",
    "    #shape: (m,Tx,depth), same shape as q,k,v\n",
    "    return z\n",
    "\n",
    "class MultiHeadAttention(tf.keras.layers.Layer):\n",
    "    def __init__(self,d_model,num_of_heads):\n",
    "        super(MultiHeadAttention,self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.num_of_heads = num_of_heads\n",
    "        self.depth = d_model//num_of_heads\n",
    "        self.wq = [tf.keras.layers.Dense(self.depth) for i in range(num_of_heads)]\n",
    "        self.wk = [tf.keras.layers.Dense(self.depth) for i in range(num_of_heads)]\n",
    "        self.wv = [tf.keras.layers.Dense(self.depth) for i in range(num_of_heads)]\n",
    "        self.wo = tf.keras.layers.Dense(d_model)\n",
    "        self.softmax = tf.keras.layers.Softmax()\n",
    "        \n",
    "    def call(self,x, attention_mask):\n",
    "        \n",
    "        multi_attn = []\n",
    "        for i in range(self.num_of_heads):\n",
    "            Q = self.wq[i](x)\n",
    "            K = self.wk[i](x)\n",
    "            V = self.wv[i](x)\n",
    "            multi_attn.append(scaled_dot_product(Q,K,V, self.softmax, attention_mask))\n",
    "            \n",
    "        multi_head = tf.concat(multi_attn,axis=-1)\n",
    "        multi_head_attention = self.wo(multi_head)\n",
    "        return multi_head_attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.283559Z",
     "iopub.status.busy": "2023-05-06T11:05:10.28276Z",
     "iopub.status.idle": "2023-05-06T11:05:10.306505Z",
     "shell.execute_reply": "2023-05-06T11:05:10.302279Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.283514Z"
    }
   },
   "outputs": [],
   "source": [
    "class Transformer(tf.keras.Model):\n",
    "    def __init__(self, num_blocks):\n",
    "        super(Transformer, self).__init__(name='transformer')\n",
    "        self.num_blocks = num_blocks\n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        self.ln_1s = []\n",
    "        self.mhas = []\n",
    "        self.ln_2s = []\n",
    "        self.mlps = []\n",
    "        # Make Transformer Blocks\n",
    "        for i in range(self.num_blocks):\n",
    "            # First Layer Normalisation\n",
    "            self.ln_1s.append(tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS))\n",
    "            # Multi Head Attention\n",
    "            self.mhas.append(MultiHeadAttention(UNITS, 8))\n",
    "            # Second Layer Normalisation\n",
    "            self.ln_2s.append(tf.keras.layers.LayerNormalization(epsilon=LAYER_NORM_EPS))\n",
    "            # Multi Layer Perception\n",
    "            self.mlps.append(tf.keras.Sequential([\n",
    "                tf.keras.layers.Dense(UNITS * MLP_RATIO, activation=GELU, kernel_initializer=INIT_GLOROT_UNIFORM),\n",
    "                tf.keras.layers.Dropout(MLP_DROPOUT_RATIO),\n",
    "                tf.keras.layers.Dense(UNITS, kernel_initializer=INIT_HE_UNIFORM),\n",
    "            ]))\n",
    "        \n",
    "    def call(self, x, attention_mask):\n",
    "        # Iterate input over transformer blocks\n",
    "        for ln_1, mha, ln_2, mlp in zip(self.ln_1s, self.mhas, self.ln_2s, self.mlps):\n",
    "            x1 = ln_1(x)\n",
    "            attention_output = mha(x1, attention_mask)\n",
    "            x2 = x1 + attention_output\n",
    "            x3 = ln_2(x2)\n",
    "            x3 = mlp(x3)\n",
    "            x = x3 + x2\n",
    "    \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Landmark Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.311171Z",
     "iopub.status.busy": "2023-05-06T11:05:10.310404Z",
     "iopub.status.idle": "2023-05-06T11:05:10.329288Z",
     "shell.execute_reply": "2023-05-06T11:05:10.328209Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.311119Z"
    }
   },
   "outputs": [],
   "source": [
    "class LandmarkEmbedding(tf.keras.Model):\n",
    "    def __init__(self, units, name):\n",
    "        super(LandmarkEmbedding, self).__init__(name=f'{name}_embedding')\n",
    "        self.units = units\n",
    "        \n",
    "    def build(self, input_shape):\n",
    "        # Embedding for missing landmark in frame, initizlied with zeros\n",
    "        self.empty_embedding = self.add_weight(\n",
    "            name=f'{self.name}_empty_embedding',\n",
    "            shape=[self.units],\n",
    "            initializer=INIT_ZEROS,\n",
    "        )\n",
    "        # Embedding\n",
    "        self.dense = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(self.units, name=f'{self.name}_dense_1', use_bias=False, kernel_initializer=INIT_GLOROT_UNIFORM, activation=GELU),\n",
    "            tf.keras.layers.Dense(self.units, name=f'{self.name}_dense_2', use_bias=False, kernel_initializer=INIT_HE_UNIFORM),\n",
    "        ], name=f'{self.name}_dense')\n",
    "\n",
    "    def call(self, x):\n",
    "        return tf.where(\n",
    "                # Checks whether landmark is missing in frame\n",
    "                tf.reduce_sum(x, axis=2, keepdims=True) == 0,\n",
    "                # If so, the empty embedding is used\n",
    "                self.empty_embedding,\n",
    "                # Otherwise the landmark data is embedded\n",
    "                self.dense(x),\n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.337136Z",
     "iopub.status.busy": "2023-05-06T11:05:10.336476Z",
     "iopub.status.idle": "2023-05-06T11:05:10.365198Z",
     "shell.execute_reply": "2023-05-06T11:05:10.361628Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.337091Z"
    }
   },
   "outputs": [],
   "source": [
    "class CustomEmbedding(tf.keras.Model):\n",
    "    def __init__(self):\n",
    "        super(CustomEmbedding, self).__init__()\n",
    "        \n",
    "    def get_diffs(self, l):\n",
    "        S = l.shape[2]\n",
    "        other = tf.expand_dims(l, 3)\n",
    "        other = tf.repeat(other, S, axis=3)\n",
    "        other = tf.transpose(other, [0,1,3,2])\n",
    "        diffs = tf.expand_dims(l, 3) - other\n",
    "        diffs = tf.reshape(diffs, [-1, cfg.INPUT_SIZE, S*S])\n",
    "        return diffs\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        # Positional Embedding, initialized with zeros\n",
    "        self.positional_embedding = tf.keras.layers.Embedding(cfg.INPUT_SIZE+1, UNITS, embeddings_initializer=INIT_ZEROS)\n",
    "        # Embedding layer for Landmarks\n",
    "        self.lips_embedding = LandmarkEmbedding(LIPS_UNITS, 'lips')\n",
    "        self.left_hand_embedding = LandmarkEmbedding(HANDS_UNITS, 'left_hand')\n",
    "        self.right_hand_embedding = LandmarkEmbedding(HANDS_UNITS, 'right_hand')\n",
    "        self.pose_embedding = LandmarkEmbedding(POSE_UNITS, 'pose')\n",
    "        # Landmark Weights\n",
    "        self.landmark_weights = tf.Variable(tf.zeros([4], dtype=tf.float32), name='landmark_weights')\n",
    "        # Fully Connected Layers for combined landmarks\n",
    "        self.fc = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(UNITS, name='fully_connected_1', use_bias=False, kernel_initializer=INIT_GLOROT_UNIFORM, activation=GELU),\n",
    "            tf.keras.layers.Dense(UNITS, name='fully_connected_2', use_bias=False, kernel_initializer=INIT_HE_UNIFORM),\n",
    "        ], name='fc')\n",
    "\n",
    "\n",
    "    def call(self, lips0, left_hand0, right_hand0, pose0, non_empty_frame_idxs, training=False):\n",
    "        # Lips\n",
    "        lips_embedding = self.lips_embedding(lips0)\n",
    "        # Left Hand\n",
    "        left_hand_embedding = self.left_hand_embedding(left_hand0)\n",
    "        # Right Hand\n",
    "        right_hand_embedding = self.right_hand_embedding(right_hand0)\n",
    "        # Pose\n",
    "        pose_embedding = self.pose_embedding(pose0)\n",
    "        # Merge Embeddings of all landmarks with mean pooling\n",
    "        x = tf.stack((lips_embedding, left_hand_embedding, right_hand_embedding, pose_embedding), axis=3)\n",
    "        # Merge Landmarks with trainable attention weights\n",
    "        x = x * tf.nn.softmax(self.landmark_weights)\n",
    "        x = tf.reduce_sum(x, axis=3)\n",
    "        # Fully Connected Layers\n",
    "        x = self.fc(x)\n",
    "        # Add Positional Embedding\n",
    "        normalised_non_empty_frame_idxs = tf.where(\n",
    "            tf.math.equal(non_empty_frame_idxs, -1.0),\n",
    "            cfg.INPUT_SIZE,\n",
    "            tf.cast(\n",
    "                non_empty_frame_idxs / tf.reduce_max(non_empty_frame_idxs, axis=1, keepdims=True) * cfg.INPUT_SIZE,\n",
    "                tf.int32,\n",
    "            ),\n",
    "        )\n",
    "        x = x + self.positional_embedding(normalised_non_empty_frame_idxs)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.36731Z",
     "iopub.status.busy": "2023-05-06T11:05:10.366665Z",
     "iopub.status.idle": "2023-05-06T11:05:10.379701Z",
     "shell.execute_reply": "2023-05-06T11:05:10.37874Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.367271Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_lr_metric(optimizer):\n",
    "    def lr(y_true, y_pred):\n",
    "        return optimizer.lr\n",
    "    return lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.381531Z",
     "iopub.status.busy": "2023-05-06T11:05:10.381003Z",
     "iopub.status.idle": "2023-05-06T11:05:10.405099Z",
     "shell.execute_reply": "2023-05-06T11:05:10.404125Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.381493Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_model():\n",
    "    # Inputs\n",
    "    frames = tf.keras.layers.Input([cfg.INPUT_SIZE, N_COLS, cfg.N_DIMS], dtype=tf.float32, name='frames')\n",
    "    non_empty_frame_idxs = tf.keras.layers.Input([cfg.INPUT_SIZE], dtype=tf.float32, name='non_empty_frame_idxs')\n",
    "    # Padding Mask\n",
    "    mask = tf.cast(tf.math.not_equal(non_empty_frame_idxs, -1), tf.float32)\n",
    "    mask = tf.expand_dims(mask, axis=2)\n",
    "    \n",
    "    x = frames\n",
    "    x = tf.slice(x, [0,0,0,0], [-1,cfg.INPUT_SIZE, N_COLS, 2])\n",
    "    # LIPS\n",
    "    lips = tf.slice(x, [0,0,LIPS_START,0], [-1,cfg.INPUT_SIZE, 40, 2])\n",
    "    lips = tf.where(\n",
    "            tf.math.equal(lips, 0.0),\n",
    "            0.0,\n",
    "            (lips - LIPS_MEAN) / LIPS_STD,\n",
    "        )\n",
    "    lips = tf.reshape(lips, [-1, cfg.INPUT_SIZE, 40*2])\n",
    "    # LEFT HAND\n",
    "    left_hand = tf.slice(x, [0,0,40,0], [-1,cfg.INPUT_SIZE, 21, 2])\n",
    "    left_hand = tf.where(\n",
    "            tf.math.equal(left_hand, 0.0),\n",
    "            0.0,\n",
    "            (left_hand - LEFT_HANDS_MEAN) / LEFT_HANDS_STD,\n",
    "        )\n",
    "    left_hand = tf.reshape(left_hand, [-1, cfg.INPUT_SIZE, 21*2])\n",
    "    # RIGHT HAND\n",
    "    right_hand = tf.slice(x, [0,0,61,0], [-1,cfg.INPUT_SIZE, 21, 2])\n",
    "    right_hand = tf.where(\n",
    "            tf.math.equal(right_hand, 0.0),\n",
    "            0.0,\n",
    "            (right_hand - RIGHT_HANDS_MEAN) / RIGHT_HANDS_STD,\n",
    "        )\n",
    "    right_hand = tf.reshape(right_hand, [-1, cfg.INPUT_SIZE, 21*2])\n",
    "    # POSE\n",
    "    pose = tf.slice(x, [0,0,82,0], [-1,cfg.INPUT_SIZE, 10, 2])\n",
    "    pose = tf.where(\n",
    "            tf.math.equal(pose, 0.0),\n",
    "            0.0,\n",
    "            (pose - POSE_MEAN) / POSE_STD,\n",
    "        )\n",
    "    pose = tf.reshape(pose, [-1, cfg.INPUT_SIZE, 10*2])\n",
    "    x = lips, left_hand, right_hand, pose\n",
    "    x = CustomEmbedding()(lips, left_hand, right_hand, pose, non_empty_frame_idxs)\n",
    "    # Encoder Transformer Blocks\n",
    "    x = Transformer(NUM_BLOCKS)(x, mask)\n",
    "    # Pooling\n",
    "    x = tf.reduce_sum(x * mask, axis=1) / tf.reduce_sum(mask, axis=1)\n",
    "    # Classification Layer\n",
    "    x = tf.keras.layers.Dense(cfg.NUM_CLASSES, activation=tf.keras.activations.softmax, kernel_initializer=INIT_GLOROT_UNIFORM)(x)\n",
    "    outputs = x\n",
    "    \n",
    "    # Create Tensorflow Model\n",
    "    model = tf.keras.models.Model(inputs=[frames, non_empty_frame_idxs], outputs=outputs)\n",
    "    \n",
    "    # Simple Categorical Crossentropy Loss\n",
    "    loss = tf.keras.losses.SparseCategoricalCrossentropy()\n",
    "    \n",
    "    # Adam Optimizer with weight decay\n",
    "    optimizer = tfa.optimizers.AdamW(learning_rate=1e-3, weight_decay=1e-5, clipnorm=1.0)\n",
    "    \n",
    "    lr_metric = get_lr_metric(optimizer)\n",
    "    metrics = [\"acc\",lr_metric]\n",
    "    model.compile(loss=loss, optimizer=optimizer, metrics=metrics)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:05:10.409646Z",
     "iopub.status.busy": "2023-05-06T11:05:10.40899Z",
     "iopub.status.idle": "2023-05-06T11:05:13.43589Z",
     "shell.execute_reply": "2023-05-06T11:05:13.434886Z",
     "shell.execute_reply.started": "2023-05-06T11:05:10.40961Z"
    }
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "model_one = get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:07:19.169503Z",
     "iopub.status.busy": "2023-05-06T11:07:19.169089Z",
     "iopub.status.idle": "2023-05-06T11:07:19.260531Z",
     "shell.execute_reply": "2023-05-06T11:07:19.259739Z",
     "shell.execute_reply.started": "2023-05-06T11:07:19.169465Z"
    }
   },
   "outputs": [],
   "source": [
    "model_one.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:07:24.418773Z",
     "iopub.status.busy": "2023-05-06T11:07:24.41835Z",
     "iopub.status.idle": "2023-05-06T11:07:24.425763Z",
     "shell.execute_reply": "2023-05-06T11:07:24.42434Z",
     "shell.execute_reply.started": "2023-05-06T11:07:24.418738Z"
    }
   },
   "outputs": [],
   "source": [
    "def lrfn(current_step, num_warmup_steps, lr_max, num_cycles=0.50, num_training_steps=cfg.N_EPOCHS):\n",
    "    \n",
    "    if current_step < num_warmup_steps:\n",
    "        if WARMUP_METHOD == 'log':\n",
    "            return lr_max * 0.10 ** (num_warmup_steps - current_step)\n",
    "        else:\n",
    "            return lr_max * 2 ** -(num_warmup_steps - current_step)\n",
    "    else:\n",
    "        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n",
    "\n",
    "        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * lr_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:07:26.32001Z",
     "iopub.status.busy": "2023-05-06T11:07:26.319523Z",
     "iopub.status.idle": "2023-05-06T11:07:26.842071Z",
     "shell.execute_reply": "2023-05-06T11:07:26.841001Z",
     "shell.execute_reply.started": "2023-05-06T11:07:26.319972Z"
    }
   },
   "outputs": [],
   "source": [
    "def plot_lr_schedule(lr_schedule, epochs):\n",
    "    fig = plt.figure(figsize=(20, 10))\n",
    "    plt.plot([None] + lr_schedule + [None])\n",
    "    # X Labels\n",
    "    x = np.arange(1, epochs + 1)\n",
    "    x_axis_labels = [i if epochs <= 40 or i % 5 == 0 or i == 1 else None for i in range(1, epochs + 1)]\n",
    "    plt.xlim([1, epochs])\n",
    "    plt.xticks(x, x_axis_labels) # set tick step to 1 and let x axis start at 1\n",
    "    \n",
    "    # Increase y-limit for better readability\n",
    "    plt.ylim([0, max(lr_schedule) * 1.1])\n",
    "    \n",
    "    # Title\n",
    "    schedule_info = f'start: {lr_schedule[0]:.1E}, max: {max(lr_schedule):.1E}, final: {lr_schedule[-1]:.1E}'\n",
    "    plt.title(f'Step Learning Rate Schedule, {schedule_info}', size=18, pad=12)\n",
    "    \n",
    "    # Plot Learning Rates\n",
    "    for x, val in enumerate(lr_schedule):\n",
    "        if epochs <= 40 or x % 5 == 0 or x is epochs - 1:\n",
    "            if x < len(lr_schedule) - 1:\n",
    "                if lr_schedule[x - 1] < val:\n",
    "                    ha = 'right'\n",
    "                else:\n",
    "                    ha = 'left'\n",
    "            elif x == 0:\n",
    "                ha = 'right'\n",
    "            else:\n",
    "                ha = 'left'\n",
    "            plt.plot(x + 1, val, 'o', color='black');\n",
    "            offset_y = (max(lr_schedule) - min(lr_schedule)) * 0.02\n",
    "            plt.annotate(f'{val:.1E}', xy=(x + 1, val + offset_y), size=12, ha=ha)\n",
    "    \n",
    "    plt.xlabel('Epoch', size=16, labelpad=5)\n",
    "    plt.ylabel('Learning Rate', size=16, labelpad=5)\n",
    "    plt.grid()\n",
    "    plt.show()\n",
    "\n",
    "# Learning rate for encoder\n",
    "LR_SCHEDULE = [lrfn(step, num_warmup_steps=cfg.N_WARMUP_EPOCHS, lr_max=cfg.LR_MAX, num_cycles=0.50) for step in range(cfg.N_EPOCHS)]\n",
    "# Plot Learning Rate Schedule\n",
    "plot_lr_schedule(LR_SCHEDULE, epochs=cfg.N_EPOCHS)\n",
    "# Learning Rate Callback\n",
    "lr_callback = tf.keras.callbacks.LearningRateScheduler(lambda step: LR_SCHEDULE[step], verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Weight Decay Callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:07:31.21075Z",
     "iopub.status.busy": "2023-05-06T11:07:31.209765Z",
     "iopub.status.idle": "2023-05-06T11:07:31.218736Z",
     "shell.execute_reply": "2023-05-06T11:07:31.217641Z",
     "shell.execute_reply.started": "2023-05-06T11:07:31.210711Z"
    }
   },
   "outputs": [],
   "source": [
    "class WeightDecayCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, wd_ratio=cfg.WD_RATIO):\n",
    "        self.step_counter = 0\n",
    "        self.wd_ratio = wd_ratio\n",
    "    \n",
    "    def on_epoch_begin(self, epoch, logs=None):\n",
    "        model_one.optimizer.weight_decay = model_one.optimizer.learning_rate * self.wd_ratio\n",
    "        print(f'learning rate: {model_one.optimizer.learning_rate.numpy():.2e}, weight decay: {model_one.optimizer.weight_decay.numpy():.2e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance Benchmark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:07:33.454446Z",
     "iopub.status.busy": "2023-05-06T11:07:33.453415Z",
     "iopub.status.idle": "2023-05-06T11:07:43.284155Z",
     "shell.execute_reply": "2023-05-06T11:07:43.282972Z",
     "shell.execute_reply.started": "2023-05-06T11:07:33.454392Z"
    }
   },
   "outputs": [],
   "source": [
    "%%timeit -n 100\n",
    "if cfg.TRAIN_MODEL:\n",
    "    # Verify model_one prediction is <<<100ms\n",
    "    model_one.predict_on_batch({ 'frames': X[:1], 'non_empty_frame_idxs': NON_EMPTY_FRAME_IDXS[:1] })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:08:02.214504Z",
     "iopub.status.busy": "2023-05-06T11:08:02.213649Z",
     "iopub.status.idle": "2023-05-06T11:08:05.93921Z",
     "shell.execute_reply": "2023-05-06T11:08:05.937968Z",
     "shell.execute_reply.started": "2023-05-06T11:08:02.214454Z"
    }
   },
   "outputs": [],
   "source": [
    "X_train = X[train_idxs]\n",
    "X_val = X[val_idxs]\n",
    "NON_EMPTY_FRAME_IDXS_TRAIN = NON_EMPTY_FRAME_IDXS[train_idxs]\n",
    "NON_EMPTY_FRAME_IDXS_VAL = NON_EMPTY_FRAME_IDXS[val_idxs]\n",
    "y_train = y[train_idxs]\n",
    "y_val = y[val_idxs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:08:07.447505Z",
     "iopub.status.busy": "2023-05-06T11:08:07.447029Z",
     "iopub.status.idle": "2023-05-06T11:08:38.988736Z",
     "shell.execute_reply": "2023-05-06T11:08:38.987679Z",
     "shell.execute_reply.started": "2023-05-06T11:08:07.447459Z"
    }
   },
   "outputs": [],
   "source": [
    "if cfg.TRAIN_MODEL:\n",
    "    run = wandb.init(project=\"kaggle-asl-signs\", config=cfg, tags=['transformer', 'final-model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:10:08.082291Z",
     "iopub.status.busy": "2023-05-06T11:10:08.081567Z",
     "iopub.status.idle": "2023-05-06T11:31:49.717108Z",
     "shell.execute_reply": "2023-05-06T11:31:49.715938Z",
     "shell.execute_reply.started": "2023-05-06T11:10:08.082249Z"
    }
   },
   "outputs": [],
   "source": [
    "if cfg.TRAIN_MODEL:\n",
    "    tf.keras.backend.clear_session()\n",
    "    callbacks=[\n",
    "            lr_callback,\n",
    "            WeightDecayCallback(),\n",
    "            wandb.keras.WandbCallback()\n",
    "    ]\n",
    "    model_one.fit(\n",
    "        x=get_train_batch_all_signs(X, y, NON_EMPTY_FRAME_IDXS),\n",
    "        steps_per_epoch=len(X) // (cfg.NUM_CLASSES * cfg.BATCH_ALL_SIGNS_N),\n",
    "        epochs=cfg.N_EPOCHS,\n",
    "        batch_size=cfg.BATCH_SIZE,\n",
    "        callbacks=callbacks,\n",
    "        verbose = 2,) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:43:11.859511Z",
     "iopub.status.busy": "2023-05-06T11:43:11.858834Z",
     "iopub.status.idle": "2023-05-06T11:43:34.440436Z",
     "shell.execute_reply": "2023-05-06T11:43:34.439185Z",
     "shell.execute_reply.started": "2023-05-06T11:43:11.859472Z"
    }
   },
   "outputs": [],
   "source": [
    "if cfg.TRAIN_MODEL:\n",
    "    model_one.save('./final_model_one')\n",
    "    model_one.save_weights('./final_model_one_weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:44:22.236031Z",
     "iopub.status.busy": "2023-05-06T11:44:22.235356Z",
     "iopub.status.idle": "2023-05-06T11:44:22.566795Z",
     "shell.execute_reply": "2023-05-06T11:44:22.565653Z",
     "shell.execute_reply.started": "2023-05-06T11:44:22.235988Z"
    }
   },
   "outputs": [],
   "source": [
    "if cfg.TRAIN_MODEL:\n",
    "    artifact = wandb.Artifact('final_model_one', type='model')\n",
    "    artifact.add_file('./final_model_one_weights.data-00000-of-00001')\n",
    "    artifact.add_file('./final_model_one_weights.index')\n",
    "    run.log_artifact(artifact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:46:13.875532Z",
     "iopub.status.busy": "2023-05-06T11:46:13.87457Z",
     "iopub.status.idle": "2023-05-06T11:46:16.997776Z",
     "shell.execute_reply": "2023-05-06T11:46:16.996475Z",
     "shell.execute_reply.started": "2023-05-06T11:46:13.875481Z"
    }
   },
   "outputs": [],
   "source": [
    "# download model\n",
    "try:\n",
    "    artifact = run.use_artifact('amanarora/asl-sings/final_model_one:v1', type='model')\n",
    "except NameError:\n",
    "    run = wandb.init(project=\"kaggle-asl-signs\")\n",
    "    artifact = run.use_artifact('amanarora/asl-sings/final_model_one:v1', type='model')\n",
    "artifact_dir = artifact.download()\n",
    "!ls {artifact_dir}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:47:07.031473Z",
     "iopub.status.busy": "2023-05-06T11:47:07.031067Z",
     "iopub.status.idle": "2023-05-06T11:47:15.497579Z",
     "shell.execute_reply": "2023-05-06T11:47:15.495617Z",
     "shell.execute_reply.started": "2023-05-06T11:47:07.031433Z"
    }
   },
   "outputs": [],
   "source": [
    "# load model weights and do sanity c\n",
    "model_one.load_weights(f\"{artifact_dir}/final_model_one_weights\")\n",
    "y_val_pred = model_one.predict({ 'frames': X_val, 'non_empty_frame_idxs': NON_EMPTY_FRAME_IDXS_VAL }, verbose=2).argmax(axis=1)\n",
    "np.mean(y_val == y_val_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:47:44.742679Z",
     "iopub.status.busy": "2023-05-06T11:47:44.74165Z",
     "iopub.status.idle": "2023-05-06T11:47:44.852197Z",
     "shell.execute_reply": "2023-05-06T11:47:44.850918Z",
     "shell.execute_reply.started": "2023-05-06T11:47:44.742634Z"
    }
   },
   "outputs": [],
   "source": [
    "# delete variables as we go to free up RAM\n",
    "del X; del y; del X_train; del X_val; del y_train; del y_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:48:01.387105Z",
     "iopub.status.busy": "2023-05-06T11:48:01.386281Z",
     "iopub.status.idle": "2023-05-06T11:48:01.411592Z",
     "shell.execute_reply": "2023-05-06T11:48:01.407474Z",
     "shell.execute_reply.started": "2023-05-06T11:48:01.387063Z"
    }
   },
   "outputs": [],
   "source": [
    "# Landmark Weights\n",
    "weights = scipy.special.softmax(model_one.get_layer('custom_embedding').weights[15])\n",
    "landmarks = ['lips_embedding', 'left_hand_embedding', 'right_hand_embedding', 'pose_embedding']\n",
    "\n",
    "# Learned attention weights, initialized at uniform 25%\n",
    "for w, lm in zip(weights, landmarks):\n",
    "    print(f'{lm} weight: {(w*100):.1f}%')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:54:19.22156Z",
     "iopub.status.busy": "2023-05-06T11:54:19.22114Z",
     "iopub.status.idle": "2023-05-06T11:55:01.465057Z",
     "shell.execute_reply": "2023-05-06T11:55:01.463802Z",
     "shell.execute_reply.started": "2023-05-06T11:54:19.221522Z"
    }
   },
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_keras_model(model_one)\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "tflite_model = converter.convert()\n",
    "with open('./model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:55:24.132716Z",
     "iopub.status.busy": "2023-05-06T11:55:24.132336Z",
     "iopub.status.idle": "2023-05-06T11:55:25.367415Z",
     "shell.execute_reply": "2023-05-06T11:55:25.366194Z",
     "shell.execute_reply.started": "2023-05-06T11:55:24.132677Z"
    }
   },
   "outputs": [],
   "source": [
    "!du -sh ./model.tflite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-05-06T11:55:40.588483Z",
     "iopub.status.busy": "2023-05-06T11:55:40.587454Z",
     "iopub.status.idle": "2023-05-06T11:55:42.090855Z",
     "shell.execute_reply": "2023-05-06T11:55:42.089421Z",
     "shell.execute_reply.started": "2023-05-06T11:55:40.588431Z"
    }
   },
   "outputs": [],
   "source": [
    "!zip submission.zip ./model.tflite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "none",
   "dataSources": [
    {
     "databundleVersionId": 5087314,
     "sourceId": 46105,
     "sourceType": "competition"
    }
   ],
   "dockerImageVersionId": 30458,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}