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 tensor computation 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.
snnTorch contains the following components:
a spiking neuron library like torch.nn, deeply integrated with autograd
variations of backpropagation commonly used with SNNs
a library for spike generation and data conversion
visualization tools for spike-based data using matplotlib and celluloid
contains popular neuromorphic datasets
optional surrogate gradient functions
dataset utility functions
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.
The following packages need to be installed to use snnTorch:
torch >= 1.2.0
torchvision >= 0.9.1
numpy >= 1.17
h5py >= 3.0.0
They are automatically installed if snnTorch is installed using the pip command. Ensure the correct versions of torch and torchvision are installed for your system to enable CUDA compatibility.
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
API & Examples¶
A complete API is available here. Examples, tutorials and Colab notebooks are provided.
Here are a few ways you can get started with snnTorch:
Additional contributions were made by Xinxin Wang, Vincent Sun, and Emre Neftci.
Several features in snnTorch were inspired by the work of Friedemann Zenke, Emre Neftci, Doo Seok Jeong, Sumit Bam Shrestha and Garrick Orchard.
License & Copyright¶
snnTorch is licensed under the GNU General Public License v3.0: https://www.gnu.org/licenses/gpl-3.0.en.html.