Improve assign_labels accuracy and performance

Resolved a bug where silent neurons defaulted to the first class label. Optimized indexing to prevent unnecessary host-to-device transfers and added robust NaN handling for firing rate proportions.

Author: Machuka

bindsnet/evaluation/evaluation.py

@@ -12,51 +12,36 @@ def assign_labels(
     rates: Optional[torch.Tensor] = None,
     alpha: float = 1.0,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    # language=rst
-    """
-    Assign labels to the neurons based on highest average spiking activity.
 
-    :param spikes: Binary tensor of shape ``(n_samples, time, n_neurons)`` of a single
-        layer's spiking activity.
-    :param labels: Vector of shape ``(n_samples,)`` with data labels corresponding to
-        spiking activity.
-    :param n_labels: The number of target labels in the data.
-    :param rates: If passed, these represent spike rates from a previous
-        ``assign_labels()`` call.
-    :param alpha: Rate of decay of label assignments.
-    :return: Tuple of class assignments, per-class spike proportions, and per-class
-        firing rates.
-    """
     n_neurons = spikes.size(2)
+    device = spikes.device # Keep everything on the same device.
 
     if rates is None:
-        rates = torch.zeros((n_neurons, n_labels), device=spikes.device)
+        rates = torch.zeros((n_neurons, n_labels), device=device)
 
     # Sum over time dimension (spike ordering doesn't matter).
-    spikes = spikes.sum(1)
-
+    summed_spikes = spikes.sum(1)
+    
     for i in range(n_labels):
         # Count the number of samples with this label.
-        n_labeled = torch.sum(labels == i).float()
+        mask = (labels == i)
+        n_labeled = mask.sum().float()
 
         if n_labeled > 0:
-            # Get indices of samples with this label.
-            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(selected_spikes, 0) / n_labeled
-            )
+            # Get indices of samples with this label (masking is faster and stays on the GPU).
+            label_sum = summed_spikes[mask].sum(0)
+            # Update rates.
+            rates[:, i] = alpha * rates[:, i] + (label_sum / n_labeled)
 
-    # Compute proportions of spike activity per class.
-    proportions = rates / rates.sum(1, keepdim=True)
-    proportions[proportions != proportions] = 0  # Set NaNs to 0
+    # 3. Compute proportions (and use 'torch.where' to avoid NaN bug).
+    total_activity = rates.sum(1, keepdim=True)
+    proportions = torch.where(total_activity > 0, rates / total_activity, torch.zeros_like(rates))
 
     # Neuron assignments are the labels they fire most for.
-    assignments = torch.max(proportions, 1)[1]
+    max_vals, assignments = torch.max(proportions, 1)
+    
+    # Set unassigned (silent) neurons to -1 instead of defaulting to 0.
+    assignments[max_vals == 0] = -1 
 
     return assignments, proportions, rates