Repository files navigation

The brain is the perfect place to look for inspiration to develop more efficient neural networks. One of the main differences with modern deep learning is that the brain encodes information in spikes rather than continuous activations.snnTorch is a Python package for performing gradient-based learning with spiking neural networks.It extends the capabilities of PyTorch, taking advantage of its GPU accelerated tensorcomputation and applying it to networks of spiking neurons. Pre-designed spiking neuron models are seamlessly integrated within the PyTorch framework and can be treated as recurrent activation units.

If you like this project, please consider starring ⭐ this repo as it is the easiest and best way to support it.

If you have issues, comments, or are looking for advice on training spiking neural networks, you can open an issue, a discussion, or chat in ourdiscord channel.

snnTorch contains the following components:

snnTorch is designed to be intuitively used with PyTorch, as though each spiking neuron were simply another activation in a sequence of layers.It is therefore agnostic to fully-connected layers, convolutional layers, residual connections, etc.

At present, the neuron models are represented by recursive functions which removes the need to store membrane potential traces for all neurons in a system in order to calculate the gradient.The lean requirements of snnTorch enable small and large networks to be viably trained on CPU, where needed.Provided that the network models and tensors are loaded onto CUDA, snnTorch takes advantage of GPU acceleration in the same way as PyTorch.

If you find snnTorch useful in your work, please cite the following source:

Jason K. Eshraghian, Max Ward, Emre Neftci, Xinxin Wang, Gregor Lenz, GirishDwivedi, Mohammed Bennamoun, Doo Seok Jeong, and Wei D. Lu “TrainingSpiking Neural Networks Using Lessons From Deep Learning”. Proceedings of the IEEE, 111(9)September 2023.

@article{eshraghian2021training,        title   =  {Training spiking neural networks using lessons from deep learning},        author  =  {Eshraghian, Jason K and Ward, Max and Neftci, Emre and Wang, Xinxin                    and Lenz, Gregor and Dwivedi, Girish and Bennamoun, Mohammed and                   Jeong, Doo Seok and Lu, Wei D},        journal = {Proceedings of the IEEE},        volume  = {111},        number  = {9},        pages   = {1016--1054},        year    = {2023}}

Let us know if you are using snnTorch in any interesting work, research or blogs, as we would love to hear more about it! Reach out atsnntorch@gmail.com.

PyTorch should be installed to use snnTorch. Ensure the correct version of torch is installed for your system to enable CUDA compatibility.

The following packages are automatically installed if using the pip command:

• numpy
• pandas

The following packages are required for using export_nir and import_nir:

• nir
• nirtorch

The following packages are required for using spikeplot:

• matplotlib

Run the following to install:

$ python$ pip install snntorch

To install snnTorch from source instead:

$ git clone https://github.com/jeshraghian/snnTorch$ cd snntorch$ python setup.py install

To install snntorch with conda:

$ conda install -c conda-forge snntorch

To install for an Intelligent Processing Units (IPU) based build using Graphcore's accelerators:

$ pip install snntorch-ipu

A complete API is availablehere. Examples, tutorials and Colab notebooks are provided.

Here are a few ways you can get started with snnTorch:

• Quickstart Notebook (Opens in Colab)
• The API Reference
• Examples
• Tutorials

For a quick example to run snnTorch, see the following snippet, or test the quickstart notebook:

importtorch,torch.nnasnnimportsnntorchassnnfromsnntorchimportsurrogatefromsnntorchimportutilsnum_steps=25# number of time stepsbatch_size=1beta=0.5# neuron decay ratespike_grad=surrogate.fast_sigmoid()# surrogate gradientnet=nn.Sequential(nn.Conv2d(1,8,5),nn.MaxPool2d(2),snn.Leaky(beta=beta,init_hidden=True,spike_grad=spike_grad),nn.Conv2d(8,16,5),nn.MaxPool2d(2),snn.Leaky(beta=beta,init_hidden=True,spike_grad=spike_grad),nn.Flatten(),nn.Linear(16*4*4,10),snn.Leaky(beta=beta,init_hidden=True,spike_grad=spike_grad,output=True)      )data_in=torch.rand(num_steps,batch_size,1,28,28)# random input dataspike_recording= []# record spikes over timeutils.reset(net)# reset/initialize hidden states for all neuronsforstepinrange(num_steps):# loop over timespike,state=net(data_in[step])# one time step of forward-passspike_recording.append(spike)# record spikes in list

If you wish to learn all the fundamentals of training spiking neural networks, from neuron models, to the neural code, up to backpropagation, the snnTorch tutorial series is a great place to begin.It consists of interactive notebooks with complete explanations that can get you up to speed.

Tutorial	Title	Colab Link
Tutorial 1	Spike Encoding with snnTorch
Tutorial 2	The Leaky Integrate and Fire Neuron
Tutorial 3	A Feedforward Spiking Neural Network
Tutorial 4	2nd Order Spiking Neuron Models (Optional)
Tutorial 5	Training Spiking Neural Networks with snnTorch
Tutorial 6	Surrogate Gradient Descent in a Convolutional SNN
Tutorial 7	Neuromorphic Datasets with Tonic + snnTorch

Advanced Tutorials	Colab Link
Population Coding
Regression: Part I - Membrane Potential Learning with LIF Neurons
Regression: Part II - Regression-based Classification with Recurrent LIF Neurons
Accelerating snnTorch on IPUs	—

If you're ready to contribute to snnTorch, instructions to do so can befound here.

snnTorch is currently maintained by theUCSC Neuromorphic Computing Group. It was initially developed byJason K. Eshraghian in theLu Group (University of Michigan).

Additional contributions were made byVincent Sun,Peng Zhou,Ridger Zhu,Alexander Henkes,Steven Abreu, Xinxin Wang, Sreyes Venkatesh,gekkom, and Emre Neftci.

snnTorch source code is published under the terms of the MIT License.snnTorch's documentation is licensed under a Creative Commons Attribution-Share Alike 3.0 Unported License (CC BY-SA 3.0).