Merge pull request #266 from NeuroBench/feat/macs_neuron

Feat: Neuron Operations

Author: jasonlyik

docs/metrics/workload_metrics/index.rst

@@ -13,4 +13,5 @@ Workload Metrics
    smape
    r2
    mse
+   neuron_operations
 

docs/metrics/workload_metrics/neuron_operations.rst

@@ -0,0 +1,85 @@
+Neuron Operations Metric
+========================
+
+The `NeuronOperations` metric is designed to measure the computational workload associated with neuron activity in spiking neural networks (SNNs). It tracks the number of operations required to update the membrane potential of neurons during the forward pass of a model. These operations include the reset mechanisms defined in the `neuron_ops_reset_operations` dictionary, such as "subtract" and "zero".
+
+Purpose
+-------
+The metric provides insights into the computational cost of neuron updates, which is critical for analyzing and optimizing the efficiency of spiking neural networks. By understanding the workload associated with different neuron types and reset mechanisms, researchers can identify bottlenecks and improve model performance.
+
+Reset Mechanisms
+----------------
+Each neuron type has specific reset mechanisms that determine how the membrane potential is updated after a spike. The two main reset mechanisms are:
+
+1. **Subtract**: The membrane potential is reduced by a certain value after a spike.
+2. **Zero**: The membrane potential is reset to zero after a spike.
+
+The computational cost of these mechanisms is defined in the `neuron_ops_reset_operations` dictionary. For example, the "Leaky" neuron type has the following costs:
+- Subtract mechanism: 4 operations
+- Zero mechanism: 4 operations
+
+Why 4 Operations for "Leaky" Neurons?
+-------------------------------------
+The computational cost of 4 operations for both the "subtract" and "zero" mechanisms in "Leaky" neurons is an abstraction that represents the number of basic mathematical operations required to perform the reset. These operations include the steps involved in updating the membrane potential, checking conditions, and writing the updated state back to memory.
+
+**Subtract Mechanism**
+If the `reset_mechanism` is set to "subtract", the membrane potential :math:`U[t+1]` will have the `threshold` subtracted from it whenever the neuron emits a spike. The update equation is:
+
+.. math::
+
+    U[t+1] = \beta U[t] + I_{\rm in}[t+1] - R U_{\rm thr}
+
+Here’s the breakdown of the 4 operations:
+
+1. **Decay Term**: Multiply the previous membrane potential :math:`U[t]` by the decay factor :math:`\beta` (1 operation).
+2. **Input Current**: Add the input current :math:`I_{\rm in}[t+1]` to the decayed potential (1 operation).
+3. **Reset Multiplication**: Multiply the reset factor :math:`R` by the threshold :math:`U_{\rm thr}` (1 operation).
+4. **Threshold Subtraction**: Subtract the result of the reset multiplication from the decayed potential and input current (1 operation).
+
+**Zero Mechanism**
+If the `reset_mechanism` is set to "zero", the membrane potential :math:`U[t+1]` will be reset to zero whenever the neuron emits a spike. The update equation is:
+
+.. math::
+
+    U[t+1] = \beta U[t] + I_{\rm syn}[t+1] - R(\beta U[t] + I_{\rm in}[t+1])
+
+Here’s the breakdown of the 4 operations:
+
+1. **Decay Term**: Multiply the previous membrane potential :math:`U[t]` by the decay factor :math:`\beta` (1 operation).
+2. **Input Current**: Add the synaptic input current :math:`I_{\rm syn}[t+1]` to the decayed potential (1 operation).
+3. **Reset Multiplication**: Multiply the reset factor :math:`R` by the sum of the decayed potential and input current :math:`(\beta U[t] + I_{\rm in}[t+1])` (1 operation).
+4. **Reset Subtraction**: Subtract the result of the reset multiplication from the decayed potential and input current (1 operation).
+
+**Why Abstract the Cost to 4 Operations?**
+
+The value of 4 operations is an abstraction that simplifies the computational workload into a consistent metric. While the actual number of operations may vary slightly depending on the implementation, this abstraction provides a way to compare the computational cost of different neuron types and reset mechanisms. It is particularly useful for analyzing and optimizing spiking neural networks across various implementations.
+
+Example: Leaky Neuron with snnTorch
+-----------------------------------
+Let’s consider an example using the "Leaky" neuron type with `snntorch`. Assume we have a layer of "Leaky" neurons, and we want to compute the workload for the "subtract" and "zero" reset mechanisms.
+
+1. **Subtract Mechanism**:
+   
+   - After a spike, the membrane potential is reduced by a fixed value.
+   - If there are 100 neurons in the layer and each neuron spikes once, the total computational cost is:
+    .. math::
+        \text{Total Cost} = \text{Number of Neurons} \times \text{Cost per Subtract}
+        = 100 \times 4 = 400 \text{ operations.}
+
+2. **Zero Mechanism**:
+   
+   - After a spike, the membrane potential is reset to zero.
+   - If the same 100 neurons spike once, the total computational cost is:
+    .. math::
+        \text{Total Cost} = \text{Number of Neurons} \times \text{Cost per Zero}
+        = 100 \times 4 = 400 \text{ operations.}
+
+Outputs
+-------
+The `NeuronOperations` metric provides two key outputs:
+
+1. **Effective Neuron Ops**: The total number of operations actually performed by neurons, normalized by the number of samples. This value accounts for the actual activity of the neurons during the forward pass, considering only the updates that occur when neurons spike.
+
+2. **Neuron Dense Ops**: The total number of operations that would be computed if all neurons were updated at every time step, regardless of whether they spiked or not. This represents the theoretical maximum workload for the network under full activity.
+
+These outputs help quantify the computational workload of the network, both in terms of actual activity and theoretical maximum activity, and can be used to optimize the model's efficiency.

examples/nehar/benchmark.py

@@ -1,4 +1,11 @@
 import os
+import sys
+import os
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath("./"))
+sys.path.append(os.path.dirname(SCRIPT_DIR))
+print("Adding path to sys.path:", os.path.dirname(SCRIPT_DIR))
+
 from neurobench.datasets import WISDM
 from training import SpikingNetwork
 from neurobench.processors.postprocessors import ChooseMaxCount
@@ -11,6 +18,7 @@
     SynapticOperations,
     ClassificationAccuracy,
     ActivationSparsityByLayer,
+    NeuronOperations
 )
 from neurobench.metrics.static import (
     ParameterCount,
@@ -34,9 +42,7 @@
     num_outputs = data_module.num_outputs
     num_steps = data_module.num_steps
 
-    spiking_network = SpikingNetwork.load_from_checkpoint(
-        model_path, map_location="cpu"
-    )
+    spiking_network = SpikingNetwork(lr=1)
 
     model = SNNTorchModel(spiking_network.model, custom_forward=True)
     test_set_loader = data_module.test_dataloader()
@@ -48,7 +54,7 @@
 
     # #
     static_metrics = [ParameterCount, Footprint, ConnectionSparsity]
-    workload_metrics = [ActivationSparsity, ActivationSparsityByLayer,MembraneUpdates, SynapticOperations, ClassificationAccuracy]
+    workload_metrics = [ActivationSparsity, ActivationSparsityByLayer,MembraneUpdates, SynapticOperations, ClassificationAccuracy, NeuronOperations]
     # #
     benchmark = Benchmark(
         model, test_set_loader, [], postprocessors, [static_metrics, workload_metrics]

neurobench/metrics/workload/__init__.py

@@ -7,6 +7,7 @@
 from .smape import SMAPE
 from .r2 import R2
 from .coco_map import CocoMap
+from .neuron_operations import NeuronOperations
 
 __stateless__ = [
     "ClassificationAccuracy",
@@ -16,6 +17,12 @@
     "ActivationSparsityByLayer",
 ]
 
-__stateful__ = ["MembraneUpdates", "SynapticOperations", "R2", "CocoMap"]
+__stateful__ = [
+    "MembraneUpdates",
+    "SynapticOperations",
+    "R2",
+    "CocoMap",
+    "NeuronOperations",
+]
 
 __all__ = __stateful__ + __stateless__

neurobench/metrics/workload/neuron_operations.py

@@ -0,0 +1,141 @@
+import torch
+from neurobench.metrics.abstract.workload_metric import AccumulatedMetric
+from collections import defaultdict
+
+
+neuron_ops_reset_operations = {
+    "Leaky": {
+        "subtract": 4,
+        "zero": 4,
+    },
+    "Synaptic": {
+        "subtract": 6,
+        "zero": 8,
+    },
+    "Lapicque": {
+        "subtract": 11,
+        "zero": 11,
+    },
+    "Alpha": {
+        "subtract": 18,
+        "zero": 24,
+    },
+}
+"""
+The `neuron_ops_reset_operations` dictionary defines the computational cost associated
+with resetting the membrane potential of neurons for different neuron types. The reset
+mechanisms are categorized into two types:
+
+1. **Subtract**: Represents a reset mechanism where the membrane potential is reduced
+   by a certain value. The value associated with this mechanism indicates the computational
+   cost (in terms of basic operations) required to perform this type of reset.
+
+2. **Zero**: Represents a reset mechanism where the membrane potential is reset to zero.
+   The value associated with this mechanism indicates the computational cost (in terms of
+   basic operations) required to perform this type of reset.
+
+### Neuron Types and Their Computational Costs:
+- **Leaky**:
+  - Subtract mechanism: 4 operations
+  - Zero mechanism: 4 operations
+- **Synaptic**:
+  - Subtract mechanism: 6 operations
+  - Zero mechanism: 8 operations
+- **Lapicque**:
+  - Subtract mechanism: 11 operations
+  - Zero mechanism: 11 operations
+- **Alpha**:
+  - Subtract mechanism: 18 operations
+  - Zero mechanism: 24 operations
+
+### Purpose:
+The values in this dictionary represent the computational cost (measured in terms of
+basic operations like addition, subtraction, etc.) required for each neuron type to
+reset its membrane potential using a specific reset mechanism.
+"""
+
+
+class NeuronOperations(AccumulatedMetric):
+    """
+    Neuron operations metric.
+
+    This metric computes the number of operations performed by neurons during the
+    forward pass of the model. The operations are tracked per neuron, per layer.
+
+    The `NeuronOperations` metric is designed to measure the computational workload
+    associated with neuron activity in spiking neural networks. Specifically, it tracks
+    the number of operations required to update the membrane potential of neurons during
+    the forward pass. These operations include the reset mechanisms defined in the
+    `neuron_ops_reset_operations` dictionary, such as "subtract" and "zero".
+
+    """
+
+    def __init__(self):
+        """Initialize the NeuronOperations metric."""
+        super().__init__(requires_hooks=True)
+        self.total_samples = 0
+        self.dense = defaultdict(int)
+        self.macs = defaultdict(int)
+
+    def reset(self):
+        """Reset the metric state for a new evaluation."""
+        self.total_samples = 0
+        self.dense = defaultdict(int)
+        self.macs = defaultdict(int)
+
+    def __call__(self, model, preds, data):
+        """
+        Accumulate the neuron operations.
+
+        Args:
+            model: A NeuroBenchModel.
+            preds: A tensor of model predictions.
+            data: A tuple of data and labels.
+        Returns:
+            float: Number of membrane potential updates.
+
+        """
+        for hook in model.activation_hooks:
+            layer_type = hook.layer.__class__.__name__
+            reset_mechanism = hook.layer._reset_mechanism
+            updates = 0
+
+            if len(hook.pre_fire_mem_potential) > 1:
+                pre_fire_mem = torch.stack(hook.pre_fire_mem_potential[1:])
+                post_fire_mem = torch.stack(hook.post_fire_mem_potential[1:])
+                updates += torch.count_nonzero(pre_fire_mem - post_fire_mem).item()
+            if hook.post_fire_mem_potential:
+                updates += hook.post_fire_mem_potential[0].numel()
+
+            self.macs[layer_type] += (
+                updates * neuron_ops_reset_operations[layer_type][reset_mechanism]
+            )
+            self.dense[layer_type] += (
+                hook.post_fire_mem_potential[0].numel()
+                * len(hook.post_fire_mem_potential)
+                * neuron_ops_reset_operations[layer_type][reset_mechanism]
+            )
+
+        self.total_samples += data[0].size(0)
+
+        return self.compute()
+
+    def compute(self):
+        """
+        Compute the total membrane updates normalized by the number of samples.
+
+        Returns:
+            float: Compute the total updates to each neuron's membrane potential within the model,
+            aggregated across all neurons and normalized by the number of samples processed.
+
+        """
+        if self.total_samples == 0:
+            return {"Effective Neuron Ops": 0, "Neuron Dense Ops": 0}
+
+        macs = sum(self.macs.values())
+        dense = sum(self.dense.values())
+
+        return {
+            "Effective Neuron Ops": macs / self.total_samples,
+            "Neuron Dense Ops": dense / self.total_samples,
+        }