fangwei123456 / spikingjelly

SpikingJelly is an open-source deep learning framework for Spiking Neural Network (SNN) based on PyTorch.
https://spikingjelly.readthedocs.io
Other
1.35k stars 239 forks source link

using seq_to_ann_forward for internal convolutional layers (not the first conv layer which accept inputs) seems to be a problem #276

Closed gwgknudayanga closed 2 years ago

gwgknudayanga commented 2 years ago

1.) I have the following network "net". To process the input "x_seq", can i use the for loop as i written? Within that for loop i am using functional.seq_to_ann_forward to process layers which don't have a state (ex: conv layers). However internal convolutional layers(in this example second conv layer) get the the inputs from neurons in previous layer which have states. Hence for internal convolutional layers we cannot parallize the processing over time by using 'seq_to_ann_forward', (that means we cannot do flatterning like 'T*N' for internal conv layers) and hence we should avoid it. Isn't it? Because the inputs for the second convolutional layer are ready only when "neuron 1" processed the inputs for time step 1 and time step 2 in a sequential manner. Can you please explain how should i proceed? thanks.

2.) Also If we call functional.ann_to_seq_forward(neuron.IFNode(step_mode='m')), in calculating outputs for each time step will the neuron get the states of from previous time step?

torch.nn as nn from spikingjelly.activation_based import neuron

net = nn.Sequential( nn.Conv2d(3, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.Conv2d(8, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.MaxPool2d(2, 2), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.Flatten(start_dim=1), nn.Linear(8 H // 2 W // 2, 10), neuron.IFNode(step_mode = 'm',backend='cupy'), )

net = net.to('cuda:0')

T = 4 N = 1 C = 3 H = 8 W = 8 x_seq = torch.rand([T, N, C, H, W]) x_seq = x_seq.to('cuda:0')

out = x_seq for i in range(len(net)): m = net[i] if isinstance(m, neuron.BaseNode): out = m(out) else: out = functional.seq_to_ann_forward(out, m)

fangwei123456 commented 2 years ago

I suggest that you refer to the tutorials:

https://spikingjelly.readthedocs.io/zh_CN/latest/activation_based_en/basic_concept.html

https://spikingjelly.readthedocs.io/zh_CN/latest/activation_based_en/container.html

To process the input "x_seq", can i use the for loop as i written?

Yes. Note that all modules inside the for loop must be stateless or step_mode = 's'.

You can also use multi_step_forward (function or container).

Hence for internal convolutional layers we cannot parallize the processing over time by using 'seq_to_ann_forward', (that means we cannot do flatterning like 'T*N' for internal conv layers) and hence we should avoid it. Isn't it?

Yes. Hence, the best way is to run the network in multi-step mode, then stateless layers can receive data with shape [T, N, *]. Then we can use seq_to_ann_forward to accelerate.

for each time step will the neuron get the states of from previous time step

Yes.

gwgknudayanga commented 2 years ago

I suggest that you refer to the tutorials: just referred. But still i have some doubts. https://spikingjelly.readthedocs.io/zh_CN/latest/activation_based_en/basic_concept.html

https://spikingjelly.readthedocs.io/zh_CN/latest/activation_based_en/container.html

To process the input "x_seq", can i use the for loop as i written?

Yes. Note that all modules inside the for loop must be stateless or step_mode = 's'. But, at there i have used for loop to iterate over the layers which are within the nn.Sequential block. I am not iterating over the time steps.So should i need to change step_mode of the neurons from 'm' to 's' ? . It is bit confusing. You can also use multi_step_forward (function or container). Here i am using torch.nn.conv2d instead of the conv layers available in spikejelly APIs.

Hence for internal convolutional layers we cannot parallize the processing over time by using 'seq_to_ann_forward', (that means we cannot do flatterning like 'T*N' for internal conv layers) and hence we should avoid it. Isn't it?

Yes. Hence, the best way is to run the network in multi-step mode, then stateless layers can receive data with shape [T, N, *]. Then we can use seq_to_ann_forward to accelerate. Once we set the mult_step_mode for the nn.Sequential block which is named as 'net' in my previous example, i have to change the step_mode of the neurons inside the nn.Sequential block to 's' instead of 'm', isn't it? -So once we set the multi_step_mode, we can call net(x_seq) to run it step by step. Then how we can call seq_to_ann_forward only for stateless layers selectively ? -Because if we call seq_to_ann_forward for the internal IFNeurons, then they will not save the states in between time steps, isn't it?

for each time step will the neuron get the states of from previous time step

Yes.

so in short how should i process the following network which contain stateless modules as well as multi step IFNeurons for the for the mentioned x_seq input which has four time steps so that it will behave as multi step state aware fashion while accelerating with seq_to_ann_forward? net = nn.Sequential( nn.Conv2d(3, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.Conv2d(8, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.MaxPool2d(2, 2), neuron.IFNode(step_mode = 'm',backend='cupy'), nn.Flatten(start_dim=1), nn.Linear(8 H // 2 W // 2, 10), neuron.IFNode(step_mode = 'm',backend='cupy'), )

Thank you !

fangwei123456 commented 2 years ago

You can use from spikingjelly.activation_based.layer import SeqToANNContainer to wrap nn.Conv2d and nn.MaxPool2d: SeqToANNContainer(nn.Conv2d(3, 8, kernel_size=3, padding=1, stride=1, bias=False))

You can also use layer.* and set the whole network to multi-step mode:

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

H = 32
W = 32

net = nn.Sequential(
    layer.Conv2d(3, 8, kernel_size=3, padding=1, stride=1, bias=False),
    neuron.IFNode(step_mode='m', backend='cupy'),
    layer.Conv2d(8, 8, kernel_size=3, padding=1, stride=1, bias=False),
    neuron.IFNode(step_mode='m', backend='cupy'),
    layer.MaxPool2d(2, 2),
    neuron.IFNode(step_mode='m', backend='cupy'),
    layer.Flatten(start_dim=1),
    layer.Linear(8 * H // 2 * W // 2, 10),
    neuron.IFNode(step_mode='m', backend='cupy'),
)

functional.set_step_mode(net, 'm')

I recommend to use the second method.

gwgknudayanga commented 2 years ago

** Thanks fang. So in the second method ,Since we set the whole network to multi_step_mode (functional.set_step_mode(net,'m'), is it ok to avoid setting the multi_step mode when declaring the neurons within the Network(neuron.IFNode) ? In the code you mentioned in previous comment i can see in adding to setting multi_mode to whole network, the neurons are also defined with step_mode='m'. That is neuron.IFNode(step_mode='m', backend='cupy').

On the other hand following code(iterate over time steps, but step_mode for neurons is 's') will give the equivalent results as the results from the second method, isn't it? Is it still fine to set the step_mode of neurons to multi_mode even when we iterating through time steps? And one more small thing, though i set the device for input(x_seq) and the networkwork(net) to gpu i cannot see the code is running in gpu when iterating over time steps. Do you have any suggestion? thank you.

x_seq = torch.rand([T, N, C, H, W]) x_seq = x_seq.to('cuda:0')

net = nn.Sequential( layer.Conv2d(3, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode='m', backend='cupy'), layer.Conv2d(8, 8, kernel_size=3, padding=1, stride=1, bias=False), neuron.IFNode(step_mode='m', backend='cupy'), layer.MaxPool2d(2, 2), neuron.IFNode(step_mode='m', backend='cupy'), layer.Flatten(start_dim=1), layer.Linear(8 H // 2 W // 2, 10), neuron.IFNode(step_mode='m', backend='cupy'), ) net.to('cuda:0')

y_seq_step_by_step = [] for t in range(T): x = x_seq[t] y = net(x) y_seq_step_by_step.append(y.unsqueeze(0))

fangwei123456 commented 2 years ago

is it ok to avoid setting the multi_step mode when declaring the neurons within the Network(neuron.IFNode) ?

Yes.

On the other hand following code(iterate over time steps, but step_mode for neurons is 's') will give the equivalent results as the results from the second method, isn't it?

No, if you use for loop manually, you should not set neuron in step_mode = 'm'.

though i set the device for input(x_seq) and the networkwork(net) to gpu i cannot see the code is running in gpu when iterating over time steps

You can print y.device, which should be cuda:0.

gwgknudayanga commented 2 years ago

Thanks fang.