Sparsity for s monitor

Author: n-shevko

bindsnet/evaluation/evaluation.py

@@ -44,8 +44,9 @@ def assign_labels(
             indices = torch.nonzero(labels == i).view(-1)
 
             # Compute average firing rates for this label.
+            selected_spikes = torch.index_select(spikes, dim=0, index=torch.tensor(indices))
             rates[:, i] = alpha * rates[:, i] + (
-                torch.sum(spikes[indices], 0) / n_labeled
+                torch.sum(selected_spikes, 0) / n_labeled
             )
 
     # Compute proportions of spike activity per class.
@@ -111,6 +112,8 @@ def all_activity(
 
     # Sum over time dimension (spike ordering doesn't matter).
     spikes = spikes.sum(1)
+    if spikes.is_sparse:
+        spikes = spikes.to_dense()
 
     rates = torch.zeros((n_samples, n_labels), device=spikes.device)
     for i in range(n_labels):
@@ -152,6 +155,8 @@ def proportion_weighting(
 
     # Sum over time dimension (spike ordering doesn't matter).
     spikes = spikes.sum(1)
+    if spikes.is_sparse:
+        spikes = spikes.to_dense()
 
     rates = torch.zeros((n_samples, n_labels), device=spikes.device)
     for i in range(n_labels):

bindsnet/network/monitors.py

@@ -45,6 +45,7 @@ def __init__(
         time: Optional[int] = None,
         batch_size: int = 1,
         device: str = "cpu",
+        sparse: Optional[bool] = False
     ):
         # language=rst
         """
@@ -62,6 +63,7 @@ def __init__(
         self.time = time
         self.batch_size = batch_size
         self.device = device
+        self.sparse = sparse
 
         # if time is not specified the monitor variable accumulate the logs
         if self.time is None:
@@ -98,11 +100,12 @@ def record(self) -> None:
         for v in self.state_vars:
             data = getattr(self.obj, v).unsqueeze(0)
             # self.recording[v].append(data.detach().clone().to(self.device))
-            self.recording[v].append(
-                torch.empty_like(data, device=self.device, requires_grad=False).copy_(
-                    data, non_blocking=True
-                )
+            record = torch.empty_like(data, device=self.device, requires_grad=False).copy_(
+                data, non_blocking=True
             )
+            if self.sparse:
+                record = record.to_sparse()
+            self.recording[v].append(record)
             # remove the oldest element (first in the list)
             if self.time is not None:
                 self.recording[v].pop(0)

examples/mnist/batch_eth_mnist.py

@@ -147,7 +147,7 @@
 spikes = {}
 for layer in set(network.layers):
     spikes[layer] = Monitor(
-        network.layers[layer], state_vars=["s"], time=int(time / dt), device=device
+        network.layers[layer], state_vars=["s"], time=int(time / dt), device=device, sparse=True
     )
     network.add_monitor(spikes[layer], name="%s_spikes" % layer)
 
@@ -165,7 +165,8 @@
 perf_ax = None
 voltage_axes, voltage_ims = None, None
 
-spike_record = torch.zeros((update_interval, int(time / dt), n_neurons), device=device)
+spike_record = [torch.zeros((batch_size, int(time / dt), n_neurons), device=device).to_sparse() for _ in range(update_interval // batch_size)]
+spike_record_idx = 0
 
 # Train the network.
 print("\nBegin training...")
@@ -197,12 +198,13 @@
             # Convert the array of labels into a tensor
             label_tensor = torch.tensor(labels, device=device)
 
+            spike_record_tensor = torch.cat(spike_record, dim=0)
             # Get network predictions.
             all_activity_pred = all_activity(
-                spikes=spike_record, assignments=assignments, n_labels=n_classes
+                spikes=spike_record_tensor, assignments=assignments, n_labels=n_classes
             )
             proportion_pred = proportion_weighting(
-                spikes=spike_record,
+                spikes=spike_record_tensor,
                 assignments=assignments,
                 proportions=proportions,
                 n_labels=n_classes,
@@ -240,7 +242,7 @@
 
             # Assign labels to excitatory layer neurons.
             assignments, proportions, rates = assign_labels(
-                spikes=spike_record,
+                spikes=spike_record_tensor,
                 labels=label_tensor,
                 n_labels=n_classes,
                 rates=rates,
@@ -261,11 +263,10 @@
 
         # Add to spikes recording.
         s = spikes["Ae"].get("s").permute((1, 0, 2))
-        spike_record[
-            (step * batch_size)
-            % update_interval : (step * batch_size % update_interval)
-            + s.size(0)
-        ] = s
+        spike_record[spike_record_idx] = s
+        spike_record_idx += 1
+        if spike_record_idx == len(spike_record):
+            spike_record_idx = 0
 
         # Get voltage recording.
         exc_voltages = exc_voltage_monitor.get("v")