├── ESSCIRC_OSN.ipynb ├── ESSCIRC_OSN_Sky130.ipynb ├── ESSCIRC_OSN_cheatsheet.ipynb ├── LICENSE ├── README.md └── graphical-abstract.jpg /ESSCIRC_OSN.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "provenance": [], 7 | "gpuType": "T4" 8 | }, 9 | "kernelspec": { 10 | "name": "python3", 11 | "display_name": "Python 3" 12 | }, 13 | "language_info": { 14 | "name": "python" 15 | }, 16 | "accelerator": "GPU" 17 | }, 18 | "cells": [ 19 | { 20 | "cell_type": "markdown", 21 | "source": [ 22 | "#Open-Source Neuromorphic Circuit Design\n", 23 | "## By Jason K. Eshraghian (www.ncg.ucsc.edu)\n", 24 | "\n", 25 | "\n", 26 | "[](https://github.com/jeshraghian/snntorch/)" 27 | ], 28 | "metadata": { 29 | "id": "8w6lhn7H8fW5" 30 | } 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": { 36 | "id": "1BlTqunB73-t" 37 | }, 38 | "outputs": [], 39 | "source": [ 40 | "!pip install snntorch --quiet # shift + enter" 41 | ] 42 | }, 43 | { 44 | "cell_type": "markdown", 45 | "source": [ 46 | "*What will I learn?*\n", 47 | "\n", 48 | "1. Train an SNN classifier using snnTorch\n", 49 | "2. Hardware Friendly Training\n", 50 | " - Weight Quantization with Brevitas\n", 51 | " - Stateful Quantization\n", 52 | "3. Manual Non-Idealities, e.g., read noise" 53 | ], 54 | "metadata": { 55 | "id": "xLlnINAI9mgJ" 56 | } 57 | }, 58 | { 59 | "cell_type": "markdown", 60 | "source": [ 61 | "# 1. Train an SNN Classifier using snnTorch\n", 62 | "## 1.1 Imports\n" 63 | ], 64 | "metadata": { 65 | "id": "-CS50cwuCW6n" 66 | } 67 | }, 68 | { 69 | "cell_type": "code", 70 | "source": [ 71 | "# snntorch imports\n", 72 | "import snntorch as snn\n", 73 | "from snntorch import functional as SF\n", 74 | "\n", 75 | "# pytorch imports\n", 76 | "import torch\n", 77 | "import torch.nn as nn\n", 78 | "from torch.utils.data import DataLoader\n", 79 | "from torchvision import datasets, transforms\n", 80 | "\n", 81 | "# data manipulation\n", 82 | "import numpy as np\n", 83 | "import itertools\n", 84 | "\n", 85 | "# plotting\n", 86 | "import matplotlib.pyplot as plt\n", 87 | "from IPython.display import HTML" 88 | ], 89 | "metadata": { 90 | "id": "H_TzogsCCcSe" 91 | }, 92 | "execution_count": null, 93 | "outputs": [] 94 | }, 95 | { 96 | "cell_type": "markdown", 97 | "source": [ 98 | "## 1.2 Boilerplate: DataLoading the MNIST Dataset" 99 | ], 100 | "metadata": { 101 | "id": "nftOdpyAGv7D" 102 | } 103 | }, 104 | { 105 | "cell_type": "code", 106 | "source": [ 107 | "# dataloader arguments\n", 108 | "batch_size = 128\n", 109 | "data_path='/data/mnist'\n", 110 | "\n", 111 | "dtype = torch.float\n", 112 | "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n", 113 | "## if you're on M1 or M2 GPU:\n", 114 | "# device = torch.device(\"mps\")" 115 | ], 116 | "metadata": { 117 | "id": "SsM2Z5NXGu5z" 118 | }, 119 | "execution_count": null, 120 | "outputs": [] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "source": [ 125 | "# Define a transform\n", 126 | "transform = transforms.Compose([\n", 127 | " transforms.Resize((28, 28)),\n", 128 | " transforms.Grayscale(),\n", 129 | " transforms.ToTensor(),\n", 130 | " transforms.Normalize((0,), (1,))])\n", 131 | "\n", 132 | "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n", 133 | "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)" 134 | ], 135 | "metadata": { 136 | "id": "XqbYptgDHUPg" 137 | }, 138 | "execution_count": null, 139 | "outputs": [] 140 | }, 141 | { 142 | "cell_type": "code", 143 | "source": [ 144 | "# Create DataLoaders\n", 145 | "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n", 146 | "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)" 147 | ], 148 | "metadata": { 149 | "id": "jSlS3gWZHXI0" 150 | }, 151 | "execution_count": null, 152 | "outputs": [] 153 | }, 154 | { 155 | "cell_type": "markdown", 156 | "source": [ 157 | "## 1.3 Construct SNN Model" 158 | ], 159 | "metadata": { 160 | "id": "SJVhfNukHbsp" 161 | } 162 | }, 163 | { 164 | "cell_type": "code", 165 | "source": [ 166 | "# Network Architecture\n", 167 | "num_inputs =\n", 168 | "num_hidden =\n", 169 | "num_outputs =\n", 170 | "\n", 171 | "# Temporal Dynamics\n", 172 | "num_steps =\n", 173 | "beta =" 174 | ], 175 | "metadata": { 176 | "id": "uu324fr_HhxV" 177 | }, 178 | "execution_count": null, 179 | "outputs": [] 180 | }, 181 | { 182 | "cell_type": "code", 183 | "source": [ 184 | "from snntorch import surrogate\n", 185 | "\n", 186 | "# Define Network\n", 187 | "class Net(nn.Module):\n", 188 | " def __init__(self):\n", 189 | " super().__init__()\n", 190 | "\n", 191 | " # Initialize layers\n", 192 | " self.fc1 =\n", 193 | " self.lif1 =\n", 194 | " self.fc2 =\n", 195 | " self.lif2 =\n", 196 | "\n", 197 | " def forward(self, x):\n", 198 | "\n", 199 | " # Initialize hidden states at t=0\n", 200 | " mem1 = self.lif1.init_leaky()\n", 201 | " mem2 = self.lif2.init_leaky()\n", 202 | "\n", 203 | " # Record the final layer\n", 204 | " spk2_rec = []\n", 205 | " mem2_rec = []\n", 206 | "\n", 207 | " # time-loop\n", 208 | " for step in range(num_steps):\n", 209 | " cur1 = self.fc1(...) # batch: 128 x 784\n", 210 | " spk1, mem1 = self.lif1(...)\n", 211 | " cur2 = self.fc2(...)\n", 212 | " spk2, mem2 = self.lif2(...)\n", 213 | "\n", 214 | " # store in list\n", 215 | " spk2_rec.append(spk2)\n", 216 | " mem2_rec.append(mem2)\n", 217 | "\n", 218 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0) # time-steps x batch x num_out\n", 219 | "\n", 220 | "# Load the network onto CUDA if available\n", 221 | "net = Net().to(device)" 222 | ], 223 | "metadata": { 224 | "id": "CkM1Z1EjHeW8" 225 | }, 226 | "execution_count": null, 227 | "outputs": [] 228 | }, 229 | { 230 | "cell_type": "markdown", 231 | "source": [ 232 | "## 1.4 Training the SNN" 233 | ], 234 | "metadata": { 235 | "id": "p8qBw03rHpn3" 236 | } 237 | }, 238 | { 239 | "cell_type": "code", 240 | "source": [ 241 | "def training_loop(model, dataloader, num_epochs=1):\n", 242 | " loss = nn.CrossEntropyLoss()\n", 243 | " optimizer = torch.optim.Adam(model.parameters(), lr=5e-4, betas=(0.9, 0.999))\n", 244 | " counter = 0\n", 245 | "\n", 246 | " # Outer training loop\n", 247 | " for epoch in range(num_epochs):\n", 248 | " train_batch = iter(dataloader)\n", 249 | "\n", 250 | " # Minibatch training loop\n", 251 | " for data, targets in train_batch:\n", 252 | " data = data.to(device)\n", 253 | " targets = targets.to(device)\n", 254 | "\n", 255 | " # forward pass\n", 256 | " model.train()\n", 257 | " spk_rec, _ = model(data)\n", 258 | "\n", 259 | " # initialize the loss & sum over time\n", 260 | " loss_val = torch.zeros((1), dtype=dtype, device=device)\n", 261 | " loss_val = loss(spk_rec.sum(0), targets) # batch x num_out\n", 262 | "\n", 263 | " # Gradient calculation + weight update\n", 264 | " optimizer.zero_grad()\n", 265 | " loss_val.backward()\n", 266 | " optimizer.step()\n", 267 | "\n", 268 | " # Print train/test loss/accuracy\n", 269 | " if counter % 10 == 0:\n", 270 | " print(f\"Iteration: {counter} \\t Train Loss: {loss_val.item()}\")\n", 271 | " counter += 1\n", 272 | "\n", 273 | " if counter == 100:\n", 274 | " break\n", 275 | "\n", 276 | "training_loop(net, train_loader)" 277 | ], 278 | "metadata": { 279 | "id": "1telBMU-HrIg" 280 | }, 281 | "execution_count": null, 282 | "outputs": [] 283 | }, 284 | { 285 | "cell_type": "code", 286 | "source": [ 287 | "def measure_accuracy(model, dataloader):\n", 288 | " with torch.no_grad():\n", 289 | " model.eval()\n", 290 | " running_length = 0\n", 291 | " running_accuracy = 0\n", 292 | "\n", 293 | " for data, targets in iter(dataloader):\n", 294 | " data = data.to(device)\n", 295 | " targets = targets.to(device)\n", 296 | "\n", 297 | " # forward-pass\n", 298 | " spk_rec, _ = model(data)\n", 299 | " spike_count = spk_rec.sum(0) # batch x num_outputs\n", 300 | " _, max_spike = spike_count.max(1)\n", 301 | "\n", 302 | " # correct classes for one batch\n", 303 | " num_correct = (max_spike == targets).sum()\n", 304 | "\n", 305 | " # total accuracy\n", 306 | " running_length += len(targets)\n", 307 | " running_accuracy += num_correct\n", 308 | "\n", 309 | " accuracy = (running_accuracy / running_length)\n", 310 | "\n", 311 | " return accuracy.item()\n" 312 | ], 313 | "metadata": { 314 | "id": "nHsdkuSlIS1E" 315 | }, 316 | "execution_count": null, 317 | "outputs": [] 318 | }, 319 | { 320 | "cell_type": "code", 321 | "source": [ 322 | "print(f\"Test set accuracy: {measure_accuracy(net, test_loader)}\")" 323 | ], 324 | "metadata": { 325 | "id": "oJHAltRCKGyx" 326 | }, 327 | "execution_count": null, 328 | "outputs": [] 329 | }, 330 | { 331 | "cell_type": "markdown", 332 | "source": [ 333 | "# 2. Hardware Friendly Training\n", 334 | "## 2.1 Weight Quantization" 335 | ], 336 | "metadata": { 337 | "id": "Z9vrb2zUD6S-" 338 | } 339 | }, 340 | { 341 | "cell_type": "code", 342 | "source": [ 343 | "!pip install brevitas --quiet" 344 | ], 345 | "metadata": { 346 | "id": "icK4WzuL-2QA" 347 | }, 348 | "execution_count": null, 349 | "outputs": [] 350 | }, 351 | { 352 | "cell_type": "markdown", 353 | "source": [ 354 | "Just replace all `nn.Linear` layers with `qnn.QuantLinear(num_inputs, num_outputs, weight_bit_width, bias)`." 355 | ], 356 | "metadata": { 357 | "id": "idiLnVjJGJAL" 358 | } 359 | }, 360 | { 361 | "cell_type": "code", 362 | "source": [ 363 | "import brevitas.nn as qnn\n", 364 | "\n", 365 | "# Define Network\n", 366 | "class QuantNet(nn.Module):\n", 367 | " def __init__(self):\n", 368 | " super().__init__()\n", 369 | "\n", 370 | " # Initialize layers\n", 371 | " self.fc1 =\n", 372 | " self.lif1 =\n", 373 | " self.fc2 =\n", 374 | " self.lif2 =\n", 375 | "\n", 376 | " def forward(self, x):\n", 377 | "\n", 378 | " # Initialize hidden states at t=0\n", 379 | " mem1 = self.lif1.init_leaky()\n", 380 | " mem2 = self.lif2.init_leaky()\n", 381 | "\n", 382 | " # Record the final layer\n", 383 | " spk2_rec = []\n", 384 | " mem2_rec = []\n", 385 | "\n", 386 | " for step in range(num_steps):\n", 387 | " cur1 = self.fc1(x.flatten(1))\n", 388 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 389 | " cur2 = self.fc2(spk1)\n", 390 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 391 | " spk2_rec.append(spk2)\n", 392 | " mem2_rec.append(mem2)\n", 393 | "\n", 394 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 395 | "\n", 396 | "# Load the network onto CUDA if available\n", 397 | "qnet = QuantNet().to(device)" 398 | ], 399 | "metadata": { 400 | "id": "PiZYTMA6D-YL" 401 | }, 402 | "execution_count": null, 403 | "outputs": [] 404 | }, 405 | { 406 | "cell_type": "code", 407 | "source": [ 408 | "training_loop(qnet, train_loader)\n", 409 | "print(f\"Test set accuracy: {measure_accuracy(qnet, test_loader)}\")" 410 | ], 411 | "metadata": { 412 | "id": "PocSa27MOOoK" 413 | }, 414 | "execution_count": null, 415 | "outputs": [] 416 | }, 417 | { 418 | "cell_type": "markdown", 419 | "source": [ 420 | "## 2.2 Stateful Quantization\n", 421 | "\n" 422 | ], 423 | "metadata": { 424 | "id": "JvRzLHdjD-6S" 425 | } 426 | }, 427 | { 428 | "cell_type": "code", 429 | "source": [ 430 | "from snntorch.functional import quant\n", 431 | "\n", 432 | "# Define Network\n", 433 | "class SquatNet(nn.Module):\n", 434 | " def __init__(self):\n", 435 | " super().__init__()\n", 436 | "\n", 437 | " # Define state quantization parameters\n", 438 | " q_lif = quant.state_quant(num_bits=4, uniform=True)\n", 439 | "\n", 440 | " # Initialize layers\n", 441 | " self.fc1 =\n", 442 | " self.lif1 =\n", 443 | " self.fc2 =\n", 444 | " self.lif2 =\n", 445 | "\n", 446 | " def forward(self, x):\n", 447 | "\n", 448 | " # Initialize hidden states at t=0\n", 449 | " mem1 = self.lif1.init_leaky()\n", 450 | " mem2 = self.lif2.init_leaky()\n", 451 | "\n", 452 | " # Record the final layer\n", 453 | " spk2_rec = []\n", 454 | " mem2_rec = []\n", 455 | "\n", 456 | " for step in range(num_steps):\n", 457 | " cur1 = self.fc1(x.flatten(1))\n", 458 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 459 | " cur2 = self.fc2(spk1)\n", 460 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 461 | " spk2_rec.append(spk2)\n", 462 | " mem2_rec.append(mem2)\n", 463 | "\n", 464 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 465 | "\n", 466 | "# Load the network onto CUDA if available\n", 467 | "sqnet = SquatNet().to(device)" 468 | ], 469 | "metadata": { 470 | "id": "tAaafUczEAu_" 471 | }, 472 | "execution_count": null, 473 | "outputs": [] 474 | }, 475 | { 476 | "cell_type": "code", 477 | "source": [ 478 | "training_loop(sqnet, train_loader)\n", 479 | "print(f\"Test set accuracy: {measure_accuracy(sqnet, test_loader)}\")" 480 | ], 481 | "metadata": { 482 | "id": "0zfo0h6AQDzS" 483 | }, 484 | "execution_count": null, 485 | "outputs": [] 486 | }, 487 | { 488 | "cell_type": "markdown", 489 | "source": [ 490 | "# 3. Manual Non-Idealities\n", 491 | "## 3.1 Noise Injection" 492 | ], 493 | "metadata": { 494 | "id": "EZAFVgHW4M2g" 495 | } 496 | }, 497 | { 498 | "cell_type": "code", 499 | "source": [ 500 | "# Define Network\n", 501 | "class MemNet(nn.Module):\n", 502 | " def __init__(self):\n", 503 | " super().__init__()\n", 504 | "\n", 505 | " q_lif = quant.state_quant(num_bits=4, uniform=True)\n", 506 | "\n", 507 | " # Initialize layers\n", 508 | " self.fc1 = qnn.QuantLinear(num_inputs, num_hidden, weight_bit_width=8, bias=False)\n", 509 | " self.lif1 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 510 | " self.fc2 = qnn.QuantLinear(num_hidden, num_outputs, weight_bit_width=8, bias=False)\n", 511 | " self.lif2 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 512 | "\n", 513 | " # Introduce readout noise\n", 514 | " self.lif1_rand = torch.randn((num_hidden), device=device)\n", 515 | " self.lif2_rand = torch.randn((num_outputs), device=device)\n", 516 | "\n", 517 | " def forward(self, x):\n", 518 | "\n", 519 | " # Initialize hidden states at t=0\n", 520 | " mem1 = self.lif1.init_leaky()\n", 521 | " mem2 = self.lif2.init_leaky()\n", 522 | "\n", 523 | " # Record the final layer\n", 524 | " spk2_rec = []\n", 525 | " mem2_rec = []\n", 526 | "\n", 527 | " for step in range(num_steps):\n", 528 | " cur1 = self.fc1(x.flatten(1))\n", 529 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 530 | " mem1 = mem1 + (mem1 * self.lif1_rand) # apply noise\n", 531 | " cur2 = self.fc2(spk1)\n", 532 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 533 | " mem2 = mem2 + (mem2 * self.lif2_rand) # apply noise\n", 534 | " spk2_rec.append(spk2)\n", 535 | " mem2_rec.append(mem2)\n", 536 | "\n", 537 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 538 | "\n", 539 | "# Load the network onto CUDA if available\n", 540 | "memnet = MemNet().to(device)" 541 | ], 542 | "metadata": { 543 | "id": "-HEhzeA02Ueg" 544 | }, 545 | "execution_count": null, 546 | "outputs": [] 547 | }, 548 | { 549 | "cell_type": "code", 550 | "source": [ 551 | "training_loop(memnet, train_loader)\n", 552 | "print(f\"Test set accuracy: {measure_accuracy(memnet, test_loader)}\")" 553 | ], 554 | "metadata": { 555 | "id": "V5havFWd3c-F" 556 | }, 557 | "execution_count": null, 558 | "outputs": [] 559 | }, 560 | { 561 | "cell_type": "markdown", 562 | "source": [ 563 | "* ADC saturation limits using `torch.clamp(input, min=None, max=None)`\n", 564 | "* Weight noise by applying random matrix to weight matrices\n", 565 | "* Activation quantization by discretizing accumulated values\n", 566 | "* Other features can be factored in much the same way, e.g., conductance drift, non-uniform quantization with ADCs, partial digital summation, etc." 567 | ], 568 | "metadata": { 569 | "id": "oPnlZxqd2KQf" 570 | } 571 | }, 572 | { 573 | "cell_type": "markdown", 574 | "source": [ 575 | "That's all folks!" 576 | ], 577 | "metadata": { 578 | "id": "pNrFs3ro-xUm" 579 | } 580 | } 581 | ] 582 | } -------------------------------------------------------------------------------- /ESSCIRC_OSN_cheatsheet.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "provenance": [], 7 | "gpuType": "T4" 8 | }, 9 | "kernelspec": { 10 | "name": "python3", 11 | "display_name": "Python 3" 12 | }, 13 | "language_info": { 14 | "name": "python" 15 | }, 16 | "accelerator": "GPU" 17 | }, 18 | "cells": [ 19 | { 20 | "cell_type": "markdown", 21 | "source": [ 22 | "#Open-Source Neuromorphic Circuit Design\n", 23 | "## By Jason K. Eshraghian (www.ncg.ucsc.edu)\n", 24 | "\n", 25 | "\n", 26 | "[](https://github.com/jeshraghian/snntorch/)" 27 | ], 28 | "metadata": { 29 | "id": "8w6lhn7H8fW5" 30 | } 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": { 36 | "id": "1BlTqunB73-t" 37 | }, 38 | "outputs": [], 39 | "source": [ 40 | "!pip install snntorch --quiet # shift + enter" 41 | ] 42 | }, 43 | { 44 | "cell_type": "markdown", 45 | "source": [ 46 | "*What will I learn?*\n", 47 | "\n", 48 | "1. Train an SNN classifier using snnTorch\n", 49 | "2. Hardware Friendly Training\n", 50 | " - Weight Quantization with Brevitas\n", 51 | " - Stateful Quantization\n", 52 | "3. Manual Non-Idealities, e.g., read noise" 53 | ], 54 | "metadata": { 55 | "id": "xLlnINAI9mgJ" 56 | } 57 | }, 58 | { 59 | "cell_type": "markdown", 60 | "source": [ 61 | "# 1. Train an SNN Classifier using snnTorch\n", 62 | "## 1.1 Imports\n" 63 | ], 64 | "metadata": { 65 | "id": "-CS50cwuCW6n" 66 | } 67 | }, 68 | { 69 | "cell_type": "code", 70 | "source": [ 71 | "# snntorch imports\n", 72 | "import snntorch as snn\n", 73 | "from snntorch import functional as SF\n", 74 | "\n", 75 | "# pytorch imports\n", 76 | "import torch\n", 77 | "import torch.nn as nn\n", 78 | "from torch.utils.data import DataLoader\n", 79 | "from torchvision import datasets, transforms\n", 80 | "\n", 81 | "# data manipulation\n", 82 | "import numpy as np\n", 83 | "import itertools\n", 84 | "\n", 85 | "# plotting\n", 86 | "import matplotlib.pyplot as plt\n", 87 | "from IPython.display import HTML" 88 | ], 89 | "metadata": { 90 | "id": "H_TzogsCCcSe" 91 | }, 92 | "execution_count": null, 93 | "outputs": [] 94 | }, 95 | { 96 | "cell_type": "markdown", 97 | "source": [ 98 | "## 1.2 Boilerplate: DataLoading the MNIST Dataset" 99 | ], 100 | "metadata": { 101 | "id": "nftOdpyAGv7D" 102 | } 103 | }, 104 | { 105 | "cell_type": "code", 106 | "source": [ 107 | "# dataloader arguments\n", 108 | "batch_size = 128\n", 109 | "data_path='/data/mnist'\n", 110 | "\n", 111 | "dtype = torch.float\n", 112 | "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n", 113 | "## if you're on M1 or M2 GPU:\n", 114 | "# device = torch.device(\"mps\")" 115 | ], 116 | "metadata": { 117 | "id": "SsM2Z5NXGu5z" 118 | }, 119 | "execution_count": null, 120 | "outputs": [] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "source": [ 125 | "# Define a transform\n", 126 | "transform = transforms.Compose([\n", 127 | " transforms.Resize((28, 28)),\n", 128 | " transforms.Grayscale(),\n", 129 | " transforms.ToTensor(),\n", 130 | " transforms.Normalize((0,), (1,))])\n", 131 | "\n", 132 | "mnist_train = datasets.MNIST(data_path, train=True, download=True, transform=transform)\n", 133 | "mnist_test = datasets.MNIST(data_path, train=False, download=True, transform=transform)" 134 | ], 135 | "metadata": { 136 | "id": "XqbYptgDHUPg" 137 | }, 138 | "execution_count": null, 139 | "outputs": [] 140 | }, 141 | { 142 | "cell_type": "code", 143 | "source": [ 144 | "# Create DataLoaders\n", 145 | "train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True, drop_last=True)\n", 146 | "test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=True, drop_last=True)" 147 | ], 148 | "metadata": { 149 | "id": "jSlS3gWZHXI0" 150 | }, 151 | "execution_count": null, 152 | "outputs": [] 153 | }, 154 | { 155 | "cell_type": "markdown", 156 | "source": [ 157 | "## 1.3 Construct SNN Model" 158 | ], 159 | "metadata": { 160 | "id": "SJVhfNukHbsp" 161 | } 162 | }, 163 | { 164 | "cell_type": "code", 165 | "source": [ 166 | "# Network Architecture\n", 167 | "num_inputs = 28*28\n", 168 | "num_hidden = 100\n", 169 | "num_outputs = 10\n", 170 | "\n", 171 | "# Temporal Dynamics\n", 172 | "num_steps = 25\n", 173 | "beta = 0.95" 174 | ], 175 | "metadata": { 176 | "id": "uu324fr_HhxV" 177 | }, 178 | "execution_count": null, 179 | "outputs": [] 180 | }, 181 | { 182 | "cell_type": "code", 183 | "source": [ 184 | "from snntorch import surrogate\n", 185 | "\n", 186 | "# Define Network\n", 187 | "class Net(nn.Module):\n", 188 | " def __init__(self):\n", 189 | " super().__init__()\n", 190 | "\n", 191 | " # Initialize layers\n", 192 | " self.fc1 = nn.Linear(num_inputs, num_hidden)\n", 193 | " self.lif1 = snn.Leaky(beta=beta)\n", 194 | " self.fc2 = nn.Linear(num_hidden, num_outputs)\n", 195 | " self.lif2 = snn.Leaky(beta=beta)\n", 196 | "\n", 197 | " def forward(self, x):\n", 198 | "\n", 199 | " # Initialize hidden states at t=0\n", 200 | " mem1 = self.lif1.init_leaky()\n", 201 | " mem2 = self.lif2.init_leaky()\n", 202 | "\n", 203 | " # Record the final layer\n", 204 | " spk2_rec = []\n", 205 | " mem2_rec = []\n", 206 | "\n", 207 | " # time-loop\n", 208 | " for step in range(num_steps):\n", 209 | " cur1 = self.fc1(x.flatten(1)) # batch: 128 x 784\n", 210 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 211 | " cur2 = self.fc2(spk1)\n", 212 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 213 | "\n", 214 | " # store in list\n", 215 | " spk2_rec.append(spk2)\n", 216 | " mem2_rec.append(mem2)\n", 217 | "\n", 218 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0) # time-steps x batch x num_out\n", 219 | "\n", 220 | "# Load the network onto CUDA if available\n", 221 | "net = Net().to(device)" 222 | ], 223 | "metadata": { 224 | "id": "CkM1Z1EjHeW8" 225 | }, 226 | "execution_count": null, 227 | "outputs": [] 228 | }, 229 | { 230 | "cell_type": "markdown", 231 | "source": [ 232 | "## 1.4 Training the SNN" 233 | ], 234 | "metadata": { 235 | "id": "p8qBw03rHpn3" 236 | } 237 | }, 238 | { 239 | "cell_type": "code", 240 | "source": [ 241 | "def training_loop(model, dataloader, num_epochs=1):\n", 242 | " loss = nn.CrossEntropyLoss()\n", 243 | " optimizer = torch.optim.Adam(model.parameters(), lr=5e-4, betas=(0.9, 0.999))\n", 244 | " counter = 0\n", 245 | "\n", 246 | " # Outer training loop\n", 247 | " for epoch in range(num_epochs):\n", 248 | " train_batch = iter(dataloader)\n", 249 | "\n", 250 | " # Minibatch training loop\n", 251 | " for data, targets in train_batch:\n", 252 | " data = data.to(device)\n", 253 | " targets = targets.to(device)\n", 254 | "\n", 255 | " # forward pass\n", 256 | " model.train()\n", 257 | " spk_rec, _ = model(data)\n", 258 | "\n", 259 | " # initialize the loss & sum over time\n", 260 | " loss_val = torch.zeros((1), dtype=dtype, device=device)\n", 261 | " loss_val = loss(spk_rec.sum(0), targets) # batch x num_out\n", 262 | "\n", 263 | " # Gradient calculation + weight update\n", 264 | " optimizer.zero_grad()\n", 265 | " loss_val.backward()\n", 266 | " optimizer.step()\n", 267 | "\n", 268 | " # Print train/test loss/accuracy\n", 269 | " if counter % 10 == 0:\n", 270 | " print(f\"Iteration: {counter} \\t Train Loss: {loss_val.item()}\")\n", 271 | " counter += 1\n", 272 | "\n", 273 | " if counter == 100:\n", 274 | " break\n", 275 | "\n", 276 | "training_loop(net, train_loader)" 277 | ], 278 | "metadata": { 279 | "id": "1telBMU-HrIg" 280 | }, 281 | "execution_count": null, 282 | "outputs": [] 283 | }, 284 | { 285 | "cell_type": "code", 286 | "source": [ 287 | "def measure_accuracy(model, dataloader):\n", 288 | " with torch.no_grad():\n", 289 | " model.eval()\n", 290 | " running_length = 0\n", 291 | " running_accuracy = 0\n", 292 | "\n", 293 | " for data, targets in iter(dataloader):\n", 294 | " data = data.to(device)\n", 295 | " targets = targets.to(device)\n", 296 | "\n", 297 | " # forward-pass\n", 298 | " spk_rec, _ = model(data)\n", 299 | " spike_count = spk_rec.sum(0) # batch x num_outputs\n", 300 | " _, max_spike = spike_count.max(1)\n", 301 | "\n", 302 | " # correct classes for one batch\n", 303 | " num_correct = (max_spike == targets).sum()\n", 304 | "\n", 305 | " # total accuracy\n", 306 | " running_length += len(targets)\n", 307 | " running_accuracy += num_correct\n", 308 | "\n", 309 | " accuracy = (running_accuracy / running_length)\n", 310 | "\n", 311 | " return accuracy.item()\n" 312 | ], 313 | "metadata": { 314 | "id": "nHsdkuSlIS1E" 315 | }, 316 | "execution_count": null, 317 | "outputs": [] 318 | }, 319 | { 320 | "cell_type": "code", 321 | "source": [ 322 | "print(f\"Test set accuracy: {measure_accuracy(net, test_loader)}\")" 323 | ], 324 | "metadata": { 325 | "id": "oJHAltRCKGyx" 326 | }, 327 | "execution_count": null, 328 | "outputs": [] 329 | }, 330 | { 331 | "cell_type": "markdown", 332 | "source": [ 333 | "# 2. Hardware Friendly Training\n", 334 | "## 2.1 Weight Quantization" 335 | ], 336 | "metadata": { 337 | "id": "Z9vrb2zUD6S-" 338 | } 339 | }, 340 | { 341 | "cell_type": "code", 342 | "source": [ 343 | "!pip install brevitas --quiet" 344 | ], 345 | "metadata": { 346 | "id": "icK4WzuL-2QA" 347 | }, 348 | "execution_count": null, 349 | "outputs": [] 350 | }, 351 | { 352 | "cell_type": "markdown", 353 | "source": [ 354 | "Just replace all `nn.Linear` layers with `qnn.QuantLinear(num_inputs, num_outputs, weight_bit_width, bias)`." 355 | ], 356 | "metadata": { 357 | "id": "idiLnVjJGJAL" 358 | } 359 | }, 360 | { 361 | "cell_type": "code", 362 | "source": [ 363 | "import brevitas.nn as qnn\n", 364 | "\n", 365 | "# Define Network\n", 366 | "class QuantNet(nn.Module):\n", 367 | " def __init__(self):\n", 368 | " super().__init__()\n", 369 | "\n", 370 | " # Initialize layers\n", 371 | " self.fc1 = qnn.QuantLinear(num_inputs, num_hidden, weight_bit_width=8, bias=False)\n", 372 | " self.lif1 = snn.Leaky(beta=beta)\n", 373 | " self.fc2 = qnn.QuantLinear(num_hidden, num_outputs, weight_bit_width=8, bias=False)\n", 374 | " self.lif2 = snn.Leaky(beta=beta)\n", 375 | "\n", 376 | " def forward(self, x):\n", 377 | "\n", 378 | " # Initialize hidden states at t=0\n", 379 | " mem1 = self.lif1.init_leaky()\n", 380 | " mem2 = self.lif2.init_leaky()\n", 381 | "\n", 382 | " # Record the final layer\n", 383 | " spk2_rec = []\n", 384 | " mem2_rec = []\n", 385 | "\n", 386 | " for step in range(num_steps):\n", 387 | " cur1 = self.fc1(x.flatten(1))\n", 388 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 389 | " cur2 = self.fc2(spk1)\n", 390 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 391 | " spk2_rec.append(spk2)\n", 392 | " mem2_rec.append(mem2)\n", 393 | "\n", 394 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 395 | "\n", 396 | "# Load the network onto CUDA if available\n", 397 | "qnet = QuantNet().to(device)" 398 | ], 399 | "metadata": { 400 | "id": "PiZYTMA6D-YL" 401 | }, 402 | "execution_count": null, 403 | "outputs": [] 404 | }, 405 | { 406 | "cell_type": "code", 407 | "source": [ 408 | "training_loop(qnet, train_loader)\n", 409 | "print(f\"Test set accuracy: {measure_accuracy(qnet, test_loader)}\")" 410 | ], 411 | "metadata": { 412 | "id": "PocSa27MOOoK" 413 | }, 414 | "execution_count": null, 415 | "outputs": [] 416 | }, 417 | { 418 | "cell_type": "markdown", 419 | "source": [ 420 | "## 2.2 Stateful Quantization\n", 421 | "\n" 422 | ], 423 | "metadata": { 424 | "id": "JvRzLHdjD-6S" 425 | } 426 | }, 427 | { 428 | "cell_type": "code", 429 | "source": [ 430 | "from snntorch.functional import quant\n", 431 | "\n", 432 | "# Define Network\n", 433 | "class SquatNet(nn.Module):\n", 434 | " def __init__(self):\n", 435 | " super().__init__()\n", 436 | "\n", 437 | " q_lif = quant.state_quant(num_bits=4, uniform=True)\n", 438 | "\n", 439 | " # Initialize layers\n", 440 | " self.fc1 = qnn.QuantLinear(num_inputs, num_hidden, weight_bit_width=8, bias=False)\n", 441 | " self.lif1 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 442 | " self.fc2 = qnn.QuantLinear(num_hidden, num_outputs, weight_bit_width=8, bias=False)\n", 443 | " self.lif2 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 444 | "\n", 445 | " def forward(self, x):\n", 446 | "\n", 447 | " # Initialize hidden states at t=0\n", 448 | " mem1 = self.lif1.init_leaky()\n", 449 | " mem2 = self.lif2.init_leaky()\n", 450 | "\n", 451 | " # Record the final layer\n", 452 | " spk2_rec = []\n", 453 | " mem2_rec = []\n", 454 | "\n", 455 | " for step in range(num_steps):\n", 456 | " cur1 = self.fc1(x.flatten(1))\n", 457 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 458 | " cur2 = self.fc2(spk1)\n", 459 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 460 | " spk2_rec.append(spk2)\n", 461 | " mem2_rec.append(mem2)\n", 462 | "\n", 463 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 464 | "\n", 465 | "# Load the network onto CUDA if available\n", 466 | "sqnet = SquatNet().to(device)" 467 | ], 468 | "metadata": { 469 | "id": "tAaafUczEAu_" 470 | }, 471 | "execution_count": null, 472 | "outputs": [] 473 | }, 474 | { 475 | "cell_type": "code", 476 | "source": [ 477 | "training_loop(sqnet, train_loader)\n", 478 | "print(f\"Test set accuracy: {measure_accuracy(sqnet, test_loader)}\")" 479 | ], 480 | "metadata": { 481 | "id": "0zfo0h6AQDzS" 482 | }, 483 | "execution_count": null, 484 | "outputs": [] 485 | }, 486 | { 487 | "cell_type": "markdown", 488 | "source": [ 489 | "# 3. Manual Non-Idealities\n", 490 | "## 3.1 Noise Injection" 491 | ], 492 | "metadata": { 493 | "id": "EZAFVgHW4M2g" 494 | } 495 | }, 496 | { 497 | "cell_type": "code", 498 | "source": [ 499 | "# Define Network\n", 500 | "class MemNet(nn.Module):\n", 501 | " def __init__(self):\n", 502 | " super().__init__()\n", 503 | "\n", 504 | " q_lif = quant.state_quant(num_bits=4, uniform=True)\n", 505 | "\n", 506 | " # Initialize layers\n", 507 | " self.fc1 = qnn.QuantLinear(num_inputs, num_hidden, weight_bit_width=8, bias=False)\n", 508 | " self.lif1 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 509 | " self.fc2 = qnn.QuantLinear(num_hidden, num_outputs, weight_bit_width=8, bias=False)\n", 510 | " self.lif2 = snn.Leaky(beta=beta, state_quant=q_lif)\n", 511 | "\n", 512 | " self.lif1_rand = torch.randn((num_hidden), device=device)\n", 513 | " self.lif2_rand = torch.randn((num_outputs), device=device)\n", 514 | "\n", 515 | " def forward(self, x):\n", 516 | "\n", 517 | " # Initialize hidden states at t=0\n", 518 | " mem1 = self.lif1.init_leaky()\n", 519 | " mem2 = self.lif2.init_leaky()\n", 520 | "\n", 521 | " # Record the final layer\n", 522 | " spk2_rec = []\n", 523 | " mem2_rec = []\n", 524 | "\n", 525 | " for step in range(num_steps):\n", 526 | " cur1 = self.fc1(x.flatten(1))\n", 527 | " spk1, mem1 = self.lif1(cur1, mem1)\n", 528 | " mem1 = mem1 + (mem1 * self.lif1_rand)\n", 529 | " cur2 = self.fc2(spk1)\n", 530 | " spk2, mem2 = self.lif2(cur2, mem2)\n", 531 | " mem2 = mem2 + (mem2 * self.lif2_rand)\n", 532 | " spk2_rec.append(spk2)\n", 533 | " mem2_rec.append(mem2)\n", 534 | "\n", 535 | " return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)\n", 536 | "\n", 537 | "# Load the network onto CUDA if available\n", 538 | "memnet = MemNet().to(device)" 539 | ], 540 | "metadata": { 541 | "id": "-HEhzeA02Ueg" 542 | }, 543 | "execution_count": null, 544 | "outputs": [] 545 | }, 546 | { 547 | "cell_type": "code", 548 | "source": [ 549 | "training_loop(memnet, train_loader)\n", 550 | "print(f\"Test set accuracy: {measure_accuracy(memnet, test_loader)}\")" 551 | ], 552 | "metadata": { 553 | "id": "V5havFWd3c-F" 554 | }, 555 | "execution_count": null, 556 | "outputs": [] 557 | }, 558 | { 559 | "cell_type": "markdown", 560 | "source": [ 561 | "* ADC saturation limits using `torch.clamp(input, min=None, max=None)`\n", 562 | "* Weight noise by applying random matrix to weight matrices\n", 563 | "* Activation quantization by discretizing accumulated values\n", 564 | "* Other features can be factored in much the same way, e.g., conductance drift, non-uniform quantization with ADCs, partial digital summation, etc." 565 | ], 566 | "metadata": { 567 | "id": "NpFA9d7_3nNV" 568 | } 569 | }, 570 | { 571 | "cell_type": "markdown", 572 | "source": [ 573 | "That's all folks!" 574 | ], 575 | "metadata": { 576 | "id": "pNrFs3ro-xUm" 577 | } 578 | } 579 | ] 580 | } -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Jason Eshraghian 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ESSCIRC/ESSDIRC 2023: Open-Source Neuromorphic Circuit Design Tutorial 2 | 3 | This repository contains the notebooks related to hardware-aware training of spiking neural networks presented at ESSCIRC/ESSDIRC 2023 (Lisbon, Portugal) for the tutorial on *Open-Source Neuromorphic Circuit Design* co-presented by Charlotte Frenkel, Jason Eshraghian, and Rajit Manohar. 4 | 5 | ![Abstract](./graphical-abstract.jpg) 6 | 7 | There are two tutorial notebooks in this repo: 8 | 9 | * One for hardware-aware training spiking neural networks using [snnTorch](https://github.com/jeshraghian/snntorch) 10 | * One for running several neuromorphic designs through the OpenLane flow using the Sky130 PDK 11 | 12 | ## Notebooks 13 | 14 | | Title | Colab Link | 15 | |-------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------| 16 | | Hardware-Aware Training of Spiking Neural Networks with [snnTorch](https://github.com/jeshraghian/snntorch) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jeshraghian/ESSCIRC23-os-neuromorphic-tutorial/blob/main/ESSCIRC_OSN.ipynb) | 17 | | **Cheat-Sheet:** Hardware-Aware Training of Spiking Neural Networks with [snnTorch](https://github.com/jeshraghian/snntorch) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jeshraghian/ESSCIRC23-os-neuromorphic-tutorial/blob/main/ESSCIRC_OSN_cheatsheet.ipynb) | 18 | | Neuromorphic Accelerator Design with Sky130 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jeshraghian/ESSCIRC23-os-neuromorphic-tutorial/blob/main/ESSCIRC_OSN_Sky130.ipynb) | 19 | -------------------------------------------------------------------------------- /graphical-abstract.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jeshraghian/ESSCIRC23-os-neuromorphic-tutorial/09eb8d94d23bcfc04b76e785ef76f21c3625d3a0/graphical-abstract.jpg --------------------------------------------------------------------------------