Merge pull request #260 from NeuroBench/dev

Dev

Author: ben9809

.bumpversion.toml

@@ -3,7 +3,7 @@
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 regex = false
-current_version = "2.0.0"
+current_version = "2.1.0"
 ignore_missing_version = false
 search = "{current_version}"
 replace = "{new_version}"

.github/workflows/python-tests.yaml

@@ -44,30 +44,38 @@ jobs:
 
   test-code:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.10', '3.11', '3.12']
 
     steps:
       - uses: actions/checkout@v4
 
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: '3.11'
+          python-version: ${{ matrix.python-version }}
 
       - name: Install Poetry
         run: |
           python -m pip install --upgrade pip
+          pip install --upgrade setuptools
           pip install poetry
 
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
           poetry install --with dev
 
       - name: Run tests
         run: |
-          poetry run pytest --cov --cov-report=xml
+          if [ "${{ matrix.python-version }}" = "3.11" ]; then
+            poetry run pytest --cov --cov-report=xml
+          else
+            poetry run pytest
+          fi
 
       - name: Upload coverage to Codecov
+        if: matrix.python-version == '3.11'
         uses: codecov/codecov-action@v4
         with:
           directory: ./coverage/reports/

docs/conf.py

@@ -14,7 +14,7 @@
 project = "NeuroBench"
 copyright = "2024, Jason Yik, Noah Pacik-Nelson, Korneel Van Den Berghe, Benedetto Leto"
 author = "Jason Yik, Noah Pacik-Nelson, Korneel Van Den Berghe"
-release = "2.0.0"
+release = "2.1.0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

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/benchmarks/benchmark.py

@@ -16,16 +16,11 @@
 import json
 import csv
 import os
-from typing import Literal, List, Type, Optional, Dict, Any
+from typing import Literal, List, Type, Optional, Dict, Any, Callable, Tuple
 import pathlib
 import snntorch
 from torch import Tensor
-
-if snntorch.__version__ >= "0.9.0":
-    from snntorch import export_to_nir
-
 import torch
-import nir
 
 
 class Benchmark:
@@ -39,18 +34,25 @@ def __init__(
         self,
         model: NeuroBenchModel,
         dataloader: Optional[DataLoader],
-        preprocessors: Optional[List[NeuroBenchPreProcessor]],
-        postprocessors: Optional[List[NeuroBenchPostProcessor]],
+        preprocessors: Optional[
+            List[
+                NeuroBenchPreProcessor
+                | Callable[[Tuple[Tensor, Tensor]], Tuple[Tensor, Tensor]]
+            ]
+        ],
+        postprocessors: Optional[
+            List[NeuroBenchPostProcessor | Callable[[Tensor], Tensor]]
+        ],
         metric_list: List[List[Type[StaticMetric | WorkloadMetric]]],
     ):
         """
         Args:
             model: A NeuroBenchModel.
             dataloader: A PyTorch DataLoader.
-            preprocessors: A list of NeuroBenchPreProcessors.
-            postprocessors: A list of NeuroBenchPostProcessors.
-            metric_list: A list of lists of strings of metrics to run.
-                First item is static metrics, second item is data metrics.
+            preprocessors: A list of NeuroBenchPreProcessors or callable functions (e.g. lambda) with matching interfaces.
+            postprocessors: A list of NeuroBenchPostProcessors or callable functions (e.g. lambda) with matching interfaces.
+            metric_list: A list of lists of StaticMetric and WorkloadMetric classes of metrics to run.
+                First item is StaticMetrics, second item is WorkloadMetrics.
         """
 
         self.model = model
@@ -66,8 +68,13 @@ def run(
         quiet: bool = False,
         verbose: bool = False,
         dataloader: Optional[DataLoader] = None,
-        preprocessors: Optional[NeuroBenchPreProcessor] = None,
-        postprocessors: Optional[NeuroBenchPostProcessor] = None,
+        preprocessors: Optional[
+            NeuroBenchPreProcessor
+            | Callable[[Tuple[Tensor, Tensor]], Tuple[Tensor, Tensor]]
+        ] = None,
+        postprocessors: Optional[
+            NeuroBenchPostProcessor | Callable[[Tensor], Tensor]
+        ] = None,
         device: Optional[str] = None,
     ) -> Dict[str, Any]:
         """
@@ -117,10 +124,10 @@ def run(
                 batch_size = data[0].size(0)
 
                 # Preprocessing data
-                data = self.processor_manager.preprocess(data)
+                input, target = self.processor_manager.preprocess(data)
 
                 # Run model on test data
-                preds = self.model(data[0])
+                preds = self.model(input)
 
                 # Postprocessing data
                 preds = self.processor_manager.postprocess(preds)
@@ -220,8 +227,18 @@ def to_nir(self, dummy_input: Tensor, filename: str, **kwargs) -> None:
                 If the installed version of `snntorch` is less than `0.9.0`.
 
         """
+        try:
+            import nir
+        except ImportError:
+            raise ImportError(
+                "Exporting to NIR requires the `nir` package. Install it using `pip install nir`."
+            )
         if snntorch.__version__ < "0.9.0":
             raise ValueError("Exporting to NIR requires snntorch version >= 0.9.0")
+
+        if snntorch.__version__ >= "0.9.0":
+            from snntorch.export_nir import export_to_nir
+
         nir_graph = export_to_nir(self.model.__net__(), dummy_input, **kwargs)
         nir.write(filename, nir_graph)
         print(f"Model exported to {filename}")

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/abstract/__init__.py

@@ -1,4 +1,4 @@
 from .static_metric import StaticMetric
-from .workload_metric import WorkloadMetric
+from .workload_metric import WorkloadMetric, AccumulatedMetric
 
-__all__ = ["StaticMetric", "WorkloadMetric"]
+__all__ = ["StaticMetric", "WorkloadMetric", "AccumulatedMetric"]

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))