velkiSNN.py

batch_size = 16
num_classes = 4
x = 50
num_steps = x
beta = 0.3
spike_grad = surrogate.fast_sigmoid(slope=10)
learning_rate = 1e-3
n_epochs = x

class MyEEGSNNModel(nn.Module):
    def __init__(self, n_outputs):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 12, 5)
        self.lif1 = snn.Leaky(beta=beta, spike_grad=spike_grad)
        self.conv2 = nn.Conv2d(12, 64, 5)
        self.lif2 = snn.Leaky(beta=beta, spike_grad=spike_grad)

        self.fc1 = nn.Linear(64 * 29 * 29, n_outputs)
        self.lif3 = snn.Leaky(beta=beta, spike_grad=spike_grad)

    def forward(self, x):
        mem1 = self.lif1.init_leaky()
        mem2 = self.lif2.init_leaky()
        mem3 = self.lif3.init_leaky()

        cur1 = F.max_pool2d(self.conv1(x), kernel_size=(2, 2))
        spk1, mem1 = self.lif1(cur1, mem1)

        cur2 = F.max_pool2d(self.conv2(spk1), kernel_size=(2, 2))
        spk2, mem2 = self.lif2(cur2, mem2)

        # Flatten
        batch_size_curr = x.shape[0]
        cur3 = self.fc1(spk2.view(batch_size_curr, -1))
        spk3, mem3 = self.lif3(cur3, mem3)

        return spk3, mem3