Is it possible to do prediciton every time step while using Surrogate Gradient?

[Dola47]

Hi,

I just checked that the spike_grad parameter needs to be passed manually for each neuron. And that the BPTT always requires hidden_init=True which means that it can be used for real-time inference. I wonder if the same can be done with Surrogate Gradient. And if yes can you provide a simple example of implementing that.

Thanks.

[jeshraghian]

Yes, this should be possible.

Constructing your network is done separately to the BPTT method, so you would add a surrogate gradient in the exact same way as outlined in the documentation: https://snntorch.readthedocs.io/en/latest/snntorch.surrogate.html

I haven't tested it in a long time, and writing up a proper tutorial for this is on my to-do list after a few more important features.

But something like this should work.
I don't remember what the forward function is supposed to return (i.e., if it should be every layer output, spk1, mem1, spk2, mem2 etc.) or if it's only the final layer (i.e., spk2, mem2). But if you figure that out, let me know and I can update the example in the docs!

from snntorch import surrogate
import torch
import torch.nn as nn

beta = 0.9

# Initialize surrogate gradient
spike_grad = surrogate.fast_sigmoid()  # passes default parameters from a closure

# Define Network
class Net(nn.Module):
 def __init__(self):
     super().__init__()

 # Initialize layers, specify the ``spike_grad`` argument
     self.fc1 = nn.Linear(num_inputs, num_hidden)
     self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)
     self.fc2 = nn.Linear(num_hidden, num_outputs)
     self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)

 def forward(self, x, syn1, mem1, spk1, syn2, mem2):
     cur1 = self.fc1(x)
     spk1, mem1 = self.lif1(cur1, mem1)
     cur2 = self.fc2(spk1)
     spk2, mem2 = self.lif2(cur2, mem2)
     return mem1, spk1, mem2, spk2

 net = Net().to(device)

optimizer = torch.optim.Adam(net.parameters(), lr=lr, betas=betas)
criterion = nn.CrossEntropyLoss()

# Time-variant input data
for input, target in dataset:
   loss = BPTT(net, input, target, num_steps, batch_size, optimizer, criterion)

# Time-invariant input data
for input, targets in dataset:
   loss = BPTT(net, input, target, num_steps, batch_size, optimizer, criterion, time_varying_data=False)