SpikingJelly

English | 中文(Chinese)

SpikingJelly 是一个基于 PyTorch 的开源脉冲神经网络（SNN）深度学习框架。

SpikingJelly 的 文档 同时提供英文和中文版本。

• 更新日志
• 安装
• 以空前简单的方式构建 SNN
• 快速便捷的 ANN-SNN 转换
• CUDA/Triton 加速的神经元
• 设备支持
• 神经形态数据集支持
• 教程
• 论文与引用
• 贡献
• 关于

更新日志

我们正在积极维护和改进 SpikingJelly。以下是我们的未来计划以及每次发布的主要亮点。

亮点

我们的新工作 通过梯度检查点和脉冲压缩实现无损、高效的脉冲神经网络训练 最近被 ICLR 2026 接受！自动训练内存优化工具现已在 spikingjelly.activation_based.memopt 中可用。更多信息请参阅 教程。

在最新版本（GitHub 版本）中：

• IFNode、LIFNode 和 ParametricLIFNode 现已配备 Triton 后端；
• FlexSN 可用于将 PyTorch 脉冲神经元动力学转换为 Triton 内核；
• SpikingSelfAttention 和 QKAttention 已上线；
• memopt 已上线；
• nir_exchange 已上线；
• op_counter 已上线；
• spikingjelly.activation_based.layer、spikingjelly.activation_based.functional 和 spikingjelly.datasets 已重构；
• 数据集实现已重构；
• 文档和教程已更新；
• spikingjelly.activation_based.functional 中的卷积-批归一化融合函数已被弃用；请改用 PyTorch 的 fuse_conv_bn_eval。

计划中

我们即将发布版本 0.0.0.1.0。

• 添加 Triton 后端，进一步加速 GPU 上的计算。
• 添加一个编译器，用于将 PyTorch 脉冲神经元转换为 Triton 内核，这将比现有的 auto_cuda 子包更加灵活。
• 添加脉冲自注意力实现。
• 更新文档和教程。

其他长期计划包括：

• 添加 NIR 支持。
• 优化训练内存消耗。
• 在华为 NPU 上加速。

对于早期实验性功能，请参阅我们的姊妹项目 flashsnn。新的想法会在 flashsnn 中进行原型验证，然后再合并到 SpikingJelly 中。

版本说明

• 奇数版本号为开发版，随 GitHub/OpenI 仓库同步更新。偶数版本号为稳定版，可在 PyPI 上获取。

• 默认文档适用于最新的开发版。如果您使用的是稳定版，请务必切换到对应版本的文档。

• 自版本 0.0.0.0.14 起，clock_driven 和 event_driven 等模块已被重命名。请参考教程 从旧版本迁移。

• 如果您使用的是较旧版本的 SpikingJelly，可能会遇到一些严重错误。更多详情请参阅 发行版中的错误历史。

不同版本的文档：

• zero
• 0.0.0.0.4
• 0.0.0.0.6
• 0.0.0.0.8
• 0.0.0.0.10
• 0.0.0.0.12
• 0.0.0.0.14
• 最新版

安装

请注意，SpikingJelly 基于 PyTorch。在安装 SpikingJelly 之前，请确保已安装 PyTorch、torchvision 和 torchaudio。需要注意的是，SpikingJelly 的最新版本要求 torch>=2.2.0，并在 torch==2.7.1 上进行了测试。

从 PyPI 安装最新稳定版：

pip install spikingjelly

从源代码安装最新开发版：

从 GitHub：

git clone https://github.com/fangwei123456/spikingjelly.git
cd spikingjelly
pip install .

从 OpenI：

git clone https://openi.pcl.ac.cn/OpenI/spikingjelly.git
cd spikingjelly
pip install .

可选依赖

要启用 cupy 后端，请安装 CuPy。

pip install cupy-cuda12x # 适用于 CUDA 12.x
pip install cupy-cuda11x # 适用于 CUDA 11.x

要启用 triton 后端，请确保已安装 Triton。通常，Triton 会随 PyTorch 2.X 一起安装。我们已在 triton==3.3.1 上测试了 Triton 后端。

pip install triton==3.3.1

要启用 nir_exchange，请安装 NIR 和 NIRTorch。

pip install nir nirtorch

以前所未有的简单方式构建 SNN

SpikingJelly 使用起来非常友好。使用 SpikingJelly 构建 SNN 就像在 PyTorch 中构建 ANN 一样简单：

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

这个简单的网络配合泊松编码器，在 MNIST 测试数据集上可以达到 92% 的准确率。更多详细信息请参阅教程。您也可以在 Python 终端中运行以下代码来训练 MNIST 分类任务：

python -m spikingjelly.activation_based.examples.lif_fc_mnist -tau 2.0 -T 100 -device cuda:0 -b 64 -epochs 100 -data-dir <MNIST 数据路径> -amp -opt adam -lr 1e-3 -j 8

快速便捷的 ANN-SNN 转换

SpikingJelly 实现了一个相对通用的 ANN-SNN 转换接口。用户可以通过 PyTorch 实现转换，并且还可以自定义转换模式。

class ANN(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1),
            nn.BatchNorm2d(32, eps=1e-3),
            nn.ReLU(),
            nn.AvgPool2d(2, 2),

            nn.Conv2d(32, 32, 3, 1),
            nn.BatchNorm2d(32, eps=1e-3),
            nn.ReLU(),
            nn.AvgPool2d(2, 2),

            nn.Conv2d(32, 32, 3, 1),
            nn.BatchNorm2d(32, eps=1e-3),
            nn.ReLU(),
            nn.AvgPool2d(2, 2),

            nn.Flatten(),
            nn.Linear(32, 10)
        )

    def forward(self,x):
        x = self.network(x)
        return x

这个简单的网络采用模拟编码，在转换后于 MNIST 测试数据集上可达到 98.44% 的准确率。更多细节请参阅教程。您也可以在 Python 终端中运行以下代码来使用转换后的模型进行 MNIST 分类训练：

>>> import spikingjelly.activation_based.ann2snn.examples.cnn_mnist as cnn_mnist
>>> cnn_mnist.main()

CUDA/Triton 加速的神经元

SpikingJelly 为多步神经元提供了多种后端。您可以使用友好的 torch 后端进行轻松的编码和调试，也可以使用 cupy 或 triton 后端以获得更快的训练速度。

下图比较了多步 LIF 神经元（float32）的 torch 和 cupy 后端的执行时间。通常情况下，triton 后端比 cupy 后端更加高效。

cupy 和 triton 后端还支持 float16 数据类型，可用于 自动混合精度训练。

要使用 cupy 后端，请安装 CuPy。要使用 triton 后端，请安装 Triton。需要注意的是，cupy 和 triton 后端仅支持 GPU，而 torch 后端则同时支持 CPU 和 GPU。

设备支持

• Nvidia GPU
• CPU
• Huawei NPU

使用起来就像使用 PyTorch 一样简单。

>>> net = nn.Sequential(layer.Flatten(), layer.Linear(28 * 28, 10, bias=False), neuron.LIFNode(tau=tau))
>>> net = net.to(device) # 可以是 CPU 或 CUDA 设备

神经形态数据集支持

SpikingJelly 包含以下神经形态数据集：

数据集	来源
ASL-DVS	基于图的神经形态视觉传感目标分类
Bullying10K	Bullying10K：面向隐私保护的欺凌识别的大规模神经形态数据集
CIFAR10-DVS	CIFAR10-DVS：用于目标分类的事件流数据集
DVS-Lip	基于事件的唇读多粒度时空特征感知网络
DVS128 Gesture	一种低功耗、完全基于事件的手势识别系统
ES-ImageNet	ES-ImageNet：用于脉冲神经网络的百万级事件流分类数据集
HARDVS	HARDVS：利用动态视觉传感器重新审视人类活动识别
N-Caltech101	使用扫视将静态图像数据集转换为脉冲神经形态数据集
N-MNIST	使用扫视将静态图像数据集转换为脉冲神经形态数据集
Nav Gesture	利用智能手机计算能力实现背景抑制的基于事件的手势识别
海德堡脉冲数字数据集（SHD）	海德堡脉冲数据集：用于脉冲神经网络系统性评估的数据集
脉冲语音命令数据集（SSC）	海德堡脉冲数据集：用于脉冲神经网络系统性评估的数据集
语音命令	语音命令：用于有限词汇量语音识别的数据集

用户可以使用原始事件数据以及由 SpikingJelly 整合的帧数据：

import torch
from torch.utils.data import DataLoader
from spikingjelly.datasets.utils import pad_sequence_collate, padded_sequence_mask
from spikingjelly.datasets import DVS128Gesture

# 设置数据集的根目录
root_dir = 'D:/datasets/DVS128Gesture'
# 加载事件数据集
event_set = DVS128Gesture(root_dir, train=True, data_type='event')
event, label = event_set[0]
# 打印事件数据中各键及其对应的值
for k in event.keys():
    print(k, event[k])

# t [80048267 80048277 80048278 ... 85092406 85092538 85092700]
# x [49 55 55 ... 60 85 45]
# y [82 92 92 ... 96 86 90]
# p [1 0 0 ... 1 0 0]
# label 0

# 加载固定帧数的数据集
fixed_frames_number_set = DVS128Gesture(root_dir, train=True, data_type='frame', frames_number=20, split_by='number')
# 随机选取两帧并打印其形状
rand_index = torch.randint(low=0, high=fixed_frames_number_set.__len__(), size=[2])
for i in rand_index:
    frame, label = fixed_frames_number_set[i]
    print(f'frame[{i}].shape=[T, C, H, W]={frame.shape}')

# frame[308].shape=[T, C, H, W]=(20, 2, 128, 128)
# frame[453].shape=[T, C, H, W]=(20, 2, 128, 128)

# 加载固定时长的帧数据集，并打印前5个样本的形状
fixed_duration_frame_set = DVS128Gesture(root_dir, data_type='frame', duration=1000000, train=True)
for i in range(5):
    x, y = fixed_duration_frame_set[i]
    print(f'x[{i}].shape=[T, C, H, W]={x.shape}')

# x[0].shape=[T, C, H, W]=(6, 2, 128, 128)
# x[1].shape=[T, C, H, W]=(6, 2, 128, 128)
# x[2].shape=[T, C, H, W]=(5, 2, 128, 128)
# x[3].shape=[T, C, H, W]=(5, 2, 128, 128)
# x[4].shape=[T, C, H, W]=(7, 2, 128, 128)

# 为固定时长的帧数据集创建数据加载器，并打印形状和序列长度
train_data_loader = DataLoader(fixed_duration_frame_set, collate_fn=pad_sequence_collate, batch_size=5)
for x, y, x_len in train_data_loader:
    print(f'x.shape=[N, T, C, H, W]={tuple(x.shape)}')
    print(f'x_len={x_len}')
    mask = padded_sequence_mask(x_len)  # mask.shape = [T, N]
    print(f'mask=\n{mask.t().int()}')
    break

# x.shape=[N, T, C, H, W]=(5, 7, 2, 128, 128)
# x_len=tensor([6, 6, 5, 5, 7])
# mask=
# tensor([[1, 1, 1, 1, 1, 1, 0],
#         [1, 1, 1, 1, 1, 1, 0],
#         [1, 1, 1, 1, 1, 0, 0],
#         [1, 1, 1, 1, 1, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1]], dtype=torch.int32)

未来还将加入更多数据集。

如果部分用户无法访问某些数据集的下载链接，可以从 OpenI 镜像站 下载。

OpenI 镜像站中保存的所有数据集均已获得许可或作者同意。

教程

SpikingJelly 提供了详尽的教程。以下是一些教程：

图片	教程
	基本概念
	神经元
	单层全连接 SNN 用于分类 MNIST
	卷积 SNN 用于分类 FMNIST
	ANN2SNN
	神经形态数据集处理
	分类 DVS 手势
	循环连接与状态突触
	STDP 学习
	强化学习

此处未列出的其他教程也可在 文档 中找到。

ZhenyuZhao 提供了 中文 Jupyter 教程笔记本。

出版物与引用

使用 SpikingJelly 的出版物记录在 出版物 中。如果您在论文中使用了 SpikingJelly，也可以通过拉取请求将其添加到此表格中。

如果您在工作中使用了 SpikingJelly，请按以下方式引用：

@article{
doi:10.1126/sciadv.adi1480,
author = {Wei Fang  and Yanqi Chen  and Jianhao Ding  and Zhaofei Yu  and Timothée Masquelier  and Ding Chen  and Liwei Huang  and Huihui Zhou  and Guoqi Li  and Yonghong Tian },
title = {SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence},
journal = {Science Advances},
volume = {9},
number = {40},
pages = {eadi1480},
year = {2023},
doi = {10.1126/sciadv.adi1480},
URL = {https://www.science.org/doi/abs/10.1126/sciadv.adi1480},
eprint = {https://www.science.org/doi/pdf/10.1126/sciadv.adi1480},
abstract = {Spiking neural networks (SNNs) aim to realize brain-inspired intelligence on neuromorphic chips with high energy efficiency by introducing neural dynamics和spike properties。As the emerging spiking deep learning paradigm attracts increasing interest, traditional programming frameworks cannot meet the demands of automatic differentiation, parallel computation acceleration, and high integration of processing neuromorphic datasets and deployment。In this work, we present the SpikingJelly framework to address the aforementioned dilemma。We contribute a full-stack toolkit for preprocessing neuromorphic datasets, building deep SNNs, optimizing their parameters, and deploying SNNs on neuromorphic chips。Compared to existing methods, the training of deep SNNs can be accelerated 11×, and the superior extensibility and flexibility of SpikingJelly enable users to accelerate custom models at low costs through multilevel inheritance and semiautomatic code generation。SpikingJelly paves the way for synthesizing truly energy-efficient SNN-based machine intelligence systems, which will enrich the ecology of neuromorphic computing。Motivation and introduction of the software framework SpikingJelly for spiking deep learning。}}

贡献

您可以阅读问题列表，了解待解决的问题及最新的开发计划。我们欢迎所有用户参与开发计划的讨论、解决问题并提交拉取请求。

并非所有 API 文档都同时提供英文和中文版本。我们欢迎用户完成翻译工作（从英文到中文或从中文到英文）。

更多信息请参阅 贡献指南。

关于

机构

北京大学数字媒体研究所多媒体学习组 (NELVT) 和 鹏城实验室 是 SpikingJelly 开发的主要机构。

主要开发者

SpikingJelly 历年来由多位主要开发者负责开发和维护。

2024.07~至今

Yifan Huang, Peng Xue

2019.12~2024.06

Wei Fang, Yanqi Chen, Jianhao Ding, Ding Chen, Liwei Huang

感谢所有贡献者

贡献者名单可在 贡献者页面 中找到。

SpikingJelly 快速上手指南

SpikingJelly 是一个基于 PyTorch 的开源脉冲神经网络（SNN）深度学习框架。它提供了极简的 SNN 构建方式、高效的 ANN-SNN 转换工具以及多种硬件加速后端。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.8+
• 核心依赖：
  • PyTorch: 必须预先安装。最新版本要求 torch>=2.2.0 (已在 torch==2.7.1 上测试通过)。
  • 相关库: 建议同时安装 torchvision 和 torchaudio。
• 可选加速依赖（根据需求安装）：
  • CuPy: 用于启用 GPU 加速后端 (cupy-cuda11x 或 cupy-cuda12x)。
  • Triton: 用于启用更高效的 Triton 后端 (通常随 PyTorch 2.x 自动安装，建议版本 triton==3.3.1)。
  • NIR: 如需使用神经形态中间表示交换功能，需安装 nir 和 nirtorch。

提示：国内用户安装 PyTorch 时，推荐使用清华或中科大镜像源以加速下载。

安装步骤

您可以选择安装稳定的发布版或最新的开发版。

1. 安装稳定版（推荐）

从 PyPI 安装最新稳定版本：

pip install spikingjelly

国内加速方案：

pip install spikingjelly -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装开发版（源码安装）

如果您需要体验最新特性（如 Triton 后端优化、新算子等），可以从 GitHub 克隆源码安装：

git clone https://github.com/fangwei123456/spikingjelly.git
cd spikingjelly
pip install .

国内加速方案（使用 Gitee 镜像或代理）： 若直接克隆 GitHub 较慢，可配置 git 代理或使用国内镜像源（如有）。安装依赖时同样建议使用国内 pip 源。

3. 安装可选后端

如需启用高性能后端，请执行以下命令：

启用 CuPy 后端：

pip install cupy-cuda12x  # 适用于 CUDA 12.x
# 或
pip install cupy-cuda11x  # 适用于 CUDA 11.x

启用 Triton 后端：

pip install triton==3.3.1

启用 NIR 支持：

pip install nir nirtorch

基本使用

SpikingJelly 的设计哲学是让构建 SNN 像构建普通 ANN 一样简单。以下是最小化的使用示例。

1. 构建一个简单的 SNN 模型

使用 nn.Sequential 即可快速搭建网络，只需将普通层替换为 SpikingJelly 提供的层和神经元节点。

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

# 定义时间步长
tau = 2.0

# 构建网络：展平层 -> 全连接层 -> LIF 神经元
net = nn.Sequential(
    layer.Flatten(),
    layer.Linear(28 * 28, 10, bias=False),
    neuron.LIFNode(tau=tau, surrogate_function=surrogate.ATan())
)

# 将模型移动到设备 (CPU 或 CUDA)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = net.to(device)

print(net)

2. 运行训练示例

框架内置了丰富的示例脚本。您可以直接在命令行运行以下命令，在 MNIST 数据集上训练一个基于 LIF 神经元的分类网络：

python -m spikingjelly.activation_based.examples.lif_fc_mnist -tau 2.0 -T 100 -device cuda:0 -b 64 -epochs 100 -data-dir <PATH to MNIST> -amp -opt adam -lr 1e-3 -j 8

参数说明：

• -data-dir: 替换为您的 MNIST 数据集存放路径。
• -device: 指定运行设备，如 cuda:0 或 cpu。
• -amp: 启用自动混合精度训练（需 GPU 支持）。

3. ANN 转 SNN (可选)

SpikingJelly 支持将训练好的普通 ANN 快速转换为 SNN。以下是一个简单的调用示例：

import spikingjelly.activation_based.ann2snn.examples.cnn_mnist as cnn_mnist

# 运行内置的 CNN-MNIST 转换示例
cnn_mnist.main()

更多详细教程、API 文档及高级用法（如自定义神经元、多步长优化等），请访问 SpikingJelly 官方文档。