Neurons in 'Temporal Backpropagation for Spiking Neural Networks with One Spike per Neuron'

[fangwei123456]

We can discuss how to implement the spiking neuron with most with one spike.

[fangwei123456]

I think the neuron is similar with that in 'Temporal Backpropagation for Spiking Neural Networks with One Spike per Neuron'. Here is my codes:

from spikingjelly.clock_driven import neuron, surrogate
import torch
import torch.nn as nn
class OneSpikeIFNode(nn.Module):
    def __init__(self, v_threshold=1.0, v_reset=0.0, surrogate_function=surrogate.Sigmoid()):
        super().__init__()
        self.v_threshold = v_threshold
        self.v_reset = v_reset
        self.surrogate_function = surrogate_function
        self.spike = None
        self.fire_mask = 0
        if self.v_reset is None:
            self.v = 0.0
        else:
            self.v = self.v_reset

    def neuronal_reset(self):
        if self.v_reset is None:
            self.v = self.v - self.spike * self.v_threshold
        else:
            self.v = (1 - self.spike) * self.v + self.spike * self.v_reset

    def neuronal_charge(self, dv: torch.Tensor):
        self.v += dv

    def neuronal_fire(self):
        self.spike = self.surrogate_function(self.v - self.v_threshold) * (1. - self.fire_mask)
        self.fire_mask += self.spike


    def forward(self, dv: torch.Tensor):
        self.neuronal_charge(dv)
        self.neuronal_fire()
        self.neuronal_reset()
        return self.spike

    def reset(self):
        if self.v_reset is None:
            self.v = 0.0
        else:
            self.v = self.v_reset

        self.spike = None
        self.fire_mask = 0

[tmasquelier]

Yes this code looks good! I notice that after the first spike you block subsequent firings, but not the integration (v will continue to accumulate). With Ali we also tried to block the integration, but I can't remember which option worked best (in s2net).

[ali-rasteh]

@tmasquelier : Yes at first we blocked the integration on V but then we reallized that it wasn't useful and finally changed the code and removed that part!
So I think @fangwei123456 's code is exactly similar to what we developed on s2net.
But don't you think that we can make it dynamic? (maximum number of spikes per neuron could be a value except 1)

[fangwei123456]

'maximum number of spikes per neuron could be a value except 1'
Yes, it can be done by
self.spike = self.surrogate_function(self.v - self.v_threshold) * self.surrogate_function(max_spike_num - self.fire_mask)

The surrogate function will be used to create binary block.

[fangwei123456]

I think the surrogate function for binary block should be different to that for spiking.

[tmasquelier]

Indeed! But for me, the most interesting case is max_spike_num = 1

[fangwei123456]

We can also discuss how to implement the latency based max pooling. But I don't know much about it. Can you describe the latency based max pooling?

[tmasquelier]

The concept is to propagate only the first spike in a spatial neighborhood. But of course, this needs to be done using Tensor operations, so that PyTorch can do its magic ;-) Ali found a nice way to do it, I'll let him explain. This reminds me that we should also check whether max pooling is really better than using stride >1 in the preceding conv layer.

[ali-rasteh]

Sorry for being late
This is the way I have implemented it :
Suppose you have a tensor x with these dimensions : (batch_size, nb_channels, nb_steps, w, h). nb_channels is the number of convolutional channels, nb_steps is the number of time steps of simulation, w and h are the 2D dimensions of the convolutional layer.

batch_size = x.shape[0]
nb_steps = x.shape[2]

Cumulative sum on x in the temporal dimension will generate 1 from the first spike time till the end
x_cumsum = x.cumsum(dim=2)
Summing on the temporal dimension will give the maximum number for the neuron with first spike
x_sum = x_cumsum.sum(dim=2, keepdim=False)
Max pooling on the given tensor will give indices for neurons with the first spike times
x_pool_2D, indices_2D = torch.nn.functional.max_pool2d(x_sum, kernel_size, stride, padding=0, dilation, return_indices=True, ceil_mode=False)

But we need to separate neurons with first spike time from x, given the indices_2D from last part and these lines give us what we want :

flattened_tensor = x.flatten(start_dim=3)
indices_3D=indices_2D.reshape((batch_size, nb_channels,1,indices_2D.shape[2],indices_2D.shape[3]))
indices_3D=indices_3D.repeat(1,1,nb_steps,1,1)
spk_rec = flattened_tensor.gather(dim=3, index=indices_3D.flatten(start_dim=3)).view_as(indices_3D)

[tmasquelier]

Hi Ali,
Thanks for the explanation.
I'm not sure I follow well the last four lines, but something seems weird. The output of the layer should have W_out ~ W_in / stride and H_out ~ H_in / stride.
Here I have the filling that W and H are unchanged.
Couldn't you instantiate a tensor full of zeros with size (batch_size, nb_channels, nb_steps, w/stride, h/stride), and then place the 1s at the positions specified by x_pool_2D using torch.Tensor.scatter_ ?

[ali-rasteh]

I have tested the code and it works correctly. In fact the gather function is changing the w,h dimensions.
You can do it with torch.Tensor.scatter_ as well but I haven't tested it.

[tmasquelier]

Great, thanks!

[fangwei123456]

Thanks for Ali's codes about pooling. I also implement a pooling layer. It is a single-step implement, and I will try to implement a multi-step version.

import torch
import torch.nn as nn
import torch.nn.functional as F

class FirstSpikePool2d(nn.MaxPool2d):
    def __init__(self, kernel_size, stride,
                 padding=0, dilation=1,
                 return_indices: bool = False, ceil_mode: bool = False) -> None:
        super().__init__(kernel_size, stride, padding, dilation, return_indices, ceil_mode)
        self.fire_mask = 0.

    def forward(self, spike: torch.Tensor):
        out_spike = super().forward(spike)
        with torch.no_grad():
            out_spike *= (1. - self.fire_mask)
            self.fire_mask += out_spike
        return out_spike

    def reset(self):
        self.fire_mask = 0.

Running example:

T = 4
spike = (torch.rand([T, 1, 1, 4, 4]) > 0.9).float()
fsp = FirstSpikePool2d(2, 2)

for t in range(T):
    print(t, 'in spike\n', spike[t], 'out spike\n', fsp(spike[t]))

Outputs:

0 in spike
 tensor([[[[0., 0., 1., 0.],
          [0., 0., 0., 1.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]]]]) out spike
 tensor([[[[0., 1.],
          [0., 0.]]]])
1 in spike
 tensor([[[[0., 0., 0., 1.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 1.]]]]) out spike
 tensor([[[[0., 0.],
          [0., 1.]]]])
2 in spike
 tensor([[[[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 1., 0.],
          [0., 0., 0., 0.]]]]) out spike
 tensor([[[[0., 0.],
          [0., 0.]]]])
3 in spike
 tensor([[[[0., 0., 0., 0.],
          [0., 0., 0., 0.],
          [0., 0., 0., 1.],
          [0., 0., 0., 0.]]]]) out spike
 tensor([[[[0., 0.],
          [0., 0.]]]])

[tmasquelier]

I see, you remove subsequent spikes with a mask. Good idea!

[fangwei123456]

A simple way to implement the multi-step version is using MultiStepContainer

[fangwei123456]

N-D max pooling:

import torch
import torch.nn as nn
import torch.nn.functional as F
from abc import abstractmethod

class FirstSpikePoolNd(nn.modules.pooling._MaxPoolNd):
    def __init__(self, kernel_size, stride,
                 padding=0, dilation=1,
                 return_indices: bool = False, ceil_mode: bool = False) -> None:
        super().__init__(kernel_size, stride, padding, dilation, return_indices, ceil_mode)
        self.fire_mask = 0.

    @abstractmethod
    def max_pool(self, *args, **kwargs):
        pass

    def forward(self, spike: torch.Tensor):
        out_spike = self.max_pool(spike, self.kernel_size, self.stride,
                            self.padding, self.dilation, self.ceil_mode,
                            self.return_indices)
        with torch.no_grad():
            out_spike *= (1. - self.fire_mask)
            self.fire_mask += out_spike
        return out_spike

    def reset(self):
        self.fire_mask = 0.


class FirstSpikePool1d(FirstSpikePoolNd):
    def max_pool(self, *args, **kwargs):
        return F.max_pool1d(*args, **kwargs)

class FirstSpikePool2d(FirstSpikePoolNd):
    def max_pool(self, *args, **kwargs):
        return F.max_pool2d(*args, **kwargs)

class FirstSpikePool3d(FirstSpikePoolNd):
    def max_pool(self, *args, **kwargs):
        return F.max_pool3d(*args, **kwargs)

[fangwei123456]

Doing some ablation experiments will be a good choice. @ali-rasteh, do the new codes work well as your origin codes?

[ali-rasteh]

@fangwei123456 Thanks for your suggestions and codes, Your codes for at most one spike per neuron works well and I got similar results with s2net! About the latency based max pooling I'm working on it and I tell you tomorrow.

[ali-rasteh]

Hi guys
@fangwei123456 : I tested both methods of maxpool with networks of same architecture and trained both for 100 epochs (with 20 time steps) and seems when 'out_spike *= (1. - self.fire_mask)' is within 'torch.no_grad()' the performance is better! (0.9196 to 0.9144). But I think we can not account on these results yet because we need to have a scheduler on training. I am adding a scheduler and trying again. But the good news is that latency based maxpool works! @tmasquelier

[fangwei123456]

Thanks for Ali's experiments! I recommend to use CosineAnnealingLR.

[ali-rasteh]

Sure! I'll start with CosineAnnealingLR.