data-generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Installation\n",
    "Start with creating a new Conda environment\n",
    "```\n",
    "conda create -n frameworks pip\n",
    "conda activate frameworks\n",
    "```\n",
    "Then install PyTorch (adjust for your CUDA version). Instructions available [here](https://pytorch.org/get-started/locally/)\n",
    "```\n",
    "conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia\n",
    "```\n",
    "Install the benchmarked frameworks from PyPI\n",
    "```\n",
    "pip install -r requirements.txt\n",
    "```\n",
    "\n",
    "This benchmark code is an adaptation of Rockpool's [benchmark script](https://gitlab.com/synsense/rockpool/-/blob/develop/rockpool/utilities/benchmarking/benchmark_utils.py?ref_type=heads). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from utils import timeit, benchmark_framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rockpool_torch():\n",
    "    from rockpool.nn.modules import LIFTorch, LinearTorch\n",
    "    from rockpool.nn.combinators import Sequential\n",
    "    import rockpool\n",
    "\n",
    "    benchmark_title = f\"Rockpool<br>v{rockpool.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        model = Sequential(\n",
    "            LinearTorch(shape=(n_neurons, n_neurons)),\n",
    "            LIFTorch(n_neurons),\n",
    "        ).to(device)\n",
    "        input_static = torch.randn(batch_size, n_steps, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        output = model(input_static)[0]\n",
    "        bench_dict[\"output\"] = output\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rockpool_exodus():\n",
    "    from rockpool.nn.modules import LIFExodus, LinearTorch\n",
    "    from rockpool.nn.combinators import Sequential\n",
    "    import rockpool\n",
    "\n",
    "    benchmark_title = f\"Rockpool EXODUS<br>v{rockpool.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        model = Sequential(\n",
    "            LinearTorch(shape=(n_neurons, n_neurons)),\n",
    "            LIFExodus(n_neurons),\n",
    "        ).to(device)\n",
    "        input_static = torch.randn(batch_size, n_steps, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        output = model(input_static)[0]\n",
    "        bench_dict[\"output\"] = output\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sinabs():\n",
    "    from sinabs.layers import LIF\n",
    "    import sinabs\n",
    "\n",
    "    benchmark_title = f\"Sinabs<br>v{sinabs.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        model = nn.Sequential(\n",
    "            nn.Linear(n_neurons, n_neurons),\n",
    "            LIF(tau_mem=torch.tensor(10.0)),\n",
    "        ).to(device)\n",
    "        input_static = torch.randn(batch_size, n_steps, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        sinabs.reset_states(model)\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sinabs_exodus():\n",
    "    from sinabs.exodus.layers import LIF\n",
    "    import sinabs\n",
    "\n",
    "    benchmark_title = f\"Sinabs EXODUS<br>v{sinabs.exodus.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        model = nn.Sequential(\n",
    "            nn.Linear(n_neurons, n_neurons),\n",
    "            LIF(tau_mem=torch.tensor(10.0)),\n",
    "        ).to(device)\n",
    "        input_static = torch.randn(batch_size, n_steps, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        sinabs.reset_states(model)\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def norse():\n",
    "    from norse.torch.module.lif import LIF\n",
    "    from norse.torch import SequentialState\n",
    "    import norse\n",
    "\n",
    "    benchmark_title = f\"Norse<br>v{norse.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        model = SequentialState(\n",
    "            nn.Linear(n_neurons, n_neurons),\n",
    "            LIF(),\n",
    "        )\n",
    "        # model = torch.compile(model, mode=\"max-autotune\")\n",
    "        model = model.to(device)\n",
    "        input_static = torch.randn(n_steps, batch_size, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        # output.sum().backward() # JIT compile everything\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        bench_dict[\"output\"] = model(input_static)[0]\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def snntorch():\n",
    "    import snntorch\n",
    "\n",
    "    benchmark_title = f\"snnTorch<br>v{snntorch.__version__}\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        class Model(nn.Module):\n",
    "            def __init__(self, beta: float = 0.95):\n",
    "                super().__init__()\n",
    "                self.fc = nn.Linear(n_neurons, n_neurons)\n",
    "                self.lif = snntorch.Leaky(beta=beta)\n",
    "                self.mem = self.lif.init_leaky()\n",
    "\n",
    "            def forward(self, x):\n",
    "                output = []\n",
    "                mem = self.mem\n",
    "                for inp in x:\n",
    "                    cur = self.fc(inp)\n",
    "                    spk, mem = self.lif(cur, mem)\n",
    "                    output.append(spk)\n",
    "                return torch.stack(output)\n",
    "\n",
    "        model = Model()\n",
    "        # model = torch.compile(model, mode=\"max-autotune\")\n",
    "        model = model.to(device)\n",
    "        input_static = torch.randn(n_steps, batch_size, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mix of https://spikingjelly.readthedocs.io/zh_CN/latest/activation_based_en/basic_concept.html#step-mode\n",
    "# and https://github.com/fangwei123456/spikingjelly/blob/master/spikingjelly/activation_based/examples/rsnn_sequential_fmnist.py\n",
    "def spikingjelly():\n",
    "    from spikingjelly.activation_based import neuron, surrogate, functional, layer\n",
    "\n",
    "    benchmark_title = f\"SpikingJelly PyTorch<br>v0.0.0.0.15\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        class Model(nn.Module):\n",
    "            def __init__(self, tau=5.0):\n",
    "                super().__init__()\n",
    "                self.model = nn.Sequential(\n",
    "                    layer.Linear(n_neurons, n_neurons),\n",
    "                    neuron.LIFNode(tau=tau, surrogate_function=surrogate.ATan(), step_mode='m'),\n",
    "                )\n",
    "\n",
    "            def forward(self, x):\n",
    "                functional.reset_net(self.model)\n",
    "                return self.model(x)\n",
    "\n",
    "        model = Model().to(device)\n",
    "        input_static = torch.randn(n_steps, batch_size, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title\n",
    "\n",
    "\n",
    "def spikingjelly_cupy():\n",
    "    from spikingjelly.activation_based import neuron, surrogate, functional, layer\n",
    "\n",
    "    benchmark_title = f\"SpikingJelly CuPy<br>v0.0.0.0.15\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        class Model(nn.Module):\n",
    "            def __init__(self, tau=5.0):\n",
    "                super().__init__()\n",
    "                self.model = nn.Sequential(\n",
    "                    layer.Linear(n_neurons, n_neurons),\n",
    "                    neuron.LIFNode(tau=tau, surrogate_function=surrogate.ATan(), step_mode='m'),\n",
    "                )\n",
    "                functional.set_backend(self.model, backend='cupy')\n",
    "\n",
    "            def forward(self, x):\n",
    "                functional.reset_net(self.model)\n",
    "                return self.model(x)\n",
    "\n",
    "        model = Model().to(device)\n",
    "        input_static = torch.randn(n_steps, batch_size, n_neurons).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lava():\n",
    "    import lava.lib.dl.slayer as slayer\n",
    "\n",
    "    benchmark_title = f\"Lava DL<br>v0.4.0.dev0\"\n",
    "\n",
    "    def prepare_fn(batch_size, n_steps, n_neurons, n_layers, device):\n",
    "        neuron_params = {\n",
    "                        'threshold'     : 0.1,\n",
    "                        'current_decay' : 1,\n",
    "                        'voltage_decay' : 0.1,\n",
    "                        'requires_grad' : True,     \n",
    "                    }\n",
    "        # slayer.block automatically add quantization.\n",
    "        # They can be disabled by setting pre_hook_fx=None\n",
    "        model = slayer.block.cuba.Dense(neuron_params, n_neurons, n_neurons, pre_hook_fx=None).to(device)\n",
    "        input_static = torch.randn(batch_size, n_neurons, n_steps).to(device)\n",
    "        with torch.no_grad():\n",
    "            model(input_static)\n",
    "        return dict(model=model, input=input_static, n_neurons=n_neurons)\n",
    "\n",
    "    def forward_fn(bench_dict):\n",
    "        model, input_static = bench_dict[\"model\"], bench_dict[\"input\"]\n",
    "        bench_dict[\"output\"] = model(input_static)\n",
    "        return bench_dict\n",
    "\n",
    "    def backward_fn(bench_dict):\n",
    "        output = bench_dict[\"output\"]\n",
    "        loss = output.sum()\n",
    "        loss.backward(retain_graph=True)\n",
    "\n",
    "    return prepare_fn, forward_fn, backward_fn, benchmark_title"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 10\n",
    "n_steps = 500\n",
    "n_layers = 1  # doesn't do anything at the moment\n",
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Benchmarking SpikingJelly CuPy<br>v0.0.0.0.15 with n_neurons = 512\n",
      "Benchmarking SpikingJelly CuPy<br>v0.0.0.0.15 with n_neurons = 4096\n",
      "Benchmarking SpikingJelly CuPy<br>v0.0.0.0.15 with n_neurons = 16384\n",
      "Benchmarking Rockpool<br>v2.6 with n_neurons = 512\n",
      "Benchmarking Rockpool<br>v2.6 with n_neurons = 4096\n",
      "Benchmarking Rockpool<br>v2.6 with n_neurons = 16384\n",
      "Benchmarking Rockpool EXODUS<br>v2.6 with n_neurons = 512\n",
      "Benchmarking Rockpool EXODUS<br>v2.6 with n_neurons = 4096\n",
      "Benchmarking Rockpool EXODUS<br>v2.6 with n_neurons = 16384\n",
      "Benchmarking Sinabs<br>v1.2.8 with n_neurons = 512\n",
      "Benchmarking Sinabs<br>v1.2.8 with n_neurons = 4096\n",
      "Benchmarking Sinabs<br>v1.2.8 with n_neurons = 16384\n",
      "Benchmarking Sinabs EXODUS<br>v1.1.2 with n_neurons = 512\n",
      "Benchmarking Sinabs EXODUS<br>v1.1.2 with n_neurons = 4096\n",
      "Benchmarking Sinabs EXODUS<br>v1.1.2 with n_neurons = 16384\n",
      "Benchmarking snnTorch<br>v0.7.0 with n_neurons = 512\n",
      "Benchmarking snnTorch<br>v0.7.0 with n_neurons = 4096\n",
      "Benchmarking snnTorch<br>v0.7.0 with n_neurons = 16384\n",
      "Benchmarking Norse<br>v1.0.0 with n_neurons = 512\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/issues/benchmark_env/lib/python3.8/site-packages/norse/torch/functional/lif.py:391: UserWarning: operator() profile_node %71 : int = prim::profile_ivalue(%69)\n",
      " does not have profile information (Triggered internally at ../third_party/nvfuser/csrc/graph_fuser.cpp:104.)\n",
      "  z, state = _lif_feed_forward_step_jit(input_spikes, state, LIFParametersJIT(*p), dt)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Benchmarking Norse<br>v1.0.0 with n_neurons = 4096\n",
      "Benchmarking Norse<br>v1.0.0 with n_neurons = 16384\n",
      "Benchmarking Lava DL<br>v0.4.0.dev0 with n_neurons = 512\n",
      "Benchmarking Lava DL<br>v0.4.0.dev0 with n_neurons = 4096\n",
      "Benchmarking Lava DL<br>v0.4.0.dev0 with n_neurons = 16384\n"
     ]
    }
   ],
   "source": [
    "data = []\n",
    "for benchmark in [spikingjelly_cupy, rockpool_torch, rockpool_exodus, sinabs, sinabs_exodus, snntorch, norse, lava,]: # spikingjelly\n",
    "# for benchmark in [ norse, snntorch,]:\n",
    "    for n_neurons in [512, 4096, 16384, ]: #  1024, 2048, 4096, 8192, 16384,\n",
    "        prepare_fn, forward_fn, backward_fn, bench_desc = benchmark()\n",
    "        print(\"Benchmarking\", bench_desc, \"with n_neurons =\", n_neurons)\n",
    "        forward_times, backward_times = benchmark_framework(\n",
    "            prepare_fn=prepare_fn,\n",
    "            forward_fn=forward_fn,\n",
    "            backward_fn=backward_fn,\n",
    "            benchmark_desc=bench_desc,\n",
    "            n_neurons=n_neurons,\n",
    "            n_layers=n_layers,\n",
    "            n_steps=n_steps,\n",
    "            batch_size=batch_size,\n",
    "            device=device,\n",
    "        )\n",
    "        data.append(\n",
    "            [\n",
    "                bench_desc,\n",
    "                np.array(forward_times).mean(),\n",
    "                np.array(backward_times).mean(),\n",
    "                n_neurons,\n",
    "            ]\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df = pd.DataFrame(data, columns=[\"framework\", \"forward\", \"backward\", \"neurons\"])\n",
    "df = df.melt(\n",
    "    id_vars=[\"framework\", \"neurons\"],\n",
    "    value_vars=[\"forward\", \"backward\"],\n",
    "    var_name=\"pass\",\n",
    "    value_name=\"time [s]\",\n",
    ")\n",
    "df.to_csv(\"data.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "frameworks",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}