Spiking RNN bug

[1y9y9l4]

There is a bug in file 'spikingjelly/clock_driven/rnn.py' line 471 and 482.
The RNN cell in PyTorch will return the 'output' and hidden state 'h_n', see:
https://pytorch.org/docs/stable/generated/torch.nn.GRU.html?highlight=gru#torch.nn.GRU

I think line 471 & 482 should be modified as:
new_states_list[0], _ = self.cells[0](x[t], states_list[0])
new_states_list[i], _ = self.cells[i](y, states_list[i])

[fangwei123456]

Thanks for you debug!

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 333 in 6364d88

def states_num():

states_num() will return the number of hidden states. For example, it returns 1 for SpikingVanillaRNN:

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 846 in 6364d88

return 1

Then its forward function also returns a tensor, rather than a tuple:

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 832 in 6364d88

return self.surrogate_function(self.linear_ih(x) + self.linear_hh(h))

So, in Line 471:

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 471 in 6364d88

new_states_list[0] = self.cells[0](x[t], states_list[0])

if self.states_num() == 1:
    new_states_list[0] = self.cells[0](x[t], states_list[0])

[fangwei123456]

When states_num() > 1, then the forward function will return a tuple:

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 701 in 6364d88

return h, c

The shape[0] of states_list is states_num():

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 418 in 6364d88

states_list = torch.zeros(size=[self.states_num(), self.num_layers, batch_size, self.hidden_size]).to(

Then the tuple is concatenated by torch.stack and converted to tensor:

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 484 in 6364d88

new_states_list[:, i] = torch.stack(self.cells[i](y, states_list[:, i]))

[1y9y9l4]

I understand what you mean. But this bug occurs in the GRU layer. Used 'pytorch==1.6.0' and 'spikingjelly==0.0.0.4', the following code reports an error:
import torch
from spikingjelly.clock_driven import rnn
rsnn = rnn.SpikingGRU(4, 64, 1)
x = torch.zeros(10, 32, 4)
y = rsnn(x)

Another thing that puzzles me is that in any case, RNN always generates LSTM Cells: (line306 & 308)

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 306 in 6364d88

cells.append(SpikingLSTMCell(self.input_size, self.hidden_size, self.bias, *args, **kwargs))

Should the correct writing be self.base_cell()(line297) instead of SpikingLSTMCell?

[fangwei123456]

Should the correct writing be self.base_cell()(line297) instead of SpikingLSTMCell?

Yes, I will fix it.

[fangwei123456]

I understand what you mean. But this bug occurs in the GRU layer. Used 'pytorch==1.6.0' and 'spikingjelly==0.0.0.4', the following code reports an error:
import torch
from spikingjelly.clock_driven import rnn
rsnn = rnn.SpikingGRU(4, 64, 1)
x = torch.zeros(10, 32, 4)
y = rsnn(x)

This problem is caused by

spikingjelly/spikingjelly/clock_driven/rnn.py

Line 306 in 6364d88

cells.append(SpikingLSTMCell(self.input_size, self.hidden_size, self.bias, *args, **kwargs))

Now the code can run correctly.

[1y9y9l4]

Thank you very much. It solved my problem perfectly.