tutorial_online_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b99721b9d6b26c10",
   "metadata": {},
   "source": [
    "# Online Inference Tutorial\n",
    "In this tutorial, we will implement online inference with event-based state-space models.\n",
    "Online inference is the process of classifying events as they arrive in real-time.\n",
    "For many edge systems, the batch size is 1, and the model has to meet a specific throughput of events per second.\n",
    "Here, you will test if your CPU is able to run real-time classification with EventSSM.\n",
    "\n",
    "The tutorial requires basic familiarity with JAX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bc0a9044321d654d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:22:32.658921Z",
     "start_time": "2024-05-27T09:22:32.654126Z"
    }
   },
   "outputs": [],
   "source": [
    "from hydra import initialize, compose\n",
    "from omegaconf import OmegaConf as om\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "from event_ssm.ssm import init_S5SSM\n",
    "from event_ssm.seq_model import ClassificationModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8b261a76014fbc7",
   "metadata": {},
   "source": [
    "## Step 1: Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7efb7b5428f7472",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:22:33.679045Z",
     "start_time": "2024-05-27T09:22:33.561733Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load configurations\n",
    "with initialize(version_base=None, config_path=\"configs\"):\n",
    "    cfg = compose(config_name=\"base.yaml\", overrides=[\"model=dvs/small\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "16eb6e254f8090cd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:22:33.771341Z",
     "start_time": "2024-05-27T09:22:33.766065Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ssm_init:\n",
      "  C_init: lecun_normal\n",
      "  dt_min: 0.001\n",
      "  dt_max: 0.1\n",
      "  conj_sym: false\n",
      "  clip_eigs: true\n",
      "ssm:\n",
      "  discretization: async\n",
      "  d_model: 128\n",
      "  d_ssm: 128\n",
      "  ssm_block_size: 16\n",
      "  num_stages: 2\n",
      "  num_layers_per_stage: 3\n",
      "  dropout: 0.25\n",
      "  classification_mode: timepool\n",
      "  prenorm: true\n",
      "  batchnorm: false\n",
      "  bn_momentum: 0.95\n",
      "  pooling_stride: 16\n",
      "  pooling_mode: timepool\n",
      "  state_expansion_factor: 2\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Print the configuration\n",
    "print(om.to_yaml(cfg.model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9806959c6627a4d5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:22:34.290002Z",
     "start_time": "2024-05-27T09:22:34.282856Z"
    }
   },
   "outputs": [],
   "source": [
    "# Set the random seed manually for reproducibility.\n",
    "key = jax.random.PRNGKey(cfg.seed)\n",
    "init_key, data_key = jax.random.split(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b936f3fdd1538bfe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:22:34.904836Z",
     "start_time": "2024-05-27T09:22:34.897859Z"
    }
   },
   "outputs": [],
   "source": [
    "# Model initialisation in flax\n",
    "ssm_init_fn = init_S5SSM(**cfg.model.ssm_init)\n",
    "\n",
    "# number of classes (dummy)\n",
    "classes = 10\n",
    "\n",
    "# number of tokens for a DVS sensor of size 128x128\n",
    "num_tokens = 128 * 128 * 2\n",
    "model = ClassificationModel(\n",
    "        ssm=ssm_init_fn,\n",
    "        num_classes=10,\n",
    "        num_embeddings=num_tokens,\n",
    "        **cfg.model.ssm,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accb046df2d07e7",
   "metadata": {},
   "source": [
    "EventSSM subsamples sequences in multiple stages to reduce the computational cost.\n",
    "Let's investigate the total subsampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3ed763820fe9f204",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:56:14.174709Z",
     "start_time": "2024-05-27T09:56:14.161702Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total subsampling: 256\n"
     ]
    }
   ],
   "source": [
    "total_subsampling = cfg.model.ssm.pooling_stride ** cfg.model.ssm.num_stages\n",
    "print(f\"Total subsampling: {total_subsampling}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e18fbb811f6c46e0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T09:56:42.653733Z",
     "start_time": "2024-05-27T09:56:38.056333Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SSM: 128 -> 128 -> 128 (stride 16 with pooling mode timepool)\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 256 -> 256 (stride 16 with pooling mode timepool)\n",
      "SSM: 256 -> 256 -> 256\n",
      "SSM: 256 -> 256 -> 256\n"
     ]
    }
   ],
   "source": [
    "# initialize model parameters\n",
    "x = jnp.zeros(total_subsampling, dtype=jnp.int32)\n",
    "t = jnp.ones(total_subsampling)\n",
    "variables = model.init(\n",
    "        {\"params\": init_key},\n",
    "        x, t, total_subsampling, False\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ed847f8098b7f53",
   "metadata": {},
   "source": [
    "## Step 2: Run the model on random data\n",
    "Generate a random list of integer tokens, jit compile the model and classify online."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9b32e55bfaf178e9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:14:54.375839Z",
     "start_time": "2024-05-27T10:14:54.360101Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequence length: 262144\n"
     ]
    }
   ],
   "source": [
    "# Generate random data\n",
    "sequence_length = 2 ** 18\n",
    "tokens = jax.random.randint(data_key, shape=(sequence_length,), minval=0, maxval=num_tokens)\n",
    "timesteps = jnp.ones(sequence_length)\n",
    "print(\"Sequence length:\", sequence_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8f49cd496d6ef30d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:15:07.170346Z",
     "start_time": "2024-05-27T10:14:55.732901Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SSM: 128 -> 128 -> 128 (stride 16 with pooling mode timepool)\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 256 -> 256 (stride 16 with pooling mode timepool)\n",
      "SSM: 256 -> 256 -> 256\n",
      "SSM: 256 -> 256 -> 256\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Array([-0.7493277 ,  0.78805447,  0.5494032 , -0.09461924, -1.7725958 ,\n",
       "       -0.00279285, -0.30725342,  1.2767574 , -2.3936124 , -1.1454588 ],      dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# jit compile the model\n",
    "from functools import partial\n",
    "jit_apply = jax.jit(partial(model.apply, length=total_subsampling, train=False))\n",
    "jit_apply(variables, x[:total_subsampling], t[:total_subsampling])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "55a885c77a44e8eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:15:11.763566Z",
     "start_time": "2024-05-27T10:15:08.166525Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looping through 262144 events with total_subsampling=256 --> 1024 iterations\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1024/1024 [00:02<00:00, 461.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken: 2.25s\n",
      "Events per second: 116318.93\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# loop through the model\n",
    "from tqdm import tqdm\n",
    "from time import time\n",
    "print(f\"Looping through {sequence_length} events with total_subsampling={total_subsampling} --> {sequence_length // total_subsampling} iterations\")\n",
    "start = time()\n",
    "for i in tqdm(range(0, sequence_length, total_subsampling)):\n",
    "    x = tokens[i:i + total_subsampling]\n",
    "    t = timesteps[i:i + total_subsampling]\n",
    "    logits = jit_apply(variables, x, t).block_until_ready()\n",
    "end = time()\n",
    "print(f\"Time taken: {end - start:.2f}s\")\n",
    "print(f\"Events per second: {sequence_length / (end - start):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "541f0afde67081f8",
   "metadata": {},
   "source": [
    "## Step 3: Optimize the inference speed\n",
    "We suggest to use [jax.lax.scan](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) instead of the for loop to further speed up the inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1318e7467cbb3b3f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:15:27.619686Z",
     "start_time": "2024-05-27T10:15:14.529552Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SSM: 128 -> 128 -> 128 (stride 16 with pooling mode timepool)\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 128 -> 128\n",
      "SSM: 128 -> 256 -> 256 (stride 16 with pooling mode timepool)\n",
      "SSM: 256 -> 256 -> 256\n",
      "SSM: 256 -> 256 -> 256\n"
     ]
    }
   ],
   "source": [
    "def step(carry, inputs):\n",
    "    x, t = inputs\n",
    "    logits = model.apply(variables, x, t, total_subsampling, False)\n",
    "    return None, logits\n",
    "tokens = tokens.reshape(-1, total_subsampling)\n",
    "timesteps = timesteps.reshape(-1, total_subsampling)\n",
    "\n",
    "# run the scan: first jit-compiles and then iterates\n",
    "logits = jax.lax.scan(step, init=None, xs=(tokens, timesteps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "aa170aadad84036d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:15:49.444621Z",
     "start_time": "2024-05-27T10:15:46.788818Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken: 1.56s\n",
      "Events per second: 168162.39\n"
     ]
    }
   ],
   "source": [
    "# measure run-time\n",
    "start = time()\n",
    "_, logits = jax.block_until_ready(jax.lax.scan(step, init=None, xs=(tokens, timesteps)))\n",
    "end = time()\n",
    "print(f\"Time taken: {end - start:.2f}s\")\n",
    "print(f\"Events per second: {sequence_length / (end - start):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "718dffb170c2df1c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-27T10:15:53.224299Z",
     "start_time": "2024-05-27T10:15:53.220810Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1024, 10)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcaba7dc4697605d",
   "metadata": {},
   "source": [
    "## Step 4: Run inference on the DVS128 Gestures dataset\n",
    "Follow the steps in the `tutorial_inference.ipynb` to run inference on the DVS128 Gestures dataset with a pretrained model.\n",
    "Plot the confidence of the model in the correct class over time"
   ]
  }
 ],
 "metadata": {
  "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}