fangwei123456/spikingjelly

SpikingJelly – PyTorch‑native framework for Spiking Neural Networks (SNNs)

What it is

• An open‑source Python library built on top of PyTorch that lets you define, train, and deploy spiking neural networks – the brain‑inspired alternative to conventional artificial neural networks.
• Targets researchers and engineers who want to work with event‑based (neuromorphic) data, convert existing ANN models to SNNs, or run large‑scale SNN experiments.

Key capabilities

Area	What SpikingJelly provides
SNN modeling	Ready‑made spiking neuron modules (LIF, IF, etc.), surrogate‑gradient functions, stateful/stateless layers, and a collection of reference SNN architectures.
Training workflows	PyTorch‑style training loops, online‑learning utilities, and tools to convert a trained ANN into an SNN (ANN2SNN).
Performance	Automatic selection of the fastest execution backend (plain PyTorch, Triton kernels, or optional native CUDA extensions). Supports mixed‑precision (e.g., fp8) and a “FlexSN” kernel generator for custom neuron kernels.
Scaling	Memory‑efficient spike compression, experimental multi‑GPU distributed training, and precision‑policy helpers for big models.
Datasets	Built‑in loaders for many neuromorphic event‑based datasets (e.g., N‑MNIST, CIFAR10‑DVS, DVS‑Gesture, Speech Commands, etc.) with raw‑event and frame‑based access.
Analysis	Profilers for FLOPs, SynOps, memory traffic, and tools to estimate inference energy.
Interchange & deployment	Export helpers for neuromorphic hardware and other frameworks (NIR, Lava, Lynxi).

Quick start example

from torch import nn
from spikingjelly.activation_based import layer, neuron, surrogate

net = nn.Sequential(
    layer.Flatten(),
    layer.Linear(28*28, 10, bias=False),
    neuron.LIFNode(tau=2.0, surrogate_function=surrogate.ATan())
)

The model is defined exactly like any PyTorch nn.Module; you can then train it on event‑based datasets or convert a conventional ANN to an SNN using the provided tutorials.

Installation

• Requires Python ≥ 3.11 and PyTorch ≥ 2.6.
• Install via the uv package manager (or pip) from PyPI:

  uv pip install spikingjelly               # stable release
  uv pip install "spikingjelly[triton]"      # include optional Triton kernels
  

• For maximum performance you can build the optional native CUDA extensions:

  SJ_BUILD_NATIVE_CUDA=1 uv pip install \
    --no-build-isolation --no-binary spikingjelly \
    --reinstall-package spikingjelly --no-cache "spikingjelly>=2.0.0"
  

• Detailed installation guide and troubleshooting are in the docs.

Typical use cases

• Research on neuromorphic vision or audio (e.g., DVS cameras, event‑based speech).
• Converting a high‑accuracy ANN into a low‑power SNN for edge devices.
• Benchmarking spiking transformer components or large‑scale SNN training on multi‑GPU clusters.
• Exporting models to neuromorphic chips or simulation frameworks via NIR/Lava.

Project health

• Actively maintained (maintainers listed, regular releases, versioning follows SemVer).
• Documentation hosted on Read the Docs, with tutorials for installation, neuron basics, dataset handling, and deployment.
• Apache‑2.0 license – permissive for both academic and commercial use.

Where to learn more

• Full docs: https://spikingjelly.readthedocs.io/
• Quick‑start tutorials, dataset guides, and conversion notebooks are linked from the README.
• Community contributions welcomed via the GitHub issue tracker and pull‑request workflow.

---

SpikingJelly is essentially the “PyTorch for spiking neural networks,” offering a familiar API while handling the unique challenges of event‑driven computation.