Merge pull request #250 from NeuroBench/feature/processors

Rafactor processors to accept callables

Author: jasonlyik

neurobench/benchmarks/benchmark.py

@@ -16,7 +16,7 @@
 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
@@ -36,18 +36,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
@@ -63,8 +70,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]:
         """
@@ -114,10 +126,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)

neurobench/processors/manager/proc_manager.py

@@ -3,7 +3,7 @@
     NeuroBenchPreProcessor,
     NeuroBenchPostProcessor,
 )
-from typing import List
+from typing import List, Callable, Tuple
 from torch import Tensor
 
 
@@ -26,8 +26,11 @@ class ProcessorManager(ABC):
 
     def __init__(
         self,
-        preprocessors: List[NeuroBenchPreProcessor],
-        postprocessors: List[NeuroBenchPostProcessor],
+        preprocessors: List[
+            NeuroBenchPreProcessor
+            | Callable[[Tuple[Tensor, Tensor]], Tuple[Tensor, Tensor]]
+        ],
+        postprocessors: List[NeuroBenchPostProcessor | Callable[[Tensor], Tensor]],
     ):
         """
         Initialize the ProcessorManager with the given preprocessors and postprocessors.
@@ -45,19 +48,31 @@ def __init__(
 
         """
 
-        if any(not isinstance(p, NeuroBenchPreProcessor) for p in preprocessors):
+        if any(
+            not (isinstance(p, NeuroBenchPreProcessor) or callable(p))
+            for p in preprocessors
+        ):
             raise TypeError(
-                "All preprocessors must be instances of NeuroBenchPreProcessor"
+                "All preprocessors must be instances of NeuroBenchPreProcessor or callable functions"
             )
-        if any(not isinstance(p, NeuroBenchPostProcessor) for p in postprocessors):
+        if any(
+            not (isinstance(p, NeuroBenchPostProcessor) or callable(p))
+            for p in postprocessors
+        ):
             raise TypeError(
-                "All postprocessors must be instances of NeuroBenchPostProcessor"
+                "All postprocessors must be instances of NeuroBenchPostProcessor or callable functions"
             )
 
         self.preprocessors = preprocessors
         self.postprocessors = postprocessors
 
-    def replace_preprocessors(self, preprocessors: List[NeuroBenchPreProcessor]):
+    def replace_preprocessors(
+        self,
+        preprocessors: List[
+            NeuroBenchPreProcessor
+            | Callable[[Tuple[Tensor, Tensor]], Tuple[Tensor, Tensor]]
+        ],
+    ):
         """
         Replace the current list of preprocessors with the provided list.
 
@@ -70,13 +85,18 @@ def replace_preprocessors(self, preprocessors: List[NeuroBenchPreProcessor]):
             NeuroBenchPreProcessor.
 
         """
-        if any(not isinstance(p, NeuroBenchPreProcessor) for p in preprocessors):
+        if any(
+            not (isinstance(p, NeuroBenchPreProcessor) or callable(p))
+            for p in preprocessors
+        ):
             raise TypeError(
-                "All preprocessors must be instances of NeuroBenchPreProcessor"
+                "All preprocessors must be instances of NeuroBenchPreProcessor or callable functions"
             )
         self.preprocessors = preprocessors
 
-    def replace_postprocessors(self, postprocessors: List[NeuroBenchPostProcessor]):
+    def replace_postprocessors(
+        self, postprocessors: List[NeuroBenchPostProcessor | Callable[[Tensor], Tensor]]
+    ):
         """
         Replace the current list of postprocessors with the provided list.
 
@@ -89,13 +109,16 @@ def replace_postprocessors(self, postprocessors: List[NeuroBenchPostProcessor]):
             NeuroBenchPostProcessor.
 
         """
-        if any(not isinstance(p, NeuroBenchPostProcessor) for p in postprocessors):
+        if any(
+            not (isinstance(p, NeuroBenchPostProcessor) or callable(p))
+            for p in postprocessors
+        ):
             raise TypeError(
-                "All postprocessors must be instances of NeuroBenchPostProcessor"
+                "All postprocessors must be instances of NeuroBenchPostProcessor or callable functions"
             )
         self.postprocessors = postprocessors
 
-    def preprocess(self, data: tuple[Tensor, Tensor]) -> tuple[Tensor, Tensor]:
+    def preprocess(self, data: Tuple[Tensor, Tensor]) -> tuple[Tensor, Tensor]:
         """
         Apply preprocessing steps to the input data.