spikingjelly.clock_driven.functional package
Module contents
- spikingjelly.clock_driven.functional.reset_net(net: 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 theirreset()
function if exists.
- spikingjelly.clock_driven.functional.spike_cluster(v: 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 thanT_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: Tensor, labels: Tensor, kernel_type='linear', loss_type='mse', *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 iflabels[i]
andlabels[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 ofsim_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
andsim_p
according toloss_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: 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 ofx[i]
andy[j]
.
Calculate inner product of
x
andy
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: BaseNode, label_one_hot: Tensor, eval_threshold=1.0, threshold0=0.9, threshold1=1.1)[源代码]
-
- 参数
- 返回
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 tothreshold0
.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: Tensor, 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.]])
- 参数
- 返回
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’se.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: 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 ofTrue
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: Tensor, single_step_module: Module)[源代码]
- 参数
x_seq (Tensor) – shape=[T, batch_size, …]
single_step_module (torch.nn.Module or list or tuple or torch.nn.Sequential) – a single-step module, or a list/tuple that contains single-step modules
- 返回
y_seq, shape=[T, batch_size, …]
- 返回类型
Tensor
See
spikingjelly.clock_driven.layer.MultiStepContainer
for more details.
- spikingjelly.clock_driven.functional.seq_to_ann_forward(x_seq: Tensor, stateless_module: Module)[源代码]
- 参数
x_seq (Tensor) – shape=[T, batch_size, …]
stateless_module (torch.nn.Module or list or tuple or torch.nn.Sequential) – 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, …]
- 返回类型
Tensor
See
spikingjelly.clock_driven.layer.SeqToANNContainer
for more details.
- spikingjelly.clock_driven.functional.fused_conv2d_weight_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d)[源代码]
- 参数
conv2d (torch.nn.Conv2d) – a Conv2d layer
bn2d (torch.nn.BatchNorm2d) – a BatchNorm2d layer
- 返回
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.clock_driven.functional.fused_conv2d_bias_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d)[源代码]
- 参数
conv2d (torch.nn.Conv2d) – a Conv2d layer
bn2d (torch.nn.BatchNorm2d) – a BatchNorm2d layer
- 返回
the bias 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.clock_driven.functional.scale_fused_conv2d_weight_of_convbn2d(conv2d: Conv2d, bn2d: BatchNorm2d, k=None, b=None)[源代码]
- 参数
conv2d (torch.nn.Conv2d) – a Conv2d layer
bn2d (torch.nn.BatchNorm2d) – a BatchNorm2d layer
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: Conv2d, bn2d: BatchNorm2d, k=None, b=None)[源代码]
- 参数
conv2d (torch.nn.Conv2d) – a Conv2d layer
bn2d (torch.nn.BatchNorm2d) – a BatchNorm2d layer
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: Conv2d, bn2d: BatchNorm2d, k=None, b=None)[源代码]
- 参数
conv2d (torch.nn.Conv2d) – a Conv2d layer
bn2d (torch.nn.BatchNorm2d) – a BatchNorm2d layer
- 返回
the fused Conv2d layer
- 返回类型
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: Tensor, target: LongTensor)[源代码]
- 参数
x_seq (Tensor) –
shape=[T, N, C, *]
, whereC
is the number of classestarget (torch.LongTensor) –
shape=[N]
, where0 <= target[i] <= C-1
- 返回
the temporal efficient training cross entropy
- 返回类型
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: 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))
Tip
The TET cross entropy is proposed by Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting.
- spikingjelly.clock_driven.functional.kaiming_normal_conv_linear_weight(net: 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
andtorch.nn.Linear
in net by the kaiming normal. Seetorch.nn.init.kaiming_normal_
for more details.