spikingjelly.activation_based.base module#

class spikingjelly.activation_based.base.StepModule[源代码]#

基类：object

API Language - 中文 | English

  • 中文

步进模式接口基类。

实现该接口的模块通过 step_mode 区分单步模式 "s" 与多步模式 "m"。

  • English

Base interface for step-mode aware modules.

Modules implementing this interface distinguish single-step mode "s" from multi-step mode "m" through step_mode.

supported_step_mode()[源代码]#

API Language - 中文 | English

  • 中文

返回:

包含支持的步进模式的tuple。"s" 代表单步模式， "m" 代表多步模式。

返回类型:

Tuple[str]

  • English

返回:

a tuple that contains the supported step mode(s). "s" is for single-step mode, and "m" is for multi-step mode.

返回类型:

Tuple[str]

property step_mode: str#

API Language - 中文 | English

  • 中文

返回:

模块当前使用的步进模式

返回类型:

str

  • English

返回:

the current step mode of this module

返回类型:

str

class spikingjelly.activation_based.base.SingleStepModule[源代码]#

基类：StepModule

API Language - 中文 | English

  • 中文

单步模式模块的接口基类。

实现该接口的模块仅支持单步模式 "s"。

  • English

Base interface for single-step mode modules.

Modules implementing this interface only support single-step mode "s".

supported_step_mode()[源代码]#

API Language - 中文 | English

  • 中文

返回:

仅包含 "s" 的 tuple

返回类型:

Tuple[str]

  • English

返回:

A tuple containing only "s"

返回类型:

Tuple[str]

class spikingjelly.activation_based.base.MultiStepModule[源代码]#

基类：StepModule

API Language - 中文 | English

  • 中文

多步模式模块的接口基类。

实现该接口的模块仅支持多步模式 "m"。

  • English

Base interface for multi-step mode modules.

Modules implementing this interface only support multi-step mode "m".

supported_step_mode()[源代码]#

API Language - 中文 | English

  • 中文

返回:

仅包含 "m" 的 tuple

返回类型:

Tuple[str]

  • English

返回:

A tuple containing only "m"

返回类型:

Tuple[str]

class spikingjelly.activation_based.base.MemoryModule[源代码]#

基类：Module, StepModule

API Language - 中文 | English

  • 中文

SpikingJelly 中所有有状态模块的基类。

MemoryModule 通过 register_memory 注册内部状态变量，并提供 reset、detach、显式 memory 提取与恢复等通用能力。

  • English

Base class of all stateful modules in SpikingJelly.

MemoryModule registers internal state variables via register_memory and provides common utilities such as reset, detach, and explicit memory extraction / restoration.

materialize_states(inputs, states, step_mode)[源代码]#

API Language - 中文 | English

  • 中文

默认实现原样返回 states。需要将标量或空状态转换为输入相关张量的 模块可以重写本方法。inputs 是当前前向传播的完整输入；实现根据 step_mode 选择合适的形状和设备参照。实现不得修改模块 memory 或 传入的 states。

参数:

  • inputs (tuple[object, ...]) -- 当前前向传播的完整输入元组

  • states (tuple[object, ...]) -- 待物化的显式状态

  • step_mode (str) -- 当前前向传播的步进模式，取 "s" 或 "m"

返回:

可传入 functional forward 的状态

返回类型:

tuple[object, ...]

  • English

The default implementation returns states unchanged. Modules that need to turn scalar or empty states into input-dependent tensors may override this method. inputs contains the complete inputs of the current forward pass; implementations select an appropriate shape and device reference according to step_mode. Implementations must not mutate module memories or the supplied states.

参数:

  • inputs (tuple[object, ...]) -- Complete input tuple of the current forward pass

  • states (tuple[object, ...]) -- Explicit states to materialize

  • step_mode (str) -- Step mode of the current forward pass, "s" or "m"

返回:

States ready for functional forward

返回类型:

tuple[object, ...]

single_step_functional_forward(inputs, states, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

使用显式状态执行单步前向传播。实现不得修改本模块注册的 memory 或传入的 states。states 的顺序与当前模块的 self._memories 一致，不包含子模块状态。复合模块的完整状态由 to_functional_forward() 负责展开。

参数:
返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

抛出:

NotImplementedError -- 当子类未实现单步 functional forward 时抛出

  • English

Run a single-step forward pass with explicit state. The implementation must not mutate this module's registered memories or the supplied states. State order follows self._memories and excludes states owned by child modules. Use to_functional_forward() to flatten the complete state of a composite module.

参数:

  • inputs (tuple[Tensor, ...]) -- Tuple of single-step input tensors

  • states (tuple[object, ...]) -- Tuple of explicit states owned by this module

  • kwargs (object)

返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

抛出:

NotImplementedError -- If the subclass does not implement single-step functional forward

multi_step_functional_forward(inputs, states, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

使用显式状态执行多步前向传播。输入张量的第 0 维为时间维。 默认实现逐时间步调用 single_step_functional_forward()；具有 专用多步实现或 kernel 的子类可以重写本方法。实现不得修改本模块 注册的 memory 或传入的 states。

参数:
返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

  • English

Run a multi-step forward pass with explicit state. Dimension 0 of each input tensor is the time dimension. The default implementation calls single_step_functional_forward() at each time step. Subclasses with a specialized sequence implementation or kernel may override this method. The implementation must not mutate this module's registered memories or the supplied states.

参数:

  • inputs (tuple[Tensor, ...]) -- Tuple of multi-step input tensors

  • states (tuple[object, ...]) -- Tuple of explicit states owned by this module

  • kwargs (object)

返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

functional_forward(inputs, states, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

先调用 materialize_states()，再按当前 step_mode 调用单步或 多步 functional forward。

参数:
返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

抛出:

ValueError -- 当 step_mode 不是 "s" 或 "m" 时抛出

  • English

Call materialize_states(), then dispatch to single-step or multi-step functional forward according to the current step_mode.

参数:

  • inputs (tuple[Tensor, ...]) -- Tuple of input tensors

  • states (tuple[object, ...]) -- Tuple of explicit states owned by this module

  • kwargs (object)

返回:

(outputs, updated_states)

返回类型:

tuple[tuple[Tensor, ...], tuple[object, ...]]

抛出:

ValueError -- If step_mode is neither "s" nor "m"

single_step_forward(x, *args, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

基于 single_step_functional_forward() 执行单步前向传播：读取 self._memories，执行显式状态转移，再写回更新后的状态。

参数:

x (Tensor) -- 输入张量，约定 shape = [N, *]，其中 N 通常为 batch 维

返回:

单步前向传播的输出

返回类型:

object

抛出:

ValueError -- 更新状态数量与已注册 memory 数量不一致时抛出

  • English

Run a single-step forward pass through single_step_functional_forward(): read self._memories, execute the explicit state transition, and write the updated states back.

参数:

x (Tensor) -- Input tensor, conventionally with shape = [N, *] where N is usually the batch dimension

返回:

Output of the single-step forward pass

返回类型:

object

抛出:

ValueError -- If the number of updated states does not match the registered memories

multi_step_forward(x_seq, *args, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

基于 multi_step_functional_forward() 执行多步前向传播：读取 self._memories，执行显式状态转移，再写回更新后的状态。

参数:

x_seq (Tensor) -- 输入序列张量，约定 shape = [T, N, *]，其中第 0 维为时间维

返回:

按时间堆叠的输出序列

返回类型:

Tensor

抛出:

ValueError -- 更新状态数量与已注册 memory 数量不一致时抛出

  • English

Run a multi-step forward pass through multi_step_functional_forward(): read self._memories, execute the explicit state transition, and write the updated states back.

参数:

x_seq (Tensor) -- Input sequence tensor, conventionally with shape = [T, N, *] and the time axis at dimension 0

返回:

Output sequence stacked along the time dimension

返回类型:

Tensor

抛出:

ValueError -- If the number of updated states does not match the registered memories

forward(*args, **kwargs)[源代码]#

API Language - 中文 | English

  • 中文

按 step_mode 调用 single_step_forward() 或 multi_step_forward()。functional 状态读写由这两个方法负责。

返回:

与当前 step_mode 对应的前向传播结果

返回类型:

object

抛出:

ValueError -- 当 self.step_mode 既不是 "s" 也不是 "m" 时抛出

  • English

Dispatch to single_step_forward() or multi_step_forward() according to step_mode. Those methods own functional state loading and committing.

返回:

Forward result selected according to the current step_mode

返回类型:

object

抛出:

ValueError -- Raised when self.step_mode is neither "s" nor "m"

register_memory(name, value)[源代码]#

API Language - 中文 | English

  • 中文

将变量存入用于保存有状态变量（例如脉冲神经元的膜电位）的字典中。这个变量将被初始化为 value 。 每次调用 self.reset() 函数后， self.name 都会被重置为 value 。

警告

若状态变量是个 torch.Tensor ，则 不应对其做原地修改操作 。

参数:

  • name (str) -- 状态变量的名字

  • value (object) -- 状态变量的初始与重制值

抛出:

AssertionError -- 当 name 已经是模块现有成员属性时抛出

  • English

Register the state variable to memory dict, which saves stateful variables (e.g., the membrane potential of a spiking neuron). The variable will be initialized as value . self.name will be set to value after calling self.reset() .

警告

Do not modify the state variable in-place if it's a torch.Tensor .

参数:

  • name (str) -- state variable's name

  • value (object) -- state variable's initial and reset value

抛出:

AssertionError -- Raised when name already exists as an attribute of the module

reset()[源代码]#

API Language - 中文 | English

  • 中文

重置所有有状态变量为重制值。

若当前状态与重制值均为同形状、同 dtype、同 device 的张量，则优先原地恢复； 否则使用复制或重新赋值恢复。

  • English

Reset all stateful variables to their reset values.

If both the current state and the reset value are tensors with the same shape, dtype, and device, the state is restored in-place whenever possible; otherwise it falls back to copy or reassignment.

set_reset_value(name, value)[源代码]#

API Language - 中文 | English

  • 中文

设置状态变量 self.name 的重制值。

参数:

  • name (str) -- 状态变量名称

  • value (object) -- 新的重制值

  • English

Set the reset value of state variable self.name .

参数:

  • name (str) -- Name of the state variable

  • value (object) -- New reset value

get_reset_value(name)[源代码]#

API Language - 中文 | English

  • 中文

获取状态变量 self.name 的重置值。

参数:

name (str) -- 状态变量名称

返回:

状态变量的重置值

返回类型:

object

抛出:

KeyError -- 当 name 不是已注册的状态变量，或没有设置重置值时抛出

  • English

Get the reset value of state variable self.name.

参数:

name (str) -- Name of the state variable

返回:

Reset value of the state variable

返回类型:

object

抛出:

KeyError -- Raised when name is not a registered state variable, or has no reset value

memories()[源代码]#

API Language - 中文 | English

  • 中文

返回:

返回一个所有状态变量的生成器

返回类型:

Generator

  • English

返回:

a generator over all stateful variables

返回类型:

Generator

named_memories()[源代码]#

API Language - 中文 | English

  • 中文

返回:

返回一个所有状态变量名称及其值的生成器

返回类型:

Generator

  • English

返回:

a generator over all stateful variables' names and values

返回类型:

Generator

detach()[源代码]#

API Language - 中文 | English

  • 中文

从计算图中分离所有有状态变量。

小技巧

可以使用这个函数实现TBPTT (Truncated Back Propagation Through Time)。

  • English

Detach all stateful variables.

Tip

We can use this function to implement TBPTT (Truncated Back Propagation Through Time).

spikingjelly.activation_based.base.named_memories(module, prefix='')[源代码]#

API Language - 中文 | English

  • 中文

递归地生成模块树中的所有状态变量。类似于 named_parameters() 方法。

参数:

  • module (Module) -- 目标模块

  • prefix (str) -- 名称前缀

返回:

状态变量名称和值的生成器

返回类型:

Generator

抛出:

RecursionError -- 若模块树存在异常递归结构，Python 递归遍历时会抛出异常

  • English

Recursively yield all memory variables in a module tree. Similar to named_parameters() .

参数:

  • module (Module) -- the target module

  • prefix (str) -- name prefix

返回:

a generator of memory variable names and values

返回类型:

Generator

抛出:

RecursionError -- Raised if traversing the module tree exceeds Python recursion limits

spikingjelly.activation_based.base.memories(module)[源代码]#

API Language - 中文 | English

  • 中文

递归地生成模块树中的所有状态变量值。类似于 parameters() 方法。

参数:

module (nn.Module) -- 目标模块

返回:

状态变量值的生成器

返回类型:

Generator

抛出:

RecursionError -- 若模块树存在异常递归结构，Python 递归遍历时会抛出异常

  • English

Recursively yield all memory variables in a module tree. Similar to parameters() .

参数:

module (nn.Module) -- the target module

返回:

a generator of memory variable values

返回类型:

Generator

抛出:

RecursionError -- Raised if traversing the module tree exceeds Python recursion limits

spikingjelly.activation_based.base.extract_memories(module)[源代码]#

API Language - 中文 | English

  • 中文

提取模块中所有的状态变量值并返回列表。

参数:

module (Module) -- 目标模块

返回:

状态变量值的列表

返回类型:

list

抛出:

RecursionError -- 若模块树存在异常递归结构，Python 递归遍历时会抛出异常

  • English

Extract all memory variable values from the module and return as a list.

参数:

module (Module) -- the target module

返回:

a list of memory variable values

返回类型:

list

抛出:

RecursionError -- Raised if traversing the module tree exceeds Python recursion limits

spikingjelly.activation_based.base.load_memories(module, memory_list)[源代码]#

API Language - 中文 | English

  • 中文

将状态变量列表加载到模块中。

参数:

  • module (Module) -- 目标模块

  • memory_list (list) -- 状态变量值列表

抛出:

ValueError -- 当 memory_list 的长度与 module 当前状态变量数量不一致时抛出

  • English

Load memory variables from a list into the module.

参数:

  • module (Module) -- the target module

  • memory_list (list) -- list of memory variable values

抛出:

ValueError -- Raised when the length of memory_list does not match the number of current memory variables in module

spikingjelly.activation_based.base.to_functional_forward(module, fn=None)[源代码]#

API Language - 中文 | English

  • 中文

给定一个可能包含隐式状态变量（记忆，memory）的模块，获取其显式状态的前向传播函数。

返回的函数签名为 (inputs, states, **kwargs) -> (outputs, updated_states) ， 其中：

  • inputs 为原始 forward 所需的常规输入参数；

  • states 为当前模块中所有状态变量的值，其顺序与 extract_memories(module) 一致；

  • outputs 为原始 forward 的输出结果；

  • new_states 为执行前向传播后更新得到的状态变量。

单个输入和输出也使用 tuple 包装。状态顺序与 extract_memories(module) 一致。

备注

转换优先直接使用 MemoryModule 的 functional 实现，并使用同一方式组合 扁平的 nn.Sequential。其他模块通过既有 memory 接口临时换入显式状态。

警告

如果某个状态变量为 torch.Tensor ，则其不应在 module.forward 中被原地修改。否则， 会导致输入给前向传播函数的状态变量被修改，导致意想不到的错误。

参数:

  • module (Module) -- 目标模块

  • fn (Optional[Callable]) -- 含隐式状态的前向传播函数。若为 None ，则默认使用 module.forward 。 该参数可用于指定特殊的前向传播函数（如， module 的父类的 forward ）。默认值 为 None 。

返回:

带有显式输入输出状态的前向传播函数

返回类型:

Callable

抛出:

ValueError -- 调用时的状态数量与模块的 memory 数量不一致

  • English

Given a module that may contain implicit state variables, get the forward function with explicit state variables.

The returned function has the signature (inputs, states, **kwargs) -> (outputs, updated_states) where:

  • inputs are the regular input arguments required by the original forward;

  • states are the current memory variable values, in the same order as returned by extract_memories(module);

  • outputs are the outputs of the original forward method;

  • new_states are the updated memory variables after the forward pass.

Single inputs and outputs are still wrapped in tuples. State order matches extract_memories(module).

备注

Conversion first uses a MemoryModule functional implementation directly and composes flat nn.Sequential modules in the same way. Other modules temporarily load explicit states through the existing memory interface.

警告

If a state variable is a torch.Tensor, it should not be modified in-place in module.forward. Otherwise, the provided states will be modified, which may lead to unexpected errors.

参数:

  • module (Module) -- the target module

  • fn (Optional[Callable]) -- the forward function to be used. If None, module.forward is used by default. This argument can be used to explicitly specify another forward function (e.g., the forward method of module's parent class). Defaults to None.

返回:

a functional-style forward function with grouped inputs and states

返回类型:

Callable

抛出:

ValueError -- If the state count differs from the module's memory count

  • 代码示例 | Example

import torch
import torch.nn as nn
from spikingjelly.activation_based import base

class StatefulModule(base.MemoryModule):
    def __init__(self):
        super().__init__()
        self.register_memory("counter", torch.tensor(0.0))
        self.linear = nn.Linear(10, 5)

    def single_step_forward(self, x):
        self.counter = self.counter + 1.0
        return self.linear(x)

module = StatefulModule()
f_forward = base.to_functional_forward(module)
x = torch.randn(3, 10)
initial_state = torch.tensor(0.0)
outputs, new_states = f_forward((x,), (initial_state,))

assert torch.equal(outputs[0], module.linear(x))
assert torch.equal(new_states[0], initial_state + 1.0)