spikingjelly.activation_based.functional package

Module contents

spikingjelly.activation_based.functional.reset_net(net: Module)[源代码]
参数:

net – 任何属于 nn.Module 子类的网络

返回:

None

将网络的状态重置。做法是遍历网络中的所有 Module，若 m `` 为 ``base.MemoryModule 函数或者是拥有 reset() 方法，则调用 m.reset()。

参数:

net – Any network inherits from nn.Module

返回:

None

Reset the whole network. Walk through every Module as m, and call m.reset() if this m is base.MemoryModule or m has reset().

spikingjelly.activation_based.functional.set_step_mode(net: Module, step_mode: str)[源代码]
参数:

  • net (nn.Module) – 一个神经网络

  • step_mode (str) – ‘s’ (单步模式) 或 ‘m’ (多步模式)

返回:

None

将 net 中所有模块的步进模式设置为 step_mode 。

备注

spikingjelly.activation_based.layer.StepModeContainer, spikingjelly.activation_based.layer.ElementWiseRecurrentContainer, spikingjelly.activation_based.layer.LinearRecurrentContainer 的子模块（不包含包装器本身）的 step_mode 不会被改变。

参数:

  • net (nn.Module) – a network

  • step_mode (str) – ‘s’ (single-step) or ‘m’ (multi-step)

返回:

None

Set step_mode for all modules in net.

Note

The submodule (not including the container itself) of spikingjelly.activation_based.layer.StepModeContainer, spikingjelly.activation_based.layer.ElementWiseRecurrentContainer, spikingjelly.activation_based.layer.LinearRecurrentContainer will not be changed.

spikingjelly.activation_based.functional.set_backend(net: ~torch.nn.modules.module.Module, backend: str, instance: object = (<class 'torch.nn.modules.module.Module'>, ))[源代码]
参数:

  • net (nn.Module) – 一个神经网络

  • backend (str) – 使用哪个后端

  • instance (nn.Module or tuple[nn.Module]) – 类型为 instance 的模块后端会被改变

返回:

None

将 net 中 所有类型为 instance 的模块后端更改为 backend

参数:

  • net (nn.Module) – a network

  • backend (str) – the backend to be set

  • instance (nn.Module or tuple[nn.Module]) – the backend of which instance will be changed

返回:

None

Sets backends of all modules whose instance is instance in net to backend

spikingjelly.activation_based.functional.detach_net(net: Module)[源代码]
参数:

net – 任何属于 nn.Module 子类的网络

返回:

None

将网络与之前的时间步的计算图断开。做法是遍历网络中的所有 Module，若 m `` 为 ``base.MemoryModule 函数或者是拥有 detach() 方法，则调用 m.detach()。

参数:

net – Any network inherits from nn.Module

返回:

None

Detach the computation graph of the whole network from previous time-steps. Walk through every Module as m, and call m.detach() if this m is base.MemoryModule or m has detach().

spikingjelly.activation_based.functional.spike_similar_loss(spikes: Tensor, labels: Tensor, kernel_type='linear', loss_type='mse', *args)[源代码]
参数:

  • spikes – shape=[N, M, T]，N个数据生成的脉冲

  • labels – shape=[N, C]，N个数据的标签，labels[i][k] == 1表示数据i属于第k类，反之亦然，允许多标签

  • kernel_type (str) – 使用内积来衡量两个脉冲之间的相似性，kernel_type是计算内积时，所使用的核函数种类

  • loss_type (str) – 返回哪种损失，可以为’mse’, ‘l1’, ‘bce’

  • args – 用于计算内积的额外参数

返回:

shape=[1]的tensor，相似损失

将N个数据输入到输出层有M个神经元的SNN，运行T步，得到shape=[N, M, T]的脉冲。这N个数据的标签为shape=[N, C]的labels。

用shape=[N, N]的矩阵sim表示实际相似度矩阵，sim[i][j] == 1表示数据i与数据j相似，反之亦然。若labels[i]与labels[j]共享至少同一个标签，则认为他们相似，否则不相似。

用shape=[N, N]的矩阵sim_p表示输出相似度矩阵，sim_p[i][j]的取值为0到1，值越大表示数据i与数据j的脉冲越相似。

使用内积来衡量两个脉冲之间的相似性，kernel_type是计算内积时，所使用的核函数种类：

  • ‘linear’，线性内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}}\)。

  • ‘sigmoid’，Sigmoid内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{sigmoid}(\alpha \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})\)，其中 \(\alpha = args[0]\)。

  • ‘gaussian’，高斯内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{exp}(- \frac{||\boldsymbol{x_{i}} - \boldsymbol{y_{j}}||^{2}}{2\sigma^{2}})\)，其中 \(\sigma = args[0]\)。

当使用Sigmoid或高斯内积时，内积的取值范围均在[0, 1]之间；而使用线性内积时，为了保证内积取值仍然在[0, 1]之间，会进行归一化：按照 \(\text{sim_p}[i][j]=\frac{\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}})}{||\boldsymbol{x_{i}}|| · ||\boldsymbol{y_{j}}||}\)。

对于相似的数据，根据输入的loss_type，返回度量sim与sim_p差异的损失：

  • ‘mse’ – 返回sim与sim_p的均方误差（也就是l2误差）。

  • ‘l1’ – 返回sim与sim_p的l1误差。

  • ‘bce’ – 返回sim与sim_p的二值交叉熵误差。

备注

脉冲向量稀疏、离散，最好先使用高斯核进行平滑，然后再计算相似度。

参数:

  • spikes – shape=[N, M, T], output spikes corresponding to a batch of N inputs

  • labels – shape=[N, C], labels of inputs, labels[i][k] == 1 means the i-th input belongs to the k-th category and vice versa. Multi-label input is allowed.

  • kernel_type (str) – Type of kernel function used when calculating inner products. The inner product is the similarity measure of two spikes.

  • loss_type (str) – Type of loss returned. Can be: ‘mse’, ‘l1’, ‘bce’

  • args – Extra parameters for inner product

返回:

shape=[1], similarity loss

A SNN consisting M neurons will receive a batch of N input data in each timestep (from 0 to T-1) and output a spike tensor of shape=[N, M, T]. The label is a tensor of shape=[N, C].

The groundtruth similarity matrix sim has a shape of [N, N]. sim[i][j] == 1 indicates that input i is similar to input j and vice versa. If and only if labels[i] and labels[j] have at least one common label, they are viewed as similar.

The output similarity matrix sim_p has a shape of [N, N]. The value of sim_p[i][j] ranges from 0 to 1, represents the similarity between output spike from both input i and input j.

The similarity is measured by inner product of two spikes. kernel_type is the type of kernel function when calculating inner product:

  • ‘linear’, Linear kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}}\).

  • ‘sigmoid’, Sigmoid kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{sigmoid}(\alpha \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})\), where \(\alpha = args[0]\).

  • ‘gaussian’, Gaussian kernel，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{exp}(- \frac{||\boldsymbol{x_{i}} - \boldsymbol{y_{j}}||^{2}}{2\sigma^{2}})\), where \(\sigma = args[0]\).

When Sigmoid or Gaussian kernel is applied, the inner product naturally lies in \([0, 1]\). To make the value consistent when using linear kernel, the result will be normalized as: \(\text{sim_p}[i][j]=\frac{\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}})}{||\boldsymbol{x_{i}}|| · ||\boldsymbol{y_{j}}||}\).

For similar data, return the specified discrepancy loss between sim and sim_p according to loss_type.

  • ‘mse’ – Return the Mean-Square Error (squared L2 norm) between sim and sim_p.

  • ‘l1’ – Return the L1 error between sim and sim_p.

  • ‘bce’ – Return the Binary Cross Entropy between sim and sim_p.

Note

Since spike vectors are usually discrete and sparse, it would be better to apply Gaussian filter first to smooth the vectors before calculating similarities.

spikingjelly.activation_based.functional.kernel_dot_product(x: Tensor, y: Tensor, kernel='linear', *args)[源代码]
参数:

  • x – shape=[N, M]的tensor，看作是N个M维向量

  • y – shape=[N, M]的tensor，看作是N个M维向量

  • kernel (str) – 计算内积时所使用的核函数

  • args – 用于计算内积的额外的参数

返回:

ret, shape=[N, N]的tensor，ret[i][j]表示x[i]和y[j]的内积

计算批量数据x和y在核空间的内积。记2个M维tensor分别为 \(\boldsymbol{x_{i}}\) 和 \(\boldsymbol{y_{j}}\)，kernel定义了不同形式的内积：

  • ‘linear’，线性内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}}\)。

  • ‘polynomial’，多项式内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = (\boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})^{d}\)，其中 \(d = args[0]\)。

  • ‘sigmoid’，Sigmoid内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{sigmoid}(\alpha \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})\)，其中 \(\alpha = args[0]\)。

  • ‘gaussian’，高斯内积，\(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{exp}(- \frac{||\boldsymbol{x_{i}} - \boldsymbol{y_{j}}||^{2}}{2\sigma^{2}})\)，其中 \(\sigma = args[0]\)。

参数:

  • x – Tensor of shape=[N, M]

  • y – Tensor of shape=[N, M]

  • kernel (str) – Type of kernel function used when calculating inner products.

  • args – Extra parameters for inner product

返回:

ret, Tensor of shape=[N, N], ret[i][j] is inner product of x[i] and y[j].

Calculate inner product of x and y in kernel space. These 2 M-dim tensors are denoted by \(\boldsymbol{x_{i}}\) and \(\boldsymbol{y_{j}}\). kernel determine the kind of inner product:

  • ‘linear’ – Linear kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}}\).

  • ‘polynomial’ – Polynomial kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = (\boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})^{d}\), where \(d = args[0]\).

  • ‘sigmoid’ – Sigmoid kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{sigmoid}(\alpha \boldsymbol{x_{i}}^{T}\boldsymbol{y_{j}})\), where \(\alpha = args[0]\).

  • ‘gaussian’ – Gaussian kernel, \(\kappa(\boldsymbol{x_{i}}, \boldsymbol{y_{j}}) = \mathrm{exp}(- \frac{||\boldsymbol{x_{i}} - \boldsymbol{y_{j}}||^{2}}{2\sigma^{2}})\), where \(\sigma = args[0]\).

spikingjelly.activation_based.functional.set_threshold_margin(output_layer: BaseNode, label_one_hot: Tensor, eval_threshold=1.0, threshold0=0.9, threshold1=1.1)[源代码]
参数:

  • output_layer – 用于分类的网络的输出层，输出层输出shape=[batch_size, C]

  • label_one_hot – one hot格式的样本标签，shape=[batch_size, C]

  • eval_threshold (float) – 输出层神经元在测试（推理）时使用的电压阈值

  • threshold0 (float) – 输出层神经元在训练时，负样本的电压阈值

  • threshold1 (float) – 输出层神经元在训练时，正样本的电压阈值

返回:

None

对于用来分类的网络，为输出层神经元的电压阈值设置一定的裕量，以获得更好的分类性能。

类别总数为C，网络的输出层共有C个神经元。网络在训练时，当输入真实类别为i的数据，输出层中第i个神经元的电压阈值会被设置成threshold1，而其他神经元的电压阈值会被设置成threshold0。而在测试（推理）时，输出层中神经元的电压阈值被统一设置成eval_threshold。

参数:

  • output_layer – The output layer of classification network, where the shape of output should be [batch_size, C]

  • label_one_hot – Labels in one-hot format, shape=[batch_size, C]

  • eval_threshold (float) – Voltage threshold of neurons in output layer when evaluating (inference)

  • threshold0 (float) – Voltage threshold of the corresponding neurons of negative samples in output layer when training

  • threshold1 (float) – Voltage threshold of the corresponding neurons of positive samples in output layer when training

返回:

None

Set voltage threshold margin for neurons in the output layer to reach better performance in classification task.

When there are C different classes, the output layer contains C neurons. During training, when the input with groundtruth label i are sent into the network, the voltage threshold of the i-th neurons in the output layer will be set to threshold1 and the remaining will be set to threshold0.

During inference, the voltage thresholds of ALL neurons in the output layer will be set to eval_threshold.

spikingjelly.activation_based.functional.redundant_one_hot(labels: Tensor, num_classes: int, n: int)[源代码]
参数:

  • labels – shape=[batch_size]的tensor，表示batch_size个标签

  • num_classes (int) – 类别总数

  • n (int) – 表示每个类别所用的编码数量

返回:

shape=[batch_size, num_classes * n]的tensor

对数据进行冗余的one-hot编码，每一类用 n 个1和 (num_classes - 1) * n 个0来编码。

示例：

>>> num_classes = 3
>>> n = 2
>>> labels = torch.randint(0, num_classes, [4])
>>> labels
tensor([0, 1, 1, 0])
>>> codes = functional.redundant_one_hot(labels, num_classes, n)
>>> codes
tensor([[1., 1., 0., 0., 0., 0.],
        [0., 0., 1., 1., 0., 0.],
        [0., 0., 1., 1., 0., 0.],
        [1., 1., 0., 0., 0., 0.]])
参数:

  • labels – Tensor of shape=[batch_size], batch_size labels

  • num_classes (int) – The total number of classes.

  • n (int) – The encoding length for each class.

返回:

Tensor of shape=[batch_size, num_classes * n]

Redundant one-hot encoding for data. Each class is encoded to n 1’s and (num_classes - 1) * n 0’s

e.g.:

>>> num_classes = 3
>>> n = 2
>>> labels = torch.randint(0, num_classes, [4])
>>> labels
tensor([0, 1, 1, 0])
>>> codes = functional.redundant_one_hot(labels, num_classes, n)
>>> codes
tensor([[1., 1., 0., 0., 0., 0.],
        [0., 0., 1., 1., 0., 0.],
        [0., 0., 1., 1., 0., 0.],
        [1., 1., 0., 0., 0., 0.]])
spikingjelly.activation_based.functional.first_spike_index(spikes: Tensor)[源代码]
参数:

spikes – shape=[*, T]，表示任意个神经元在t=0, 1, …, T-1，共T个时刻的输出脉冲

返回:

index, shape=[*, T]，为 True 的位置表示该神经元首次释放脉冲的时刻

输入若干个神经元的输出脉冲，返回一个与输入相同shape的 bool 类型的index。index为 True 的位置，表示该神经元首次释放脉冲的时刻。

示例：

>>> spikes = (torch.rand(size=[2, 3, 8]) >= 0.8).float()
>>> spikes
tensor([[[0., 0., 0., 0., 0., 0., 0., 0.],
 [1., 0., 0., 0., 0., 0., 1., 0.],
 [0., 1., 0., 0., 0., 1., 0., 1.]],

[[0., 0., 1., 1., 0., 0., 0., 1.],
 [1., 1., 0., 0., 1., 0., 0., 0.],
 [0., 0., 0., 1., 0., 0., 0., 0.]]])
>>> first_spike_index(spikes)
tensor([[[False, False, False, False, False, False, False, False],
 [ True, False, False, False, False, False, False, False],
 [False,  True, False, False, False, False, False, False]],

[[False, False,  True, False, False, False, False, False],
 [ True, False, False, False, False, False, False, False],
 [False, False, False,  True, False, False, False, False]]])
参数:

spikes – shape=[*, T], indicates the output spikes of some neurons when t=0, 1, …, T-1.

返回:

index, shape=[*, T], the index of True represents the moment of first spike.

Return an index tensor of the same shape of input tensor, which is the output spike of some neurons. The index of True represents the moment of first spike.

e.g.:

>>> spikes = (torch.rand(size=[2, 3, 8]) >= 0.8).float()
>>> spikes
tensor([[[0., 0., 0., 0., 0., 0., 0., 0.],
 [1., 0., 0., 0., 0., 0., 1., 0.],
 [0., 1., 0., 0., 0., 1., 0., 1.]],

[[0., 0., 1., 1., 0., 0., 0., 1.],
 [1., 1., 0., 0., 1., 0., 0., 0.],
 [0., 0., 0., 1., 0., 0., 0., 0.]]])
>>> first_spike_index(spikes)
tensor([[[False, False, False, False, False, False, False, False],
 [ True, False, False, False, False, False, False, False],
 [False,  True, False, False, False, False, False, False]],

[[False, False,  True, False, False, False, False, False],
 [ True, False, False, False, False, False, False, False],
 [False, False, False,  True, False, False, False, False]]])
spikingjelly.activation_based.functional.multi_step_forward(x_seq: Tensor, single_step_module: Module)[源代码]
参数:
返回:

shape=[T, batch_size, ...] 的输出tensor

返回类型:

torch.Tensor

在单步模块 single_step_module 上使用多步前向传播。

参数:
返回:

the output tensor with shape=[T, batch_size, ...]

返回类型:

torch.torch.Tensor

Applies multi-step forward on single_step_module.

spikingjelly.activation_based.functional.chunk_multi_step_forward(split_size: int, x_seq: Tensor, multi_step_module: Module)[源代码]
参数:

  • split_size (int) – 分割的尺寸

  • x_seq (Tensor) – 输入

  • multi_step_module (nn.Module) – 一个使用多步传播模式的网络

返回:

输出

返回类型:

Tensor

将 shape = [T, *] 的输入 x_seq 拆分成多个 shape = [split_size, *] 的小tensor(若 T % split_size != 0，最后 一个tensor的 shape[0] 会小于 split_size)，然后逐个输入到 multi_step_module 中，再将输出重新拼接为 shape = [split_size, *]。 chunk_multi_step_forward 可以在使用很大的 T 进行不带梯度的推理(例如ANN2SNN)时使用，能够减少内存消耗量。

示例代码：

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

net = nn.Sequential(
    layer.Linear(8, 4),
    neuron.IFNode(step_mode='m'),
    layer.Linear(4, 2),
    neuron.IFNode(step_mode='m'),
)

x_seq = torch.rand([1024, 8])
with torch.no_grad():
    y_seq = functional.chunk_multi_step_forward(16, x_seq, net)
    print(y_seq.shape)
    # torch.Size([1024, 2])
参数:

  • split_size (int) – the split size

  • x_seq (Tensor) – the input tensor

  • multi_step_module (nn.Module) –

返回:

the output tensor

返回类型:

Tensor

Splits the input x_seq with shape = [T, *] to many tensor chunks with shape = [split_size, *] (if T % split_size != 0, shape[0] of the last tensor chunk will be smaller than split_size), and sends chunks to multi_step_module, then concatenates the outputs to shape = [split_size, *].

chunk_multi_step_forward can be used for inference with a large T (e.g., ANN2SNN) to reduce the memory consumption.

Codes example:

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

net = nn.Sequential(
    layer.Linear(8, 4),
    neuron.IFNode(step_mode='m'),
    layer.Linear(4, 2),
    neuron.IFNode(step_mode='m'),
)

x_seq = torch.rand([1024, 8])
with torch.no_grad():
    y_seq = functional.chunk_multi_step_forward(16, x_seq, net)
    print(y_seq.shape)
    # torch.Size([1024, 2])
spikingjelly.activation_based.functional.seq_to_ann_forward(x_seq: Tensor, stateless_module: Module)[源代码]
参数:
返回:

the output tensor with shape=[T, batch_size, ...]

返回类型:

Tensor

参数:
返回:

the output tensor with shape=[T, batch_size, ...]

返回类型:

Tensor

Applied forward on stateless modules

spikingjelly.activation_based.functional.fused_conv2d_weight_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d)[源代码]
参数:
返回:

the weight of this fused module

返回类型:

Tensor

{Conv2d-BatchNorm2d} 模块可以合并为一个单个的 {Conv2d}，其中``BatchNorm2d`` 的参数会被吸收进 Conv2d。 本函数返回合并后的卷积的权重。

备注

这里按照 conv2d.bias 为 None 进行处理。原因参见 Disable bias for convolutions directly followed by a batch norm 。

参数:
返回:

the weight of this fused module

返回类型:

Tensor

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d ‘s parameters being absorbed into Conv2d. This function returns the weight of this fused module.

Note

We assert conv2d.bias is None. See Disable bias for convolutions directly followed by a batch norm for more details.

spikingjelly.activation_based.functional.fused_conv2d_bias_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d)[源代码]
参数:
返回:

the weight of this fused module

返回类型:

Tensor

{Conv2d-BatchNorm2d} 模块可以合并为一个单个的 {Conv2d}，其中``BatchNorm2d`` 的参数会被吸收进 Conv2d。 本函数返回合并后的卷积的偏置项。

备注

这里按照 conv2d.bias 为 None 进行处理。原因参见 Disable bias for convolutions directly followed by a batch norm 。

参数:
返回:

the weight of this fused module

返回类型:

Tensor

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d ‘s parameters being absorbed into Conv2d. This function returns the bias of this fused module.

Note

We assert conv2d.bias is None. See Disable bias for convolutions directly followed by a batch norm for more details.

spikingjelly.activation_based.functional.scale_fused_conv2d_weight_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d, k=None, b=None)[源代码]
参数:
返回:

the weight of this fused module

返回类型:

Tensor

{Conv2d-BatchNorm2d} 模块可以合并为一个单个的 {Conv2d}，其中``BatchNorm2d`` 的参数会被吸收进 Conv2d。 本函数对 {Conv2d-BatchNorm2d} 模块整体的等效权重进行 weight = k * weight + b 的线性变换。

备注

这里按照 conv2d.bias 为 None 进行处理。原因参见 Disable bias for convolutions directly followed by a batch norm 。

参数:
返回:

the weight of this fused module

返回类型:

Tensor

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d ‘s parameters being absorbed into Conv2d. This function applies a linear transform weight = k * weight + b on the equivalent weight of the whole {Conv2d-BatchNorm2d}.

Note

We assert conv2d.bias is None. See Disable bias for convolutions directly followed by a batch norm for more details.

spikingjelly.activation_based.functional.scale_fused_conv2d_bias_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d, k=None, b=None)[源代码]
参数:
返回:

the weight of this fused module

返回类型:

Tensor

{Conv2d-BatchNorm2d} 模块可以合并为一个单个的 {Conv2d}，其中``BatchNorm2d`` 的参数会被吸收进 Conv2d。 本函数对 {Conv2d-BatchNorm2d} 模块整体的等效偏置项进行 bias = k * bias + b 的线性变换。

备注

这里按照 conv2d.bias 为 None 进行处理。原因参见 Disable bias for convolutions directly followed by a batch norm 。

参数:
返回:

the weight of this fused module

返回类型:

Tensor

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d ‘s parameters being absorbed into Conv2d. This function applies a linear transform bias = k * bias + b on the equivalent bias of the whole {Conv2d-BatchNorm2d}.

Note

We assert conv2d.bias is None. See Disable bias for convolutions directly followed by a batch norm for more details.

spikingjelly.activation_based.functional.fuse_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d)[源代码]
参数:
返回:

the weight of this fused module

返回类型:

Tensor

{Conv2d-BatchNorm2d} 模块可以合并为一个单个的 {Conv2d}，其中``BatchNorm2d`` 的参数会被吸收进 Conv2d。 本函数对返回这个等效的合并后的 {Conv2d}。

备注

这里按照 conv2d.bias 为 None 进行处理。原因参见 Disable bias for convolutions directly followed by a batch norm 。

参数:
返回:

the weight of this fused module

返回类型:

Tensor

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d ‘s parameters being absorbed into Conv2d. This function returns the fused {Conv2d} merged by {Conv2d-BatchNorm2d}.

Note

We assert conv2d.bias is None. See Disable bias for convolutions directly followed by a batch norm for more details.

spikingjelly.activation_based.functional.temporal_efficient_training_cross_entropy(x_seq: Tensor, target: Tensor)[源代码]
参数:

  • x_seq (torch.Tensor) – shape=[T, N, C, *] 的预测值，其中 C 是类别总数

  • target (torch.Tensor) – shape=[N] 的真实值，其中 target[i] 是真实类别

返回:

the temporal efficient training cross entropy

返回类型:

torch.Tensor

Temporal efficient training (TET) 交叉熵损失, 是每个时间步的交叉熵损失的平均。

示例代码：

def tet_ce_for_loop_version(x_seq: torch.Tensor, target: torch.LongTensor):
    loss = 0.
    for t in range(x_seq.shape[0]):
        loss += F.cross_entropy(x_seq[t], target)
    return loss / x_seq.shape[0]


T = 8
N = 4
C = 10
x_seq = torch.rand([T, N, C])
target = torch.randint(low=0, high=C - 1, size=[N])
print(f'max error = {(tet_ce_for_loop_version(x_seq, target) - temporal_efficient_training_cross_entropy(x_seq, target)).abs().max()}')
# max error < 1e-6

备注

TET交叉熵是 Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting 一文提出的。

参数:

  • x_seq (torch.Tensor) – the predicted value with shape=[T, N, C, *], where C is the number of classes

  • target (torch.Tensor) – the ground truth tensor with shape=[N], where target[i] is the label

返回:

the temporal efficient training cross entropy

返回类型:

torch.Tensor

The temporal efficient training (TET) cross entropy, which is the mean of cross entropy of each time-step.

Codes example:

def tet_ce_for_loop_version(x_seq: torch.Tensor, target: torch.LongTensor):
    loss = 0.
    for t in range(x_seq.shape[0]):
        loss += F.cross_entropy(x_seq[t], target)
    return loss / x_seq.shape[0]


T = 8
N = 4
C = 10
x_seq = torch.rand([T, N, C])
target = torch.randint(low=0, high=C - 1, size=[N])
print(f'max error = {(tet_ce_for_loop_version(x_seq, target) - temporal_efficient_training_cross_entropy(x_seq, target)).abs().max()}')
# max error < 1e-6

Note

The TET cross entropy is proposed by Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting.

spikingjelly.activation_based.functional.kaiming_normal_conv_linear_weight(net: Module)[源代码]
参数:

net – 任何属于 nn.Module 子类的网络

返回:

None

使用kaiming normal初始化 ``net` 中的所有 :class:`torch.nn._ConvNd 和 torch.nn.Linear 的权重（不包括偏置项）。参见 torch.nn.init.kaiming_normal_。

参数:

net – Any network inherits from nn.Module

返回:

None

initialize all weights (not including bias) of torch.nn._ConvNd and torch.nn.Linear in net by the kaiming normal. See torch.nn.init.kaiming_normal_ for more details.

spikingjelly.activation_based.functional.delay(x_seq: Tensor, delay_steps: int)[源代码]
参数:

  • x_seq (torch.Tensor) – 输入的序列，shape = [T, *]

  • delay_steps (int) – 延迟的时间步数

返回:

延迟后的序列

返回类型:

torch.Tensor

延迟函数，可以用来延迟输入，使得 y[t] = x[t - delay_steps]。缺失的数据用0填充。

代码示例：

x = torch.rand([5, 2])
x[3:].zero_()
x.requires_grad = True
y = delay(x, 1)
print('x=')
print(x)
print('y=')
print(y)
y.sum().backward()
print('x.grad=')
print(x.grad)

输出为：

x=
tensor([[0.1084, 0.5698],
        [0.4563, 0.3623],
        [0.0556, 0.4704],
        [0.0000, 0.0000],
        [0.0000, 0.0000]], requires_grad=True)
y=
tensor([[0.0000, 0.0000],
        [0.1084, 0.5698],
        [0.4563, 0.3623],
        [0.0556, 0.4704],
        [0.0000, 0.0000]], grad_fn=<CatBackward0>)
x.grad=
tensor([[1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.],
        [0., 0.]])
参数:

  • x_seq (torch.Tensor) – the input sequence with shape = [T, *]

  • delay_steps (int) – the number of delayed time-steps

返回:

the delayed sequence

返回类型:

torch.Tensor

A delay function that can delay inputs and makes y[t] = x[t - delay_steps]. The nonexistent data will be regarded as 0.

Codes example:

x = torch.rand([5, 2])
x[3:].zero_()
x.requires_grad = True
y = delay(x, 1)
print('x=')
print(x)
print('y=')
print(y)
y.sum().backward()
print('x.grad=')
print(x.grad)

The outputs are:

x=
tensor([[0.1084, 0.5698],
        [0.4563, 0.3623],
        [0.0556, 0.4704],
        [0.0000, 0.0000],
        [0.0000, 0.0000]], requires_grad=True)
y=
tensor([[0.0000, 0.0000],
        [0.1084, 0.5698],
        [0.4563, 0.3623],
        [0.0556, 0.4704],
        [0.0000, 0.0000]], grad_fn=<CatBackward0>)
x.grad=
tensor([[1., 1.],
        [1., 1.],
        [1., 1.],
        [1., 1.],
        [0., 0.]])
spikingjelly.activation_based.functional.fptt_online_training_init_w_ra(optimizer: Optimizer) → list[源代码]
spikingjelly.activation_based.functional.fptt_online_training(model: Module, optimizer: Optimizer, x_seq: Tensor, target_seq: Tensor, f_loss_t: Callable, alpha: float, w_ra: list) → None[源代码]
参数:

The FPTT online learning method proposed by Training Recurrent Neural Networks via Forward Propagation Through Time and used for SNN in Accurate online training of dynamical spiking neural networks through Forward Propagation Through Time .

Example:

from spikingjelly.activation_based import neuron

net = nn.Sequential(
    nn.Linear(8, 4),
    neuron.IFNode(),
    nn.Linear(4, 2),
    neuron.IFNode()
)

optimizer = torch.optim.SGD(net.parameters(), lr=0.1)

T = 4
N = 2
w_ra = fptt_online_training_init_w_ra(optimizer)
for epoch in range(2):

    x_seq = torch.rand([T, N, 8])
    target_seq = torch.rand([T, N, 2])

    fptt_online_training(model=net, optimizer=optimizer, x_seq=x_seq, target_seq=target_seq, f_loss_t=F.mse_loss, alpha=0.1, w_ra=w_ra)
    functional.reset_net(net)
spikingjelly.activation_based.functional.ottt_online_training(model: Module, optimizer: Optimizer, x_seq: Tensor, target_seq: Tensor, f_loss_t: Callable, online: bool) → None[源代码]
参数:

  • model (nn.Module) – the neural network

  • optimizer (torch.optim.Optimizer) – the optimizer for the network

  • x_seq (torch.Tensor) – the input sequence

  • target_seq (torch.Tensor) – the output sequence

  • f_loss_t (Callable) – the loss function, which should has the formulation of def f_loss_t(x_t, y_t) -> torch.Tensor

  • online (bool) – whether online update parameters or accumulate gradients through time steps

The OTTT online training method is proposed by Online Training Through Time for Spiking Neural Networks. This function can also be used for SLTT training method proposed by Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural Networks .

Example:

from spikingjelly.activation_based import neuron, layer, functional

net = layer.OTTTSequential(
    nn.Linear(8, 4),
    neuron.OTTTLIFNode(),
    nn.Linear(4, 2),
    neuron.LIFNode()
)

optimizer = torch.optim.SGD(net.parameters(), lr=0.1)

T = 4
N = 2
online = True
for epoch in range(2):

    x_seq = torch.rand([N, T, 8])
    target_seq = torch.rand([N, T, 2])

    functional.ottt_online_training(model=net, optimizer=optimizer, x_seq=x_seq, target_seq=target_seq, f_loss_t=F.mse_loss, online=online)
    functional.reset_net(net)