fangwei123456/spikingjelly

SpikingJelly – 用于脉冲神经网络（SNN）的 PyTorch 原生框架

它是什么

• 一个基于 PyTorch 构建的开源 Python 库，允许您定义、训练和部署脉冲神经网络——这是传统人工神经网络的受大脑启发的替代方案。
• 面向希望处理基于事件（神经形态）数据、将现有 ANN 模型转换为 SNN 或运行大规模 SNN 实验的研究人员和工程师。

关键能力

领域	SpikingJelly 提供的内容
SNN 建模	现成的脉冲神经元模块（LIF、IF 等）、代理梯度函数、有状态/无状态层以及参考 SNN 架构集合。
训练工作流	PyTorch 风格的训练循环、在线学习工具以及将训练好的 ANN 转换为 SNN（ANN2SNN）的工具。
性能	自动选择最快的执行后端（纯 PyTorch、Triton 内核或可选的原生 CUDA 扩展）。支持混合精度（例如 fp8）以及用于自定义神经元内核的“FlexSN”内核生成器。
扩展	内存高效的脉冲压缩、实验性多 GPU 分布式训练以及用于大型模型的精度策略辅助工具。
数据集	内置加载器，支持多种神经形态事件数据集（例如 N-MNIST、CIFAR10-DVS、DVS-Gesture、Speech Commands 等），并提供原始事件和基于帧的访问。
分析	用于 FLOPs、SynOps、内存流量和推理能耗估算的工具。
互操作与部署	用于神经形态硬件和其他框架（NIR、Lava、Lynxi）的导出辅助工具。

快速入门示例

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())
)

该模型的定义方式与任何 PyTorch nn.Module 完全相同；然后您可以使用提供的教程在事件数据集上训练它，或将传统 ANN 转换为 SNN。

安装

• 需要 Python ≥ 3.11 和 PyTorch ≥ 2.6。
• 通过 uv 包管理器（或 pip）从 PyPI 安装：

  uv pip install spikingjelly               # 稳定版本
  uv pip install "spikingjelly[triton]"      # 包含可选的 Triton 内核
  

• 为了获得最大性能，您可以构建可选的原生 CUDA 扩展：

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

• 详细的安装指南和故障排除在文档中。

典型用例

• 神经形态视觉或音频研究（例如 DVS 相机、基于事件的语音）。
• 将高精度 ANN 转换为低功耗 SNN 用于边缘设备。
• 在多 GPU 集群上对脉冲 Transformer 组件或大规模 SNN 训练进行基准测试。
• 通过 NIR/Lava 将模型导出到神经形态芯片或仿真框架。

项目健康

• 积极维护（维护者已列出，定期发布，版本遵循 SemVer）。
• 文档托管在 Read the Docs 上，提供安装、神经元基础、数据集处理和部署教程。
• Apache-2.0 许可证——对学术和商业使用均宽松。

了解更多

• 完整文档：https://spikingjelly.readthedocs.io/
• 快速入门教程、数据集指南和转换笔记本可从 README 链接。
• 欢迎通过 GitHub 问题跟踪器和拉取请求工作流贡献社区。

---

SpikingJelly 本质上是“用于脉冲神经网络的 PyTorch”，提供熟悉的 API，同时处理事件驱动计算的独特挑战。