CIFAR10 Norse sim updated

Norse/CIFAR-10/CIFAR10_conv_bench.py

@@ -84,21 +84,20 @@ for maxTh, Alpha, reinforcement, compression, layerwise in np.array(list(itertoo
             def __init__(self, num_channels=3, feature_size=32, method="super", alpha=100):
                 super(ConvNet, self).__init__()
 
-                self.features = int(((feature_size - 4) / 2 - 4) / 2)
-                self.conv1_out_channels = 32
-                self.conv2_out_channels = 128
+                self.conv1_out_channels = 128
+                self.conv2_out_channels = 256
                 self.fc1_out_channels = 1024
                 self.out_channels = 10
                 self.conv1 = torch.nn.Conv2d(num_channels, self.conv1_out_channels, 5, 1, bias=False)
                 self.conv2 = torch.nn.Conv2d(self.conv1_out_channels, self.conv2_out_channels, 5, 1, bias=False)
-                self.fc1 = torch.nn.Linear(self.features**2  * self.conv2_out_channels, self.fc1_out_channels, bias=False)
-                self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
-                self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
-                self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.25))
-                self.out = LILinearCell(self.fc1_out_channels, self.out_channels)
+                self.fc1 = torch.nn.Linear(256*5*5, self.fc1_out_channels, bias=False)
+                self.fc2 = torch.nn.Linear(self.fc1_out_channels, 1000, bias=False)
+                self.lif0 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.3))
+                self.lif1 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.3))
+                self.lif2 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.3))
+                self.lif3 = LIFCell(p=LIFParameters(method=method, alpha=alpha, v_th=0.3))
+                self.out = LILinearCell(1000, self.out_channels,p=LIFParameters(method=method, alpha=alpha, v_th=0.3))
                 
-                #LILinearCell(self.fc1_out_channels, self.out_channels)
-
             def forward(self, x):
                 seq_length = x.shape[0]
                 batch_size = x.shape[1]
@@ -111,16 +110,22 @@ for maxTh, Alpha, reinforcement, compression, layerwise in np.array(list(itertoo
                 )
 
                 for ts in range(seq_length):
-                    z = self.conv1(x[ts, :])
+                    z = torch.nn.functional.relu(self.conv1(x[ts, :]))
+
                     z, s0 = self.lif0(z, s0)
                     z = torch.nn.functional.max_pool2d(z, 2, 2)
-                    z = self.out_channels * self.conv2(z)
+                    z = 3 * torch.nn.functional.relu(self.conv2(z))
+
                     z, s1 = self.lif1(z, s1)
                     z = torch.nn.functional.max_pool2d(z, 2, 2)
-                    z = z.view(-1, self.features**2 * self.conv2_out_channels)
+
+                    z = z.view(-1, 256*5*5)
                     z = self.fc1(z)
                     z, s2 = self.lif2(z, s2)
-                    v, so = self.out(torch.nn.functional.relu(z), so)
+                    z = self.fc2(z)
+                    z, s3 = self.lif3(z, s3)
+                    v, so = self.out(z, so)
+
                     voltages[ts, :, :] = v
 
                 return voltages
@@ -171,7 +176,7 @@ for maxTh, Alpha, reinforcement, compression, layerwise in np.array(list(itertoo
         torch.autograd.set_detect_anomaly(True)
 
         T = 35
-        LR = 0.001
+        LR = 3e-5
         EPOCHS = 100  # Increase this for improved accuracy
 
         if torch.cuda.is_available():
@@ -180,7 +185,7 @@ for maxTh, Alpha, reinforcement, compression, layerwise in np.array(list(itertoo
             DEVICE = torch.device("cpu")
 
         model = Model(
-            encoder=encode.SpikeLatencyLIFEncoder(T), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
+            encoder=encode.SpikeLatencyLIFEncoder(T,p=LIFParameters(v_th=0.3)), snn=ConvNet(alpha=80), decoder=decode).to(DEVICE)
 
         optimizer = torch.optim.Adam(model.parameters(), lr=LR)