Update model evaluator to check multiple outputs

Author: SaoirseARM

backends/arm/test/misc/test_model_evaluator.py

@@ -37,14 +37,13 @@ def test_get_model_error(self):
             example_input,
             "tmp/output_tag0.tosa",
         )
-        max_error, max_absolute_error, max_percentage_error, mae = (
-            evaluator.get_model_error()
-        )
 
-        self.assertEqual(max_error, 1.0)
-        self.assertEqual(max_absolute_error, 1.0)
-        self.assertEqual(max_percentage_error, 25.0)
-        self.assertEqual(mae, 0.25)
+        model_error_dict = evaluator.get_model_error()
+
+        self.assertEqual(model_error_dict["max_error"], [1.0])
+        self.assertEqual(model_error_dict["max_absolute_error"], [1.0])
+        self.assertEqual(model_error_dict["max_percentage_error"], [25.0])
+        self.assertEqual(model_error_dict["mean_absolute_error"], [0.25])
 
     def test_get_compression_ratio(self):
         with tempfile.NamedTemporaryFile(delete=True) as temp_bin:

backends/arm/util/arm_model_evaluator.py

@@ -7,11 +7,26 @@
 import os
 import tempfile
 import zipfile
-from typing import Any, Optional, Tuple
+from collections import defaultdict
+from typing import Optional, Tuple
 
 import torch
 
 
+def flatten_args(args) -> tuple | list:
+    flattened_args: list = []
+    if isinstance(args, torch.Tensor):
+        return [args]
+
+    for arg in args:
+        if isinstance(arg, (tuple, list)):
+            flattened_args.extend(arg)
+        else:
+            flattened_args.append(arg)
+
+    return tuple(flattened_args)
+
+
 class GenericModelEvaluator:
     def __init__(
         self,
@@ -32,31 +47,34 @@ def __init__(
         else:
             self.tosa_output_path = None
 
-    def get_model_error(self) -> tuple[float, float, float, float]:
+    def get_model_error(self) -> defaultdict:
         """
-        Returns the following metrics between the outputs of the FP32 and INT8 model:
+        Returns a dict containing the following metrics between the outputs of the FP32 and INT8 model:
         - Maximum error
         - Maximum absolute error
         - Maximum percentage error
         - Mean absolute error
         """
-        fp32_output = self.fp32_model(*self.example_input)
-        int8_output = self.int8_model(*self.example_input)
-
-        difference = fp32_output - int8_output
-        percentage_error = torch.div(difference, fp32_output) * 100
-
-        max_error = torch.max(difference).item()
-        max_absolute_error = torch.max(torch.abs(difference)).item()
-        max_percentage_error = torch.max(percentage_error).item()
-        mean_absolute_error = torch.mean(torch.abs(difference).float()).item()
-
-        return (
-            float(max_error),
-            float(max_absolute_error),
-            float(max_percentage_error),
-            float(mean_absolute_error),
-        )
+        fp32_outputs = flatten_args(self.fp32_model(*self.example_input))
+        int8_outputs = flatten_args(self.int8_model(*self.example_input))
+
+        model_error_dict = defaultdict(list)
+
+        for fp32_output, int8_output in zip(fp32_outputs, int8_outputs):
+            difference = fp32_output - int8_output
+            percentage_error = torch.div(difference, fp32_output) * 100
+            model_error_dict["max_error"].append(torch.max(difference).item())
+            model_error_dict["max_absolute_error"].append(
+                torch.max(torch.abs(difference)).item()
+            )
+            model_error_dict["max_percentage_error"].append(
+                torch.max(percentage_error).item()
+            )
+            model_error_dict["mean_absolute_error"].append(
+                torch.mean(torch.abs(difference).float()).item()
+            )
+
+        return model_error_dict
 
     def get_compression_ratio(self) -> float:
         """Compute the compression ratio of the outputted TOSA flatbuffer."""
@@ -72,19 +90,10 @@ def get_compression_ratio(self) -> float:
 
         return compression_ratio
 
-    def evaluate(self) -> dict[str, Any]:
-        max_error, max_absolute_error, max_percent_error, mean_absolute_error = (
-            self.get_model_error()
-        )
-        output_metrics = {
-            "name": self.model_name,
-            "metrics": {
-                "max_error": max_error,
-                "max_absolute_error": max_absolute_error,
-                "max_percentage_error": max_percent_error,
-                "mean_absolute_error": mean_absolute_error,
-            },
-        }
+    def evaluate(self) -> dict[any]:
+        model_error_dict = self.get_model_error()
+
+        output_metrics = {"name": self.model_name, "metrics": dict(model_error_dict)}
 
         if self.tosa_output_path:
             # We know output_metrics["metrics"] is list since we just defined it, safe to ignore.

examples/arm/aot_arm_compiler.py

@@ -328,9 +328,11 @@ def get_args():
     )
     args = parser.parse_args()
 
-    if args.evaluate and (args.quantize is None or args.intermediates is None):
+    if args.evaluate and (
+        args.quantize is None or args.intermediates is None or (not args.delegate)
+    ):
         raise RuntimeError(
-            "--evaluate requires --quantize and --intermediates to be enabled."
+            "--evaluate requires --quantize, --intermediates and --delegate to be enabled."
         )
 
     if args.debug:
@@ -378,6 +380,9 @@ def get_args():
         # Wrap quantized model back into an exported_program
         exported_program = torch.export.export_for_training(model, example_inputs)
 
+    if args.intermediates:
+        os.makedirs(args.intermediates, exist_ok=True)
+
     if args.delegate:
         # As we can target multiple output encodings from ArmBackend, one must
         # be specified.