SpikingJelly
中文 | English
Contents
Why SpikingJelly
SpikingJelly is a PyTorch-native framework for spiking neural networks (SNNs), with support for large-scale SNN training and inference.
- Beginner-friendly API
- ANN2SNN conversion
- Event-based datasets
- Automatic operator execution: Torch, Triton and optional native CUDA
- Memory-efficient training, distributed execution, precision control
- Hardware deployment and framework exchange
Installation
Use uv in an activated virtual environment: Python >=3.11, Torch >=2.6.
First install PyTorch.
| Install | Command |
|---|
| PyPI release | uv pip install spikingjelly |
| PyPI pre-release | uv pip install --pre spikingjelly |
| Optional Triton (if missing) | uv pip install "spikingjelly[triton]" |
| Latest development source | uv pip install git+https://github.com/fangwei123456/spikingjelly.git |
Native CUDA (requires a toolchain; available after the corresponding V2 PyPI release):
SJ_BUILD_NATIVE_CUDA=1 uv pip install \
--no-build-isolation --no-binary spikingjelly \
--reinstall-package spikingjelly --no-cache "spikingjelly>=2.0.0"
See the installation guide
for prerequisites, decision diagrams, other extras and troubleshooting.
Unreleased changes require the development source.
Quick Start
Define an SNN in the same way that you would define any PyTorch model:
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())
)
Next steps:
Core Capabilities
| Area | What SpikingJelly provides |
|---|
| SNN modeling | Activation-based SNN components: spiking neurons, surrogate gradients, stateful and stateless modules. Predefined SNN models. |
| Training workflows | PyTorch-native training flows, online-learning utilities, and ANN2SNN conversion |
| Performance | automatic Torch/Triton/native CUDA execution, FlexSN for customized neuron kernels, and mixed-precision training utilities (e.g., fp8) |
| Scaling | Memory-efficient training, and distributed training |
| Datasets | Neuromorphic datasets, and data preprocessing pipelines |
| Analysis | FLOPs / SynOps / memory-access profiling, and inference energy estimation |
| Interchange and deployment | NIR, Lava, and Lynxi-oriented exchange interfaces for neuromorphic workflows |
Operator Execution
Neuron execution is selected automatically from tensor device and execution path:
CPU uses Torch; CUDA eager prefers compatible native CUDA, while Inductor
expansion prefers Triton. CUDA Graph retains its pre-capture choice. No public
neuron backend argument and no CuPy dependency are required. Low-precision
state profiles can use Torch reference execution.
Pure Python wheels include Torch reference and optional Triton execution.
Optional native neuron, fused IF/LIF-Linear and packed/sparse projection builds
are described under Installation. Missing extensions use Torch reference
for projections; fused no-intermediate-spike performance requires the extension.
Use ordinary Linear/Conv plus memopt instead of retired SpikeLinear/SpikeConv.
Large-Scale SNN Systems
For large-scale SNN systems, SpikingJelly provides:
- Memory-efficient training with spike compression (
memopt)
- Experimental distributed execution for multi-GPU workloads
- Precision policy tools for large-scale training and inference
- Spiking transformer components
Datasets
SpikingJelly includes the following event-based and neuromorphic datasets:
This is the beginning of the README. Read the rest at the original. Read the full README ↗
