安装指南#
English version: Installation Guide
准备环境#
V2 要求 Python >= 3.11、Torch >= 2.6。先创建环境,再按 PyTorch 官方安装入口 为目标设备安装 Torch、torchvision 和 torchaudio:
uv venv --python 3.11
source .venv/bin/activate
本文使用 uv 管理环境和安装。Torch、torchvision、torchaudio 应按官方命令配套
安装,避免安装器为缺失的配套包重新选择 Torch;安装 SJ 后可运行
uv pip check 检查依赖一致性。
本文按当前开发源码说明。PyPI 安装只包含已发布的改动;使用旧版本时,请切换到 对应版本的文档。最低版本要求不代表所有版本和设备都已经验收;已有验证包括 Torch 2.7.1,以及 Torch 2.11.0+cu128 / Triton 3.6.0 的 GPU 环境。
V2 使用兼容 PEP 440 的语义化版本号。此前的 0.0.0.0.X 是历史版本方案,
奇数 X 为开发版,偶数为 PyPI 稳定版。
选择安装方式#
CPU 不需要 Triton 或原生 CUDA。NVIDIA CUDA 用户可按需要安装 Triton、手动 构建原生扩展,或同时准备两者。Triton 保持可选,不会由普通安装默认添加。 普通 GPU 神经元可先使用匹配的 Triton;融合投影,或实测 eager 原生实现有收益时, 再考虑编译原生扩展。
已有与 Torch 匹配、满足 SJ 最低版本要求的 Triton 时,无需额外安装或升级。
可用 uv run --no-sync python -c "import triton; print(triton.__version__)"
查看版本;导入成功不能证明所有 kernel 或编译组合都受支持。缺失时再安装 extra。
两种加速实现可以同时安装,模块无需配置 backend。#
安装方式 |
命令 |
|---|---|
PyPI 发布版 |
|
PyPI 先行版 |
|
可选 Triton |
|
最新源码开发版 |
|
若开发版需要 Triton,可执行:
uv pip install "spikingjelly[triton] @ git+https://github.com/fangwei123456/spikingjelly.git"
源码也可从 OpenI 获取。
已有源码 checkout 的开发者可执行 uv pip install --editable ".[triton]"。
根目录 ops/ 经安装映射为 spikingjelly._ops,只设置 PYTHONPATH
不能代替安装。开发环境约定见仓库 CONTRIBUTING.md。
手动构建原生 CUDA#
常规 wheel 不含预编译原生动态库。当前推荐手动源码构建,默认从 PyPI 下载 sdist 并在本机编译,无需克隆仓库。以下命令需等待包含这些算子的 V2 版本发布到 PyPI,旧发行版不能据此构建当前算子。
先准备 CUDA 版 Torch、与其匹配的 CUDA Toolkit(含 nvcc)和 C++ 编译器。
若工具链不在默认位置,设置 CUDA_HOME;无可见 GPU 的构建需设置
TORCH_CUDA_ARCH_LIST,指定实际目标架构。
uv pip install "setuptools>=77.0.3" ninja
SJ_BUILD_NATIVE_CUDA=1 uv pip install \
--no-build-isolation --no-binary spikingjelly \
--reinstall-package spikingjelly --no-cache "spikingjelly>=2.0.0"
构建成功后,在可用的目标 NVIDIA GPU 上启动新进程,检查默认 eager 绑定:
uv run --no-sync python -c \
'import torch; from spikingjelly.activation_based import functional; print(functional.neuron_implementation("lif", torch.device("cuda:0")))'
兼容的原生扩展预期返回 implementation: cuda;若为 triton 或 torch,
检查 unavailable 中的原生加载原因。查询不运行神经元;状态精度等配置仍会影响
实际调用路径。完整检查见下文。
选项 |
作用 |
|---|---|
|
本次源码构建要求编译原生扩展;默认不编译。 |
|
下载 SpikingJelly sdist;其他依赖仍可使用 wheel。 |
|
使用当前环境的 Torch 和构建依赖。默认隔离环境不包含我们的构建所需 Torch。 |
|
即使已安装,也重新安装 SpikingJelly。 |
|
不复用先前构建的 wheel,避免构建设置变化后仍使用旧产物。 |
该命令不会默认添加 Triton;需要时另行安装 spikingjelly[triton]。
显式要求构建时,缺 CUDA 版 Torch、CUDA Toolkit、编译器或必要架构配置会立即报错;
实际编译失败也会终止安装。未设置环境变量时仍构建纯 Python 包。
该检查只在源码构建中生效;安装已有 wheel 或跳过已有安装不会执行构建检查。
命令保留重安装及禁用缓存参数,因为 --no-binary 仍可能复用 uv 缓存的 wheel。
PyPI 下载不可用或需要开发源码时,可从 OpenI 获取 checkout,沿用同一手动构建:
git clone https://git.openi.org.cn/OpenI/spikingjelly.git
cd spikingjelly
uv pip install "setuptools>=77.0.3" ninja
SJ_BUILD_NATIVE_CUDA=1 uv pip install --no-build-isolation .
也可使用 GitHub checkout。uv 对命令行明确传入的本地目录重新构建安装,因此无需
sdist 的 --no-binary 等参数。editable 修改 .cu 或头文件后需要再次执行
原生构建;不会像 Python 源码那样即时生效。
原生扩展运行时只加载二进制,不调用编译器。当前加载器检查算子 ABI、完整 Torch 版本、CUDA 版本及目标 GPU 支持;更换环境后可能需要重建。Triton 保留自己的 首次 JIT 和缓存,不能假定所有平台的 CUDA Torch 都附带可用 Triton。
安装后如何执行#
将模块和输入移到同一设备即可。普通注册神经元的选择流程如下;CUDA 首次使用 按设备及执行路径检查候选,随后复用绑定,不在线测速。
图中 CUDA 箭头表示候选顺序。CUDA Graph 沿用捕获函数已经选择的实现。#
安装某种实现,不代表每个调用都使用它:
功能 |
启用条件 |
|---|---|
普通注册神经元 |
受支持的输入、FP32 状态及内置替代梯度可进入融合路径。普通模块状态跟随输入 dtype,低精度状态或自定义 surrogate 可使用 Torch 参考公式。 |
IF/LIF/PLIF 显式精度配置 |
要求支持该组合的 CUDA Triton;不能以其他实现替换明确指定的数值策略。 |
FlexSN CUDA 多步 |
可编译的 core 使用 Triton;已知不支持的组合可使用 Torch/HOP。单步使用 Torch。 |
融合 IF/LIF-Linear、packed/sparse 投影 |
原生扩展兼容时使用对应内核,缺失时使用 Torch 参考执行;Triton 不提供这些原生融合投影的性能特性。 |
确认安装与排查#
先检查安装位置和 CUDA Torch:
import torch
import spikingjelly
print(spikingjelly.__file__)
print(torch.__version__, torch.version.cuda, torch.cuda.is_available())
有可用 NVIDIA GPU 时,再查询普通神经元的绑定:
from spikingjelly.activation_based import functional
device = torch.device("cuda:0")
print(functional.neuron_implementation("lif", device, execution="eager"))
print(functional.neuron_implementation("lif", device, execution="compile"))
查询初始化选择但不计算神经元输出、不推进状态。返回的 implementation 表示
绑定,unavailable 解释更高优先级候选不可用的原因;低精度状态等配置仍可能
走参考路径。缺失依赖或已知不兼容可以检查下一候选,未知 JIT、kernel、OOM 或
梯度错误会直接报告。日志默认静默,启用方法与严格诊断环境变量见
神经元自动执行。改变依赖、扩展或配置后重启进程。
其他可选依赖#
功能 |
安装命令 |
|---|---|
|
|
Lightning 集成 |
|
Transformer Engine 精度功能 |
|
其他 extras 及版本约束见仓库 pyproject.toml。当前包不依赖 CuPy;旧安装项
迁移见 从老版本迁移。