spikingjelly.clock_driven.functional package

Module contents

spikingjelly.clock_driven.functional.reset_net(net: torch.nn.modules.module.Module)[源代码]
  • API in English

参数

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

返回

None

将网络的状态重置。做法是遍历网络中的所有 Module,若含有 reset() 函数,则调用。

  • 中文API

参数

net – Any network inherits from nn.Module

返回

None

Reset the whole network. Walk through every Module and call their reset() function if exists.

spikingjelly.clock_driven.functional.spike_cluster(v: torch.Tensor, v_threshold, T_in: int)[源代码]
参数
  • v – shape=[T, N],N个神经元在 t=[0, 1, …, T-1] 时刻的电压值

  • v_threshold (float or tensor) – 神经元的阈值电压,float或者是shape=[N]的tensor

  • T_in – 脉冲聚类的距离阈值。一个脉冲聚类满足,内部任意2个相邻脉冲的距离不大于T_in,而其内部任一脉冲与外部的脉冲距离大于T_in

返回

一个元组,包含

  • N_o – shape=[N],N个神经元的输出脉冲的脉冲聚类的数量

  • k_positive – shape=[N],bool类型的tensor,索引。需要注意的是,k_positive可能是一个全False的tensor

  • k_negative – shape=[N],bool类型的tensor,索引。需要注意的是,k_negative可能是一个全False的tensor

返回类型

(Tensor, Tensor, Tensor)

STCA: Spatio-Temporal Credit Assignment with Delayed Feedback in Deep Spiking Neural Networks一文提出的脉冲聚类方法。如果想使用该文中定义的损失,可以参考如下代码:

v_k_negative = out_v * k_negative.float().sum(dim=0)
v_k_positive = out_v * k_positive.float().sum(dim=0)
loss0 = ((N_o > N_d).float() * (v_k_negative - 1.0)).sum()
loss1 = ((N_o < N_d).float() * (1.0 - v_k_positive)).sum()
loss = loss0 + loss1
参数
  • v – shape=[T, N], membrane potentials of N neurons when t=[0, 1, …, T-1]

  • v_threshold (float or tensor) – Threshold voltage(s) of the neurons, float or tensor of the shape=[N]

  • T_in – Distance threshold of the spike clusters. A spike cluster satisfies that the distance of any two adjacent spikes within cluster is NOT greater than T_in and the distance between any internal and any external spike of cluster is greater than T_in.

返回

A tuple containing

  • N_o – shape=[N], numbers of spike clusters of N neurons’ output spikes

  • k_positive – shape=[N], tensor of type BoolTensor, indexes. Note that k_positive can be a tensor filled with False

  • k_negative – shape=[N], tensor of type BoolTensor, indexes. Note that k_negative can be a tensor filled with False

返回类型

(Tensor, Tensor, Tensor)

A spike clustering method proposed in STCA: Spatio-Temporal Credit Assignment with Delayed Feedback in Deep Spiking Neural Networks. You can refer to the following code if this form of loss function is needed:

v_k_negative = out_v * k_negative.float().sum(dim=0)
v_k_positive = out_v * k_positive.float().sum(dim=0)
loss0 = ((N_o > N_d).float() * (v_k_negative - 1.0)).sum()
loss1 = ((N_o < N_d).float() * (1.0 - v_k_positive)).sum()
loss = loss0 + loss1
spikingjelly.clock_driven.functional.spike_similar_loss(spikes: torch.Tensor, labels: torch.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,返回度量simsim_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.clock_driven.functional.kernel_dot_product(x: torch.Tensor, y: torch.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]的内积

计算批量数据xy在核空间的内积。记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.clock_driven.functional.set_threshold_margin(output_layer: spikingjelly.clock_driven.neuron.BaseNode, label_one_hot: torch.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.clock_driven.functional.redundant_one_hot(labels: torch.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.clock_driven.functional.first_spike_index(spikes: torch.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.clock_driven.functional.multi_step_forward(x_seq: torch.Tensor, multi_step_module: torch.nn.modules.module.Module)[源代码]
参数
  • x_seq (torch.Tensor) – shape=[T, batch_size, …]

  • multi_step_module (torch.nn.Module or list or tuple) – a multi-step module, or a list/tuple that contains multi-step modules

返回

y_seq, shape=[T, batch_size, …]

返回类型

torch.Tensor

See spikingjelly.clock_driven.layer.MultiStepContainer for more details.

spikingjelly.clock_driven.functional.seq_to_ann_forward(x_seq: torch.Tensor, stateless_module: torch.nn.modules.module.Module)[源代码]
参数
  • x_seq (torch.Tensor) – shape=[T, batch_size, …]

  • multi_step_module (torch.nn.Module or list or tuple) – a stateless module, e.g., ‘torch.nn.Conv2d’ or a list contains stateless modules, e.g., ‘[torch.nn.Conv2d, torch.nn.BatchNorm2d]

返回

y_seq, shape=[T, batch_size, …]

返回类型

torch.Tensor

See spikingjelly.clock_driven.layer.SeqToANNContainer for more details.

spikingjelly.clock_driven.functional.fused_conv2d_weight_of_convbn2d(conv2d: torch.nn.modules.conv.Conv2d, bn2d: torch.nn.modules.batchnorm.BatchNorm2d)[源代码]
参数
返回

the weight of this fused module

返回类型

torch.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.clock_driven.functional.fused_conv2d_bias_of_convbn2d(conv2d: torch.nn.modules.conv.Conv2d, bn2d: torch.nn.modules.batchnorm.BatchNorm2d)[源代码]
参数
返回

the bias of this fused module

返回类型

torch.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.clock_driven.functional.scale_fused_conv2d_weight_of_convbn2d(conv2d: torch.nn.modules.conv.Conv2d, bn2d: torch.nn.modules.batchnorm.BatchNorm2d, k=None, b=None)[源代码]
参数

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d’s parameters being absorbed into Conv2d. This function sets the weight of this fused module to weight * k + b.

Note

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

spikingjelly.clock_driven.functional.scale_fused_conv2d_bias_of_convbn2d(conv2d: torch.nn.modules.conv.Conv2d, bn2d: torch.nn.modules.batchnorm.BatchNorm2d, k=None, b=None)[源代码]
参数

A {Conv2d-BatchNorm2d} can be fused to a {Conv2d} module with BatchNorm2d’s parameters being absorbed into Conv2d. This function sets the bias of this fused module to bias * k + b.

Note

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

spikingjelly.clock_driven.functional.fuse_convbn2d(conv2d: torch.nn.modules.conv.Conv2d, bn2d: torch.nn.modules.batchnorm.BatchNorm2d, k=None, b=None)[源代码]
参数
返回

the fused Conv2d layer

返回类型

torch.nn.Conv2d

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

Note

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

spikingjelly.clock_driven.functional.temporal_efficient_training_cross_entropy(x_seq: torch.Tensor, target: torch.LongTensor)[源代码]
参数
  • x_seq (torch.Tensor) – shape=[T, N, C, *], where C is the number of classes

  • target (torch.LongTensor) – shape=[N], where 0 <= target[i] <= C-1

返回

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(tet_ce_for_loop_version(x_seq, target))
print(temporal_efficient_training_cross_entropy(x_seq, target))
spikingjelly.clock_driven.functional.kaiming_normal_conv_linear_weight(net: torch.nn.modules.module.Module)[源代码]
  • API in English

参数

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

返回

None

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

  • 中文API

参数

net – Any network inherits from nn.Module

返回

None

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