Merge pull request #252 from NeuroBench/feature/act-sparsity-layer

Layer Wise Activation Sparsity

Author: jasonlyik

examples/nehar/benchmark.py

@@ -9,7 +9,8 @@
     ActivationSparsity,
     MembraneUpdates,
     SynapticOperations,
-    ClassificationAccuracy
+    ClassificationAccuracy,
+    ActivationSparsityByLayer,
 )
 from neurobench.metrics.static import (
     ParameterCount,
@@ -47,7 +48,7 @@
 
     # #
     static_metrics = [ParameterCount, Footprint, ConnectionSparsity]
-    workload_metrics = [ActivationSparsity, MembraneUpdates, SynapticOperations, ClassificationAccuracy]
+    workload_metrics = [ActivationSparsity, ActivationSparsityByLayer,MembraneUpdates, SynapticOperations, ClassificationAccuracy]
     # #
     benchmark = Benchmark(
         model, test_set_loader, [], postprocessors, [static_metrics, workload_metrics]

neurobench/hooks/neuron.py

@@ -10,7 +10,7 @@ class NeuronHook(ABC):
 
     """
 
-    def __init__(self, layer):
+    def __init__(self, layer, name=None):
         """
         Initializes the class.
 
@@ -24,6 +24,7 @@ def __init__(self, layer):
         self.activation_inputs = []
         self.pre_fire_mem_potential = []
         self.post_fire_mem_potential = []
+        self.name = name
         if layer is not None:
             self.hook = layer.register_forward_hook(self.hook_fn)
             self.hook_pre = layer.register_forward_pre_hook(self.pre_hook_fn)

neurobench/metrics/workload/__init__.py

@@ -2,12 +2,19 @@
 from .activation_sparsity import ActivationSparsity
 from .synaptic_operations import SynapticOperations
 from .classification_accuracy import ClassificationAccuracy
+from .activation_sparsity_by_layer import ActivationSparsityByLayer
 from .mse import MSE
 from .smape import SMAPE
 from .r2 import R2
 from .coco_map import CocoMap
 
-__stateless__ = ["ClassificationAccuracy", "ActivationSparsity", "MSE", "SMAPE"]
+__stateless__ = [
+    "ClassificationAccuracy",
+    "ActivationSparsity",
+    "MSE",
+    "SMAPE",
+    "ActivationSparsityByLayer",
+]
 
 __stateful__ = ["MembraneUpdates", "SynapticOperations", "R2", "CocoMap"]
 

neurobench/metrics/workload/activation_sparsity_by_layer.py

@@ -0,0 +1,64 @@
+from neurobench.metrics.abstract.workload_metric import AccumulatedMetric
+from collections import defaultdict
+import torch
+
+
+class ActivationSparsityByLayer(AccumulatedMetric):
+    """
+    Sparsity layer-wise of model activations.
+
+    Calculated as the number of zero activations over the number of activations layer by
+    layer, over all timesteps, samples in data.
+
+    """
+
+    def __init__(self):
+        """Initialize the ActivationSparsityByLayer metric."""
+        self.layer_sparsity = defaultdict(float)
+        self.layer_neuro_num = defaultdict(int)
+        self.layer_spike_num = defaultdict(int)
+
+    def reset(self):
+        """Reset the metric."""
+        self.layer_sparsity = defaultdict(float)
+        self.layer_neuro_num = defaultdict(int)
+        self.layer_spike_num = defaultdict(int)
+
+    def __call__(self, model, preds, data):
+        """
+        Compute activation sparsity layer by layer.
+
+        Args:
+            model: A NeuroBenchModel.
+            preds: A tensor of model predictions.
+            data: A tuple of data and labels.
+        Returns:
+            float: Activation sparsity
+
+        """
+
+        for hook in model.activation_hooks:
+            name = hook.name
+            if name is None:
+                continue
+            for output in hook.activation_outputs:
+                spike_num, neuro_num = torch.count_nonzero(
+                    output.dequantize() if output.is_quantized else output
+                ).item(), torch.numel(output)
+
+                self.layer_spike_num[name] += spike_num
+                self.layer_neuro_num[name] += neuro_num
+
+        return self.compute()
+
+    def compute(self):
+        """Compute the activation sparsity layer by layer."""
+        for key in self.layer_neuro_num:
+            sparsity = (
+                (self.layer_neuro_num[key] - self.layer_spike_num[key])
+                / self.layer_neuro_num[key]
+                if self.layer_neuro_num[key] != 0
+                else 0.0
+            )
+            self.layer_sparsity[key] = sparsity
+        return dict(self.layer_sparsity)

neurobench/models/neurobench_model.py

@@ -79,9 +79,9 @@ def is_activation_layer(module):
         def find_activation_layers(module):
             """Recursively find activation layers in a module."""
             layers = []
-            for child in module.children():
+            for child_name, child in module.named_children():
                 if is_activation_layer(child):
-                    layers.append(child)
+                    layers.append({"layer_name": child_name, "layer": child})
                 elif list(child.children()):  # Check for nested submodules
                     layers.extend(find_activation_layers(child))
             return layers
@@ -135,7 +135,9 @@ def register_hooks(self):
 
         # Registered activation hooks
         for layer in self.activation_layers():
-            self.activation_hooks.append(NeuronHook(layer))
+            layer_name = layer["layer_name"]
+            layer = layer["layer"]
+            self.activation_hooks.append(NeuronHook(layer, layer_name))
 
         for layer in self.connection_layers():
             self.connection_hooks.append(LayerHook(layer))

tests/models/model_list.py

@@ -1,10 +1,13 @@
 import torch.nn as nn
 import snntorch as snn
+from torch import manual_seed
 import snntorch.surrogate as surrogate
 
 beta = 0.9
+manual_seed(0)
 spike_grad = surrogate.fast_sigmoid()
 
+
 net = nn.Sequential(
     nn.Flatten(),
     nn.Linear(20, 256),

tests/test_metrics.py

@@ -8,6 +8,7 @@
     ActivationSparsity,
     SynapticOperations,
     MembraneUpdates,
+    ActivationSparsityByLayer,
 )
 from neurobench.metrics.static import (
     Footprint,
@@ -104,6 +105,7 @@ def setUp(self):
         self.activation_sparsity = ActivationSparsity()
         self.synaptic_operations = SynapticOperations()
         self.mem_updates = MembraneUpdates()
+        self.activation_sparsity_by_layer = ActivationSparsityByLayer()
 
     def test_classification_accuracy(self):
         model = SNNTorchModel(self.dummy_net)
@@ -426,9 +428,24 @@ def test_membrane_potential_updates(self):
 
         out = model(inp)
         tot_mem_updates = self.mem_updates(model, out, (inp, 0))
+        self.mem_updates.reset()
 
         self.assertEqual(tot_mem_updates, 50)
 
+    def test_activation_sparsity_by_layer(self):
+
+        inp = torch.ones(5, 10, 20)  # batch size, time steps, input size
+
+        model = SNNTorchModel(self.net_snn)
+
+        model.register_hooks()
+
+        out = model(inp)
+        act_sparsity_by_layer = self.activation_sparsity_by_layer(model, out, (inp, 0))
+        self.activation_sparsity_by_layer.reset()
+
+        self.assertEqual(act_sparsity_by_layer["1"], 0.96)
+
 
 # TODO: refactor this metric if needed
 # def test_neuron_update_metric():