Skip to content

Commit

Permalink
Merge branch 'nonseq_dvs' of github.com:synsense/sinabs into nonseq_dvs
Browse files Browse the repository at this point in the history
  • Loading branch information
bauerfe committed Oct 31, 2024
2 parents 02ca1b2 + 577509a commit c23acf8
Show file tree
Hide file tree
Showing 3 changed files with 387 additions and 10 deletions.
333 changes: 333 additions & 0 deletions examples/dynapcnn_network/snn_DVSLayer_given_followed_by_pool.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,333 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"metadata": {}
},
"outputs": [
{
"data": {
"text/plain": [
"<torch._C.Generator at 0x73c31f4e9090>"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"from sinabs.backend.dynapcnn import DynapcnnNetwork\n",
"from sinabs.backend.dynapcnn import DVSLayer\n",
"from sinabs.layers import Merge, IAFSqueeze, SumPool2d\n",
"from sinabs.activation.surrogate_gradient_fn import PeriodicExponential\n",
"import sinabs.layers as sl\n",
"\n",
"from torch.nn import CrossEntropyLoss\n",
"from torch.optim import Adam\n",
"\n",
"from tonic.datasets.nmnist import NMNIST\n",
"from tonic.transforms import ToFrame\n",
"from torch.utils.data import DataLoader\n",
"import numpy as np\n",
"from tqdm.notebook import tqdm\n",
"from statistics import mode\n",
"\n",
"device = torch.device('cpu')\n",
"torch.manual_seed(0)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"channels = 1\n",
"height = 34\n",
"width = 34\n",
"batch_size = 1\n",
"\n",
"input_shape = (channels, height, width)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {
"metadata": {}
},
"outputs": [
{
"data": {
"text/plain": [
"SNN(\n",
" (dvs): DVSLayer(\n",
" (pool_layer): SumPool2d(norm_type=1, kernel_size=(1, 1), stride=None, ceil_mode=False)\n",
" (crop_layer): Crop2d((0, 34), (0, 34))\n",
" (flip_layer): FlipDims()\n",
" )\n",
" (dvs_pool): AvgPool2d(kernel_size=1, stride=1, padding=0)\n",
" (conv1): Conv2d(1, 10, kernel_size=(2, 2), stride=(1, 1), bias=False)\n",
" (iaf1): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(1.), min_v_mem=Parameter containing:\n",
" tensor(-1.), batch_size=1, num_timesteps=-1)\n",
" (pool1): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
" (conv2): Conv2d(10, 10, kernel_size=(4, 4), stride=(1, 1), bias=False)\n",
" (iaf2): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(1.), min_v_mem=Parameter containing:\n",
" tensor(-1.), batch_size=1, num_timesteps=-1)\n",
" (conv3): Conv2d(10, 1, kernel_size=(2, 2), stride=(1, 1), bias=False)\n",
" (iaf3): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(1.), min_v_mem=Parameter containing:\n",
" tensor(-1.), batch_size=1, num_timesteps=-1)\n",
" (fc1): Linear(in_features=144, out_features=200, bias=False)\n",
" (iaf4): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(1.), min_v_mem=Parameter containing:\n",
" tensor(-1.), batch_size=1, num_timesteps=-1)\n",
" (fc2): Linear(in_features=200, out_features=10, bias=False)\n",
" (iaf5): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(1.), min_v_mem=Parameter containing:\n",
" tensor(-1.), batch_size=1, num_timesteps=-1)\n",
" (flat): Flatten(start_dim=1, end_dim=-1)\n",
")"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"class SNN(nn.Module):\n",
" def __init__(self, input_shape) -> None:\n",
" super().__init__()\n",
"\n",
" self.dvs = DVSLayer(input_shape=(input_shape[1], input_shape[2]))\n",
" self.dvs_pool = nn.AvgPool2d(1,1)\n",
" # -- chip core A --\n",
" self.conv1 = nn.Conv2d(1, 10, 2, 1, bias=False)\n",
" self.iaf1 = IAFSqueeze(batch_size=batch_size, min_v_mem=-1.0, spike_threshold=1.0, surrogate_grad_fn=PeriodicExponential())\n",
" self.pool1 = nn.AvgPool2d(2,2)\n",
" # -- chip core B --\n",
" self.conv2 = nn.Conv2d(10, 10, 4, 1, bias=False)\n",
" self.iaf2 = IAFSqueeze(batch_size=batch_size, min_v_mem=-1.0, spike_threshold=1.0, surrogate_grad_fn=PeriodicExponential())\n",
" # -- chip core C --\n",
" self.conv3 = nn.Conv2d(10, 1, 2, 1, bias=False)\n",
" self.iaf3 = IAFSqueeze(batch_size=batch_size, min_v_mem=-1.0, spike_threshold=1.0, surrogate_grad_fn=PeriodicExponential())\n",
" # -- chip core D --\n",
" self.fc1 = nn.Linear(144, 200, bias=False)\n",
" self.iaf4 = IAFSqueeze(batch_size=batch_size, min_v_mem=-1.0, spike_threshold=1.0, surrogate_grad_fn=PeriodicExponential())\n",
" # -- chip core E --\n",
" self.fc2 = nn.Linear(200, 10, bias=False)\n",
" self.iaf5 = IAFSqueeze(batch_size=batch_size, min_v_mem=-1.0, spike_threshold=1.0, surrogate_grad_fn=PeriodicExponential())\n",
"\n",
" # -- layers ignored during deployment --\n",
" self.flat = nn.Flatten()\n",
"\n",
" def init_weights(self):\n",
" for name, layer in self.named_modules():\n",
" if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):\n",
" nn.init.xavier_normal_(layer.weight.data)\n",
"\n",
" def detach_neuron_states(self):\n",
" for name, layer in self.named_modules():\n",
" if name != '':\n",
" if isinstance(layer, sl.StatefulLayer):\n",
" for name, buffer in layer.named_buffers():\n",
" buffer.detach_()\n",
"\n",
" def forward(self, x):\n",
" \n",
" dvs_out = self.dvs(x) # 0\n",
"\n",
" dvs_pool_out = self.dvs_pool(dvs_out)\n",
" \n",
" con1_out = self.conv1(dvs_pool_out) # 4\n",
" iaf1_out = self.iaf1(con1_out) # 5\n",
" pool1_out = self.pool1(iaf1_out) # 6\n",
"\n",
" conv2_out = self.conv2(pool1_out) # 7\n",
" iaf2_out = self.iaf2(conv2_out) # 8\n",
"\n",
" conv3_out = self.conv3(iaf2_out) # 9\n",
" iaf3_out = self.iaf3(conv3_out) # 10\n",
"\n",
" flat_out = self.flat(iaf3_out) # 15\n",
" \n",
" fc1_out = self.fc1(flat_out) # 11\n",
" iaf4_out = self.iaf4(fc1_out) # 12\n",
" fc2_out = self.fc2(iaf4_out) # 13\n",
" iaf5_out = self.iaf5(fc2_out) # 14\n",
"\n",
" return iaf5_out\n",
" \n",
"snn = SNN(input_shape)\n",
"snn.init_weights()\n",
"snn.to(device)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"metadata": {}
},
"outputs": [],
"source": [
"hw_model = DynapcnnNetwork(\n",
" snn=snn,\n",
" input_shape=input_shape,\n",
" batch_size=batch_size,\n",
" discretize=True\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"input_dummy = torch.randn((batch_size, *input_shape))"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"out = hw_model(input_dummy)"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"tensor([[[[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]],\n",
"\n",
" [[0.]]]], grad_fn=<PowBackward0>)\n"
]
}
],
"source": [
"print(out)"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Network is valid\n"
]
},
{
"data": {
"text/plain": [
"DynapcnnNetwork(\n",
" (_dynapcnn_module): DynapcnnNetworkModule(\n",
" (_dynapcnn_layers): ModuleDict(\n",
" (1): DynapcnnLayer(\n",
" (_conv): Conv2d(1, 10, kernel_size=(2, 2), stride=(1, 1), bias=False)\n",
" (_spk): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(174.), min_v_mem=Parameter containing:\n",
" tensor(-174.), batch_size=1, num_timesteps=-1)\n",
" )\n",
" (5): DynapcnnLayer(\n",
" (_conv): Conv2d(200, 10, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
" (_spk): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(419.), min_v_mem=Parameter containing:\n",
" tensor(-419.), batch_size=1, num_timesteps=-1)\n",
" )\n",
" (3): DynapcnnLayer(\n",
" (_conv): Conv2d(10, 1, kernel_size=(2, 2), stride=(1, 1), bias=False)\n",
" (_spk): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(303.), min_v_mem=Parameter containing:\n",
" tensor(-303.), batch_size=1, num_timesteps=-1)\n",
" )\n",
" (2): DynapcnnLayer(\n",
" (_conv): Conv2d(10, 10, kernel_size=(4, 4), stride=(1, 1), bias=False)\n",
" (_spk): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(130.), min_v_mem=Parameter containing:\n",
" tensor(-130.), batch_size=1, num_timesteps=-1)\n",
" )\n",
" (4): DynapcnnLayer(\n",
" (_conv): Conv2d(1, 200, kernel_size=(12, 12), stride=(1, 1), bias=False)\n",
" (_spk): IAFSqueeze(spike_threshold=Parameter containing:\n",
" tensor(365.), min_v_mem=Parameter containing:\n",
" tensor(-365.), batch_size=1, num_timesteps=-1)\n",
" )\n",
" (0): DVSLayer(\n",
" (pool_layer): SumPool2d(norm_type=1, kernel_size=(1, 1), stride=None, ceil_mode=False)\n",
" (crop_layer): Crop2d((0, 34), (0, 34))\n",
" (flip_layer): FlipDims()\n",
" )\n",
" )\n",
" (merge_layer): Merge()\n",
" )\n",
")"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"hw_model.to(device=\"speck2fdevkit\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "speck-rescnn",
"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.3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
2 changes: 0 additions & 2 deletions sinabs/backend/dynapcnn/connectivity_specs.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,6 @@
(Pooling, Weight): "pooling-weight",
# Dvs can be followed by weight layer of next core
(DVS, Weight): "dvs-weight",
# Dvs can be followed by pooling layers
(DVS, Pooling): "dvs-pooling",
}

# Unpack dict
Expand Down
Loading

0 comments on commit c23acf8

Please sign in to comment.