Core Neuron Modules#

SpikingJelly 的 核心神经元模块 提供了规范神经元抽象。 这些神经元是 SNN 领域中被广泛接受和使用的模型，旨在作为研究与实际应用中的基础构建单元。

纳入该类别的主要标准包括：

  • 神经元在概念上具有通用性，不依赖于某一特定论文、任务或训练策略。

  • 该神经元可以被推荐用于下游项目中的通用建模需求。

除非明确需要某种特定的研究型行为，否则在构建新模型时，建议用户优先使用核心神经元模块。

SpikingJelly's core neuron modules provide canonical and stable neuron abstractions. These neurons represent widely accepted models in SNN (SNN) literature and are designed to serve as fundamental building blocks for both research and practical applications.

The main criteria for inclusion in this category are:

  • The neuron is conceptually general and not tied to a specific paper, task, or training strategy.

  • The neuron can be recommended for general use in downstream projects.

Users are encouraged to preferentially use core neuron modules when building new models, unless a specific research-oriented behavior is explicitly required.

Base Classes#

class spikingjelly.activation_based.neuron.base_node.BaseNode(v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s', backend='torch', store_v_seq: bool = False)[源代码]#

基类：MemoryModule

API Language: 中文 | English

  • 中文

可微分SNN神经元的基类神经元。

参数:

  • v_threshold (float) -- 神经元的阈值电压

  • v_reset (Optional[float]) -- 神经元的重置电压。如果不为 None，当神经元释放脉冲后，电压会被重置为 v_reset； 如果设置为 None，当神经元释放脉冲后，电压会被减去 v_threshold

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • detach_reset (bool) -- 是否将reset过程的计算图分离

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • backend (str) -- 使用哪种后端。不同的 step_mode 可能会带有不同的后端。可以通过打印 self.supported_backends 查看当前 使用的步进模式支持的后端。在支持的情况下，使用 'cupy' 或 'triton' 后端速度更快。

  • store_v_seq (bool) -- 在使用 step_mode = 'm' 时，给与 shape = [T, N, *] 的输入后，是否保存中间过程的 shape = [T, N, *] 的各个时间步的电压值 self.v_seq 。设置为 False 时计算完成后只保留最后一个时刻的电压，即 shape = [N, *] 的 self.v 。 通常设置成 False ，可以节省内存

  • English

This class is the base class of differentiable spiking neurons.

参数:

  • v_threshold (float) -- threshold of this neurons layer

  • v_reset (Optional[float]) -- reset voltage of this neurons layer. If not None, the neuron's voltage will be set to v_reset after firing a spike. If None, the neuron's voltage will subtract v_threshold after firing a spike

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • detach_reset (bool) -- whether detach the computation graph of reset in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

  • backend (str) -- backend fot this neurons layer. Different step_mode may support for different backends. The user can print self.supported_backends and check what backends are supported by the current step_mode. If supported, using 'cupy' or 'triton' backend will be faster

  • store_v_seq (bool) -- when using step_mode = 'm' and given input with shape = [T, N, *], this option controls whether storing the voltage at each time-step to self.v_seq with shape = [T, N, *]. If set to False, only the voltage at last time-step will be stored to self.v with shape = [N, *], which can reduce the memory consumption

返回:

None

返回类型:

None

abstractmethod neuronal_charge(x: Tensor)[源代码]#

API Language: 中文 | English

  • 中文

定义神经元的充电差分方程。子类必须实现这个函数。

  • English

Define the charge difference equation. The sub-class must implement this function.

neuronal_fire()[源代码]#

API Language: 中文 | English

  • 中文

根据当前神经元的电压、阈值，计算输出脉冲。

  • English

Calculate out spikes of neurons by their current membrane potential and threshold voltage.

neuronal_reset(spike)[源代码]#

API Language: 中文 | English

  • 中文

根据当前神经元释放的脉冲，对膜电位进行重置。

  • English

Reset the membrane potential according to neurons' output spikes.

single_step_forward(x: Tensor)[源代码]#

API Language: 中文 | English

  • 中文

按照充电、放电、重置的顺序进行前向传播。

参数:

x (torch.Tensor) -- 输入到神经元的电压增量

返回:

神经元的输出脉冲

返回类型:

torch.Tensor

  • English

Forward by the order of neuronal_charge, neuronal_fire, and neuronal_reset.

参数:

x (torch.Tensor) -- increment of voltage inputted to neurons

返回:

out spikes of neurons

返回类型:

torch.Tensor

multi_step_forward(x_seq: Tensor)[源代码]#
class spikingjelly.activation_based.neuron.base_node.NonSpikingBaseNode(decode: str | None = None)[源代码]#

基类：Module, MultiStepModule

参数:

decode (Optional[str]) -- 解码方式。若不为 None，在 forward 中将使用该方式对膜电位序列进行解码

返回:

None

返回类型:

None

abstractmethod neuronal_charge(x: Tensor)[源代码]#
forward(x_seq: Tensor)[源代码]#
class spikingjelly.activation_based.neuron.base_node.SimpleBaseNode(v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s')[源代码]#

基类：MemoryModule

API Language: 中文 | English

  • 中文

BaseNode 的简化版，便于用户修改或扩展神经元。

参数:

  • v_threshold (float) -- 神经元的阈值电压

  • v_reset (Optional[float]) -- 神经元的重置电压

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • detach_reset (bool) -- 是否将 reset 过程的计算图分离

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • English

A simple version of BaseNode. Users can modify this neuron easily.

参数:

  • v_threshold (float) -- threshold of this neurons layer

  • v_reset (Optional[float]) -- reset voltage of this neurons layer

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • detach_reset (bool) -- whether detach the computation graph of reset in backward

  • step_mode (str) -- the step mode, which can be 's' (single-step) or 'm' (multi-step)

返回:

None

返回类型:

None

single_step_forward(x: Tensor)[源代码]#
neuronal_charge(x: Tensor)[源代码]#
neuronal_fire()[源代码]#
neuronal_reset(spike)[源代码]#

Integrate-and-fire (IF) Neurons#

class spikingjelly.activation_based.neuron.integrate_and_fire.SimpleIFNode(v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s')[源代码]#

基类：SimpleBaseNode

API Language: 中文 | English

  • 中文

IFNode 的简化版实现。

参数:

  • v_threshold (float) -- 神经元阈值电压

  • v_reset (Optional[float]) -- 神经元重置电压

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 替代梯度函数

  • detach_reset (bool) -- 是否在反向传播时分离 reset 计算图

  • step_mode (str) -- 步进模式，可为 "s" 或 "m"

  • English

A simple version of IFNode.

参数:

  • v_threshold (float) -- Threshold voltage of the neuron

  • v_reset (Optional[float]) -- Reset voltage of the neuron

  • surrogate_function (surrogate.SurrogateFunctionBase) -- Surrogate gradient function

  • detach_reset (bool) -- Whether to detach reset graph in backward

  • step_mode (str) -- Step mode, either "s" or "m"

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#

API Language: 中文 | English

  • 中文

  • 中文

神经元充电的微分方程：

\[H[t] = V[t-1] + X[t]\]
参数:

x (torch.Tensor) -- 输入电压

返回:

None（膜电位更新存储在 self.v 中）

返回类型:

None

  • English

  • English

The differential equation for neuronal charge:

\[H[t] = V[t-1] + X[t]\]
参数:

x (torch.Tensor) -- Input voltage

返回:

None (membrane potential is stored in self.v)

返回类型:

None

class spikingjelly.activation_based.neuron.integrate_and_fire.IFNode(v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s', backend='torch', store_v_seq: bool = False)[源代码]#

基类：BaseNode

API Language: 中文 | English

  • 中文

Integrate-and-Fire 神经元模型，可以看作理想积分器，无输入时电压保持恒定，不会像 LIF 神经元那样衰减。其阈下神经动力学方程为：

\[H[t] = V[t-1] + X[t]\]
参数:

  • v_threshold (float) -- 神经元的阈值电压

  • v_reset (Optional[float]) -- 神经元的重置电压。如果不为 None，当神经元释放脉冲后，电压会被重置为 v_reset； 如果设置为 None，当神经元释放脉冲后，电压会被减去 v_threshold

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • detach_reset (bool) -- 是否将 reset 过程的计算图分离

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • backend (str) -- 使用哪种后端。不同的 step_mode 可能会带有不同的后端。可以通过打印 self.supported_backends 查看当前 使用的步进模式支持的后端。在支持的情况下，使用 'cupy' 或 'triton' 后端速度更快。

  • store_v_seq (bool) -- 在使用 step_mode = 'm' 时，给与 shape = [T, N, *] 的输入后，是否保存中间过程的 shape = [T, N, *] 的各个时间步的电压值 self.v_seq 。设置为 False 时计算完成后只保留最后一个时刻的电压，即 shape = [N, *] 的 self.v 。 通常设置成 False ，可以节省内存

  • English

The Integrate-and-Fire neuron, which can be seen as an ideal integrator. The voltage of the IF neuron will not decay as that of the LIF neuron. The sub-threshold neural dynamics of it is as followed:

\[H[t] = V[t-1] + X[t]\]
参数:

  • v_threshold (float) -- threshold of this neurons layer

  • v_reset (Optional[float]) -- reset voltage of this neurons layer. If not None, the neuron's voltage will be set to v_reset after firing a spike. If None, the neuron's voltage will subtract v_threshold after firing a spike

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • detach_reset (bool) -- whether detach the computation graph of reset in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

  • backend (str) -- backend fot this neurons layer. Different step_mode may support for different backends. The user can print self.supported_backends and check what backends are supported by the current step_mode. If supported, using 'cupy' or 'triton' backend will be faster

  • store_v_seq (bool) -- when using step_mode = 'm' and given input with shape = [T, N, *], this option controls whether storing the voltage at each time-step to self.v_seq with shape = [T, N, *]. If set to False, only the voltage at last time-step will be stored to self.v with shape = [N, *], which can reduce the memory consumption

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#
multi_step_forward(x_seq: Tensor)[源代码]#
single_step_forward(x: Tensor)[源代码]#
class spikingjelly.activation_based.neuron.integrate_and_fire.NonSpikingIFNode(decode: str | None = None)[源代码]#

基类：NonSpikingBaseNode

API Language: 中文 | English

  • 中文

  • 中文

不发放脉冲的 IF 节点，输出膜电位（或根据 decode 进行解码）。

参数:

decode (Optional[str]) -- 非脉冲输出解码方式，见 NonSpikingBaseNode

  • English

  • English

Non-spiking IF node that outputs membrane potential (or decoded outputs specified by decode).

参数:

decode (Optional[str]) -- Decoding mode for non-spiking outputs, see NonSpikingBaseNode

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#

Leaky Integrate-and-fire (LIF) Neurons#

class spikingjelly.activation_based.neuron.lif.SimpleLIFNode(tau: float, decay_input: bool, v_threshold: float = 1.0, v_reset: float = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s')[源代码]#

基类：SimpleBaseNode

API Language: 中文 | English

  • 中文

LIFNode 的简化版实现。

  • English

A simple version of LIFNode.

参数:

  • tau (float) -- 膜电位时间常数（详见父类 LIFNode）

  • decay_input (bool) -- 输入是否参与衰减（详见父类）

  • v_threshold (float) -- 神经元的阈值电压（详见父类）

  • v_reset (float) -- 神经元的重置电压（详见父类）

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 替代梯度函数（详见父类）

  • detach_reset (bool) -- 是否将 reset 过程的计算图分离

  • step_mode (str) -- 步进模式，可为 "s" 或 "m"

  • tau -- Membrane time constant (see parent class LIFNode)

  • decay_input -- Whether input participates in decay (see parent)

  • v_threshold -- Threshold voltage of the neuron (see parent)

  • v_reset -- Reset voltage of the neuron (see parent)

  • surrogate_function -- Surrogate gradient function (see parent)

  • detach_reset -- Whether to detach reset graph in backward

  • step_mode -- Step mode, either "s" or "m"

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#

If decay_input == True:

\[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]

If decay_input == False:

\[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]
class spikingjelly.activation_based.neuron.lif.LIFNode(tau: float = 2.0, decay_input: bool = True, v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s', backend='torch', store_v_seq: bool = False)[源代码]#

基类：BaseNode

API Language: 中文 | English

  • 中文

Leaky Integrate-and-Fire 神经元模型，可以看作是带漏电的积分器。其阈下神经动力学方程为：

若 decay_input == True:

\[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]

若 decay_input == False:

\[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]
参数:

  • tau (float) -- 膜电位时间常数

  • decay_input (bool) -- 输入是否也会参与衰减

  • v_threshold (float) -- 神经元的阈值电压

  • v_reset (Optional[float]) -- 神经元的重置电压。如果不为 None，当神经元释放脉冲后，电压会被重置为 v_reset； 如果设置为 None，当神经元释放脉冲后，电压会被减去 v_threshold

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • detach_reset (bool) -- 是否将 reset 过程的计算图分离

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • backend (str) -- 使用哪种后端。不同的 step_mode 可能会带有不同的后端。可以通过打印 self.supported_backends 查看当前 使用的步进模式支持的后端。在支持的情况下，使用 'cupy' 或 'triton' 后端速度更快。

  • store_v_seq (bool) -- 在使用 step_mode = 'm' 时，给与 shape = [T, N, *] 的输入后，是否保存中间过程的 shape = [T, N, *] 的各个时间步的电压值 self.v_seq 。设置为 False 时计算完成后只保留最后一个时刻的电压，即 shape = [N, *] 的 self.v 。 通常设置成 False ，可以节省内存

  • English

The Leaky Integrate-and-Fire neuron, which can be seen as a leaky integrator. The subthreshold neural dynamics of it is as followed:

If decay_input == True:

\[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]

If decay_input == False:

\[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]
参数:

  • tau (float) -- membrane time constant

  • decay_input (bool) -- whether the input will decay

  • v_threshold (float) -- threshold of this neurons layer

  • v_reset (Optional[float]) -- reset voltage of this neurons layer. If not None, the neuron's voltage will be set to v_reset after firing a spike. If None, the neuron's voltage will subtract v_threshold after firing a spike

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • detach_reset (bool) -- whether detach the computation graph of reset in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

  • backend (str) -- backend fot this neurons layer. Different step_mode may support for different backends. The user can print self.supported_backends and check what backends are supported by the current step_mode. If supported, using 'cupy' or 'triton' backend will be faster

  • store_v_seq (bool) -- when using step_mode = 'm' and given input with shape = [T, N, *], this option controls whether storing the voltage at each time-step to self.v_seq with shape = [T, N, *]. If set to False, only the voltage at last time-step will be stored to self.v with shape = [N, *], which can reduce the memory consumption

返回:

None

返回类型:

None

extra_repr()[源代码]#
neuronal_charge(x: Tensor)[源代码]#
single_step_forward(x: Tensor)[源代码]#
multi_step_forward(x_seq: Tensor)[源代码]#
class spikingjelly.activation_based.neuron.lif.NonSpikingLIFNode(tau: float = 2.0, decode: str | None = None)[源代码]#

基类：NonSpikingBaseNode

Non-spiking version of LIFNode that outputs continuous-valued membrane potentials instead of spikes.

See also: spikingjelly.activation_based.layer.misc.SynapseFilter.

参数:

  • tau (float) -- 膜电位时间常数

  • decode (Optional[str]) -- 解码方式

  • tau -- Membrane time constant

  • decode -- Decoding method

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#

Parametric Leaky Integrate-and-fire (PLIF) Neurons#

class spikingjelly.activation_based.neuron.plif.ParametricLIFNode(init_tau: float = 2.0, decay_input: bool = True, v_threshold: float = 1.0, v_reset: float | None = 0.0, surrogate_function: SurrogateFunctionBase = Sigmoid(alpha=4.0, spiking=True), detach_reset: bool = False, step_mode='s', backend='torch', store_v_seq: bool = False)[源代码]#

基类：BaseNode

API Language: 中文 | English

  • 中文

Parametric Leaky Integrate-and-Fire (PLIF) 神经元模型，提出自 Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural Networks。可以看作是带漏电的积分器。其阈下神经动力学方程为：

若 decay_input == True:

\[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]

若 decay_input == False:

\[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]

其中 \(\frac{1}{\tau} = {\rm Sigmoid}(w)\)，\(w\) 是可学习的参数。

参数:

  • init_tau (float) -- 膜电位时间常数的初始值

  • decay_input (bool) -- 输入是否也会参与衰减

  • v_threshold (float) -- 神经元的阈值电压

  • v_reset (Optional[float]) -- 神经元的重置电压。如果不为 None，当神经元释放脉冲后，电压会被重置为 v_reset； 如果设置为 None，当神经元释放脉冲后，电压会被减去 v_threshold

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • detach_reset (bool) -- 是否将 reset 过程的计算图分离

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • backend (str) -- 使用哪种后端。不同的 step_mode 可能会带有不同的后端。可以通过打印 self.supported_backends 查看当前 使用的步进模式支持的后端。在支持的情况下，使用 'cupy' 或 'triton' 后端速度更快。

  • store_v_seq (bool) -- 在使用 step_mode = 'm' 时，给与 shape = [T, N, *] 的输入后，是否保存中间过程的 shape = [T, N, *] 的各个时间步的电压值 self.v_seq 。设置为 False 时计算完成后只保留最后一个时刻的电压，即 shape = [N, *] 的 self.v 。 通常设置成 False ，可以节省内存

  • English

The Parametric Leaky Integrate-and-Fire (PLIF) neuron, proposed in Incorporating Learnable Membrane Time Constant to Enhance Learning of Spiking Neural Networks, can be seen as a leaky integrator. The subthreshold neural dynamics of it is as followed:

IF decay_input == True:

\[H[t] = V[t-1] + \frac{1}{\tau}(X[t] - (V[t-1] - V_{reset}))\]

IF decay_input == False:

\[H[t] = V[t-1] - \frac{1}{\tau}(V[t-1] - V_{reset}) + X[t]\]

where \(\frac{1}{\tau} = {\rm Sigmoid}(w)\), \(w\) is a learnable parameter.

参数:

  • init_tau (float) -- the initial value of membrane time constant

  • decay_input (bool) -- whether the input will decay

  • v_threshold (float) -- threshold of this neurons layer

  • v_reset (Optional[float]) -- reset voltage of this neurons layer. If not None, the neuron's voltage will be set to v_reset after firing a spike. If None, the neuron's voltage will subtract v_threshold after firing a spike

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • detach_reset (bool) -- whether detach the computation graph of reset in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

  • backend (str) -- backend for this neurons layer. Different step_mode may support for different backends. The user can print self.supported_backends and check what backends are supported by the current step_mode. If supported, using 'cupy' or 'triton' backend will have the fastest training speed

  • store_v_seq (bool) -- when using step_mode = 'm' and given input with shape = [T, N, *], this option controls whether storing the voltage at each time-step to self.v_seq with shape = [T, N, *]. If set to False, only the voltage at last time-step will be stored to self.v with shape = [N, *], which can reduce the memory consumption

返回:

None

返回类型:

None

neuronal_charge(x: Tensor)[源代码]#
multi_step_forward(x_seq: Tensor)[源代码]#

Parallel Spiking Neuron Family#

class spikingjelly.activation_based.neuron.psn.PSN(T: int, surrogate_function: SurrogateFunctionBase = ATan(alpha=2.0, spiking=True))[源代码]#

基类：Module, MultiStepModule

API Language: 中文 | English

  • 中文

并行脉冲神经元（Parallel Spiking Neuron，PSN），由 Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability 提出。神经元动力学定义如下：

\[\begin{split}H &= WX, ~~~~~~~~~~~~~~~W \in \mathbb{R}^{T \times T}, X \in \mathbb{R}^{T \times N}\\ S &= \Theta(H - B), ~~~~~B \in \mathbb{R}^{T}, S\in \{0, 1\}^{T \times N}\end{split}\]

其中 \(W\) 是可学习的权重矩阵，\(B\) 是可学习的阈值。

注意

PSN 仅支持多步模式。

参数:

  • English

The Parallel Spiking Neuron (PSN), proposed in Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability. The neuronal dynamics are defined as:

\[\begin{split}H &= WX, ~~~~~~~~~~~~~~~W \in \mathbb{R}^{T \times T}, X \in \mathbb{R}^{T \times N}\\ S &= \Theta(H - B), ~~~~~B \in \mathbb{R}^{T}, S\in \{0, 1\}^{T \times N}\end{split}\]

where \(W\) is the learnable weight matrix, and \(B\) is the learnable threshold.

Note

The PSN only supports the multi-step mode.

参数:

  • T (int) -- the number of time-steps

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

返回:

None

返回类型:

None

forward(x_seq: Tensor)[源代码]#
class spikingjelly.activation_based.neuron.psn.MaskedPSN(k: int, T: int, lambda_init: float = 0.0, surrogate_function: SurrogateFunctionBase = ATan(alpha=2.0, spiking=True), step_mode: str = 's')[源代码]#

基类：MemoryModule

API Language: 中文 | English

  • 中文

Masked Parallel Spiking Neuron，由 Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability 提出。神经元动力学定义如下：

\[\begin{split}H &= (W \cdot {M}_{k})X, ~~~~~~~~~~~~~~~W \in \mathbb{R}^{T \times T}, {M}_{k} \in \mathbb{R}^{T \times T}, X \in \mathbb{R}^{T \times N} \\ S &= \Theta(H - B), ~~~~~B \in \mathbb{R}^{T}, S\in \{0, 1\}^{T \times N}\end{split}\]

其中 \(W\) 是可学习权重矩阵，\(B\) 是可学习阈值，\({M}_{k}\) 定义为：

\[\begin{split}{M}_{k}[i][j] = \begin{cases} 1, ~~ j \leq i \leq j + k - 1 \\ 0, \mathrm{otherwise} \end{cases}.\end{split}\]

\(\lambda\) 用于调节逐步掩码过程：

\[M_{k}(\lambda) = \lambda \cdot M_{k} + (1 - \lambda) \cdot J,\]

其中 \(J\) 为全 1 矩阵。用户可以在训练中通过 self.lambda_ = ... 设置 \(\lambda\)。

注意

Masked PSN 支持单步模式和多步模式，但多步模式比单步模式快得多。

参数:

  • k (int) -- Masked PSN 的阶数

  • T (int) -- 时间步数

  • lambda_init (float) -- \(\lambda\) 的初始值，用于调节逐步掩码过程

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • English

Masked Parallel Spiking Neuron (Masked PSN), proposed in Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability. The neuronal dynamics are defined as:

\[\begin{split}H &= (W \cdot {M}_{k})X, ~~~~~~~~~~~~~~~W \in \mathbb{R}^{T \times T}, {M}_{k} \in \mathbb{R}^{T \times T}, X \in \mathbb{R}^{T \times N} \\ S &= \Theta(H - B), ~~~~~B \in \mathbb{R}^{T}, S\in \{0, 1\}^{T \times N}\end{split}\]

where \(W\) is the learnable weight matrix, \(B\) is the learnable threshold, and \({M}_{k}\) is defined as:

\[\begin{split}{M}_{k}[i][j] = \begin{cases} 1, ~~ j \leq i \leq j + k - 1 \\ 0, \mathrm{otherwise} \end{cases}.\end{split}\]

\(\lambda\) is used to adjust the progressive masking process:

\[M_{k}(\lambda) = \lambda \cdot M_{k} + (1 - \lambda) \cdot J,\]

where \(J\) is an all-one matrix. Users can set \(\lambda\) during training by calling self.lambda_ = ....

Note

The masked PSN supports both single-step and multi-step mode. Multi-step mode is much faster than single-step mode.

参数:

  • k (int) -- the order of the Masked PSN

  • T (int) -- the number of time-steps

  • lambda_init (float) -- the initial value of \(\lambda\) to adjust the progressive masking process

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

返回:

None

返回类型:

None

static gen_masked_weight(lambda_: Tensor, mask0: Tensor, mask1: Tensor, weight: Tensor)[源代码]#
masked_weight()[源代码]#
single_step_forward(x: Tensor)[源代码]#
multi_step_forward(x_seq: Tensor)[源代码]#
property lambda_#
class spikingjelly.activation_based.neuron.psn.SlidingPSN(k: int, exp_init: bool = True, surrogate_function: SurrogateFunctionBase = ATan(alpha=2.0, spiking=True), step_mode: str = 's', backend: str = 'gemm')[源代码]#

基类：MemoryModule

API Language: 中文 | English

  • 中文

Sliding Parallel Spiking Neuron，由 Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability 提出。神经元动力学定义如下：

\[\begin{split}H[t] &= \sum_{i=0}^{k-1} W_i \cdot X[t - k + 1 + i], \\ S[t] &= \Theta(H[t] - B),\end{split}\]

其中 \(W = [W_0, W_1, ..., W_{k-1}] \in \mathbb{R}^{T}\) 是可学习权重，\(B\) 是可学习阈值。

注意

Sliding PSN 支持单步模式和多步模式，但多步模式比单步模式快得多。

参数:

  • k (int) -- Sliding PSN 的阶数

  • exp_init (bool) -- 如果为 True，权重初始化为 (..., 1/4, 1/2, 1)；如果为 False，权重使用 Kaiming uniform 初始化

  • surrogate_function (surrogate.SurrogateFunctionBase) -- 反向传播时用来计算脉冲函数梯度的替代函数

  • step_mode (str) -- 步进模式，可以为 's' (单步) 或 'm' (多步)

  • backend (str) -- 神经元层使用的后端，可以为 "gemm" 或 "conv"。此选项仅在多步模式下生效

  • English

Sliding Parallel Spiking Neuron (Sliding PSN), proposed in Parallel Spiking Neurons with High Efficiency and Long-term Dependencies Learning Ability. The neuronal dynamics are defined as:

\[\begin{split}H[t] &= \sum_{i=0}^{k-1} W_i \cdot X[t - k + 1 + i], \\ S[t] &= \Theta(H[t] - B),\end{split}\]

where \(W = [W_0, W_1, ..., W_{k-1}] \in \mathbb{R}^{T}\) is the learnable weight, and \(B\) is the learnable threshold.

Note

Sliding PSN supports both single-step and multi-step mode. Multi-step mode is much faster than single-step mode.

参数:

  • k (int) -- the order of the Sliding PSN

  • exp_init (bool) -- if True, the weight will be initialized as (..., 1/4, 1/2, 1); if False, the weight will be initialized by Kaiming uniform

  • surrogate_function (surrogate.SurrogateFunctionBase) -- the function for calculating surrogate gradients of the heaviside step function in backward

  • step_mode (str) -- the step mode, which can be s (single-step) or m (multi-step)

  • backend (str) -- backend for this neuron layer, which can be "gemm" or "conv". This option only works for multi-step mode

返回:

None

返回类型:

None

gen_gemm_weight(T: int)[源代码]#
single_step_forward(x: Tensor)[源代码]#
multi_step_forward(x_seq: Tensor)[源代码]#

FlexSN#

class spikingjelly.activation_based.neuron.flexsn.FlexSNKernel(core: Callable, num_inputs: int, num_states: int, num_outputs: int, example_inputs: Tuple[Tensor] | None = None, requires_grad: Tuple[bool] | None = None)[源代码]#

基类：object

API Language: 中文 | English

  • 中文

FlexSNKernel 可以根据自定义的 PyTorch 单步函数 core 生成 Triton 多步脉冲神经元核。 不同于 FlexSN ， FlexSNKernel 是对底层 custom op 调度的轻量 Callable 封装。

实例化后， FlexSNKernel 对象接受的输入参数为 [*input_seqs, *states] ，其中 input_seqs 是 num_inputs 个输入序列，states 是 num_states 个初始状态；返回值为 [*output_seqs, *state_seqs] ， 其中 output_seqs 是 num_outputs 个输出序列，state_seqs 是 num_states 个状态序列。

参数:

  • core (Callable) -- 描述单步前向推理的函数，签名应为 [*inputs, *states] -> [*outputs, *updated_states]，其中输入、输出和状态均为张量。

  • num_inputs (int) -- 输入序列的数量。

  • num_states (int) -- 初始状态张量的数量，同时也应与返回的更新后状态数量一致。

  • num_outputs (int) -- 输出序列的数量。

  • example_inputs (Optional[Tuple[torch.Tensor]]) -- 传给 core 的示例张量，形式为 [*inputs, *states]，用于辅助构建推理与训练 kernel。 若为 None，则由底层构建器使用默认示例张量。

  • requires_grad (Optional[Tuple[bool]]) -- 指示 core 各个输入参数是否需要梯度的布尔元组，仅用于训练 kernel 的构建。 若为 None，则由底层构建器采用默认行为。

抛出:

RuntimeError -- 当前环境未启用 CUDA 时抛出，因为 FlexSNKernel 仅支持运行在 CUDA 设备上的 Triton 内核。

  • English

FlexSNKernel can generate Triton multi-step spiking neuron kernels from a customized PyTorch single-step function core via FlexSN's triton / inductor backend. It is a lightweight Callable wrapper over the underlying FlexSN custom-op dispatch path.

The input arguments of a FlexSNKernel object is [*input_seqs, *states] , where input_seqs is a list of input sequences, states is a list of initial states; the return value is [*output_seqs, *state_seqs] , where output_seqs is a list of output sequences, and state_seqs is a list of state sequences.

参数:

  • core (Callable) -- function describing the single-step inference dynamics with signature [*inputs, *states] -> [*outputs, *updated_states]. Inputs, outputs, and states should all be tensors.

  • num_inputs (int) -- number of input sequences.

  • num_states (int) -- number of initial state tensors, which should also match the number of updated states returned by core.

  • num_outputs (int) -- number of output sequences.

  • example_inputs (Optional[Tuple[torch.Tensor]]) -- example tensors passed to core in the form [*inputs, *states]. They are used to help build inference and training kernels. If None, the backend builders use their default example tensors.

  • requires_grad (Optional[Tuple[bool]]) -- tuple indicating whether each argument of core requires gradients. It is only used when building the training kernels. If None, the backend builders use their default behavior.

抛出:

RuntimeError -- raised when CUDA is unavailable, because FlexSNKernel only supports the Triton kernels on CUDA devices.

返回:

None

返回类型:

None

class spikingjelly.activation_based.neuron.flexsn.FlexSN(core: Callable, num_inputs: int, num_states: int, num_outputs: int, example_inputs: Tuple[Tensor] | None = None, requires_grad: Tuple[bool] | None = None, step_mode: str = 'm', backend: str = 'triton', store_state_seqs: bool = False, example_outputs: Tuple[Tensor] | None = None)[源代码]#

基类：MemoryModule

API Language: 中文 | English

  • 中文

FlexSN 可以根据自定义的 PyTorch 单步函数 core 生成 Triton 多步脉冲神经元。 FlexSN 在 FlexSNKernel 的基础上，进一步实现了其他 SpikingJelly 神经元的功能。 实例化后，FlexSN 对象输入和输出的语义与其他 SpikingJelly 神经元一致，取决于步进模式 step_mode 。

参数:

  • core (Callable) -- 描述单步前向推理的函数，签名应为 [*inputs, *states] -> [*outputs, *updated_states]， 其中输入与输出均为张量。“inputs”和“outputs”的数量任意，需用 num_inputs 和 num_outputs 指明。 “states”的数量任意，与“updated_states”数量一致，且需用 num_states 指明。

  • num_inputs (int) -- 输入的数量，应严格对应 core 参数及 example_inputs 中“inputs”的数量。

  • num_states (int) -- 状态的数量，应严格对应 core 的参数、返回值及 example_inputs 中的“states”的数量。

  • num_outputs (int) -- 输出的数量，应严格对应 core 返回值中“outputs”的数量。

  • example_inputs (Optional[Tuple[torch.Tensor]]) -- 提供给 core 的示例输入，形式为 [*inputs, *states]。 这些张量都会自动被放置到 cuda 设备上。若为 None ，则自动生成 num_inputs + num_states 个仅含一个元素的张量。默认为 None 。

  • requires_grad (Optional[Tuple[bool]]) -- 指示 core 的参数 (即 [*inputs, *states]) 是否需要梯度。 用于生成前向和反向计算图。长度应与 core 的参数及 example_inputs 对应。 若为 None，则所有参数均需梯度。默认 None。

  • step_mode (str) -- 步进模式。"triton" 和 "inductor" 内核仅在 "m" 模式下可用。默认 "m"。

  • backend (str) -- 使用的后端。"triton"、"inductor" 和 "hop" 仅在 step_mode="m" 时可用; "torch" 始终可用。"triton" 与 "inductor" 在 FlexSN 中是并列且等价的 Triton 类标签， 当前共享同一条维护中的 Triton 执行路径。默认 "triton"。

  • store_state_seqs (bool) -- 是否保存状态序列。如果为 True，用户可以通过 state_seqs 属性访问。 state_seqs 是个列表，每个元素是形状为 [T, ...] 的张量。默认 False。

  • example_outputs (Optional[Tuple[torch.Tensor]]) -- core 的单步输出模板，形式为 tuple([*outputs])。 当 backend="torch" 且输入为空序列 T == 0 时, 需要用它来构造输出张量的 形状和 dtype, 从而避免为了推断输出而执行 core。对于 "triton" 与 "inductor" 这两个等价的后端, 若提供该参数, 每个模板张量都必须与第一个 example_inputs 张量的单步形状和 dtype 相匹配。"hop" 后端会保留任意 输出模板, 并在空序列/HOP 路径中按运行时输入设备物化它们, 不执行上述形状/dtype 校验。若不需要空序列模板, 则可以为 None。默认 None。

  • English

FlexSN can generate Triton multi-step spiking neuron from a customized PyTorch single-step function core . FlexSN is built upon FlexSNKernel and further implements other features of SpikingJelly neurons. The input / output semantics of a FlexSN object is similar to those of other SpikingJelly neurons, depending on step_mode .

参数:

  • core (Callable) -- a function describing the single-step inference dynamics of the spiking neuron. Its signature should be [*inputs, *states] -> [*outputs, *updated_states], and the arguments and return values should all be tensors. There can be arbitrary number of inputs and outputs (specified by num_inputs and num_outputs). There can be arbitrary number of states (specified by num_states), and the number of updated states should match the number of states.

  • num_inputs (int) -- number of inputs. It should strictly match the number of inputs" in core's arguments and example_inputs.

  • num_states (int) -- number of states. It should strictly match the number of "states" in core's arguments, core's return values, and example_inputs.

  • num_outputs (int) -- number of outputs. It should strictly match the number of "outputs" in core's return values.

  • example_inputs (Optional[Tuple[torch.Tensor]]) -- example inputs to core with the form of [*inputs, *states]. These tensors will be moved to cuda device. If None, example_inputs will be num_inputs + num_states tensors with single element. Defaults to None.

  • requires_grad (Optional[Tuple[bool]]) -- whether the core's arguments (i.e. [*inputs, *states]) requires gradients. This info is used to generate the forward and backward graphs. Its length should match the number of core's arguments and the length of example_inputs. If None, all argument tensors require grad. Defaults to None.

  • step_mode (str) -- step mode. "triton" and "inductor" backends are available only in "m" mode. Defaults to "m".

  • backend (str) -- backend to use. "triton", "inductor", and "hop" are available only when step_mode="m". "torch" is always available. In FlexSN, "triton" and "inductor" are peer labels and currently dispatch the same maintained Triton execution path. Defaults to "triton".

  • store_state_seqs (bool) -- whether to store the state sequences. If True, users can access the state sequences via state_seqs property. state_seqs is a list of tensors with shape [T, ...]. Defaults to False.

  • example_outputs (Optional[Tuple[torch.Tensor]]) -- per-step output templates for core with the form of tuple([*outputs]). When backend="torch" and the input sequence is empty (T == 0), these templates are required to materialize output shapes and dtypes without executing core. For the equivalent "triton" / "inductor" backends, each provided template must match the first example_inputs tensor's per-step shape and dtype. The "hop" backend intentionally allows arbitrary output templates and materializes them on the runtime input device for empty-sequence/HOP paths, so it does not enforce that scan-backend shape/dtype check. Defaults to None when empty-sequence output templates are not needed.

返回:

None

返回类型:

None

property kernel#
property backend#
static init_states(num_states: int, step_mode: str, *args) → List[Tensor][源代码]#

API Language: 中文 | English

  • 中文

初始化神经元的状态张量。用户可以通过重写此方法来自定义状态初始化规则。默认情况下，所有 状态均被初始化为零张量。

参数:

  • num_states (int) -- 状态变量的数量。应与 core 的“states”部分中的张量数量一致。

  • step_mode (str) -- 本模块当前所处的步进模式。可选值为 "s" 或 "m" 。

  • args (Sequence[torch.Tensor]) -- forward 的输入，即 [*inputs]。用户应根据 args 和 FlexSN 的 step_mode 等信息来决定状态张量的初始化方式。

返回:

初始化后的状态张量列表，顺序对应了 core 参数的“states”部分。

返回类型:

List[torch.Tensor]

  • English

Initialize the neuron state tensors. Users can override this method to customize the state initialization rules. By default, all state tensors are initialized to zero.

The state tensors are stored in the states attribute, which is a list of tensors, whose order corresponds to the "states" part of core.

参数:

  • num_states (int) -- number of states. It should strictly match the number of "states" in core's return values.

  • step_mode (str) -- the current step mode of this module. It can be "s" or "m" .

  • args (Sequence[torch.Tensor]) -- the input of forward, i.e., [*inputs]. Users should initialize state tensors based on args and step_mode.

返回:

the list of initialized state tensors, whose order corresponds to the "states" part of core.

返回类型:

List[torch.Tensor]

single_step_forward(*args)[源代码]#
multi_step_forward(*args)[源代码]#
forward(*args)[源代码]#