Adding model stats to aot_arm_compiler

* Adds GenericModelEvaluator, which gathers metrics applicable to all models
* Adds --evaluate option to enable gathering quantization metrics

Signed-off-by: Tom Allsop <[email protected]>

Change-Id: Ia9b591841f188870fa5e62d0568169498301393d

Author: SaoirseARM

backends/arm/scripts/arm_model_evaluator.py

@@ -0,0 +1,89 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import tempfile
+import zipfile
+from typing import Optional, Tuple, Union
+
+import torch
+
+
+class GenericModelEvaluator:
+    def __init__(
+        self,
+        model_name: str,
+        fp32_model: torch.nn.Module,
+        int8_model: torch.nn.Module,
+        example_input: Tuple[torch.Tensor],
+        tosa_output_path: Optional[str],
+    ) -> None:
+        self.model_name = model_name
+
+        self.fp32_model = fp32_model
+        self.int8_model = int8_model
+        self.example_input = example_input
+
+        if tosa_output_path:
+            self.tosa_output_path = tosa_output_path
+        else:
+            self.tosa_output_path = None
+
+    def get_model_error(self) -> Union[float, float, float, float]:
+        """
+        Returns 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 max_error, max_absolute_error, max_percentage_error, mean_absolute_error
+
+    def get_compression_ratio(self) -> float:
+        """Compute the compression ratio of the outputted TOSA flatbuffer."""
+        with tempfile.NamedTemporaryFile(delete=True, suffix=".zip") as temp_zip:
+            with zipfile.ZipFile(
+                temp_zip.name, "w", compression=zipfile.ZIP_DEFLATED
+            ) as f:
+                f.write(self.tosa_output_path)
+
+            compression_ratio = os.path.getsize(
+                self.tosa_output_path
+            ) / os.path.getsize(temp_zip.name)
+
+        return compression_ratio
+
+    def evaluate(self) -> dict[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,
+            },
+        }
+
+        if self.tosa_output_path:
+            output_metrics["metrics"][
+                "compression_ratio"
+            ] = self.get_compression_ratio()
+
+        return output_metrics

backends/arm/test/misc/test_model_evaluator.py

@@ -0,0 +1,67 @@
+# Copyright 2024 Arm Limited and/or its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+import tempfile
+import unittest
+
+import torch
+from executorch.backends.arm.scripts.arm_model_evaluator import GenericModelEvaluator
+
+random.seed(0)
+
+# Create an input that is hard to compress
+COMPRESSION_RATIO_TEST = bytearray(random.getrandbits(8) for _ in range(1000000))
+
+
+def mocked_model_1(input: torch.Tensor) -> torch.Tensor:
+    return torch.tensor([1.0, 2.0, 3.0, 4.0])
+
+
+def mocked_model_2(input: torch.Tensor) -> torch.Tensor:
+    return torch.tensor([1.0, 2.0, 3.0, 3.0])
+
+
+class TestGenericModelEvaluator(unittest.TestCase):
+    """Tests the GenericModelEvaluator class."""
+
+    def test_get_model_error(self):
+        example_input = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+        evaluator = GenericModelEvaluator(
+            "dummy_model",
+            mocked_model_1,
+            mocked_model_2,
+            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)
+
+    def test_get_compression_ratio(self):
+        with tempfile.NamedTemporaryFile(delete=True) as temp_bin:
+            temp_bin.write(COMPRESSION_RATIO_TEST)
+
+            # As the size of the file is quite small we need to call flush()
+            temp_bin.flush()
+            temp_bin_name = temp_bin.name
+
+            example_input = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+            evaluator = GenericModelEvaluator(
+                "dummy_model",
+                mocked_model_1,
+                mocked_model_2,
+                example_input,
+                temp_bin_name,
+            )
+
+            ratio = evaluator.get_compression_ratio()
+            self.assertAlmostEqual(ratio, 1.0, places=2)

examples/arm/aot_arm_compiler.py

@@ -8,17 +8,21 @@
 # Example script for exporting simple models to flatbuffer
 
 import argparse
+import json
 import logging
 import os
 
-import torch
+from pathlib import Path
+from typing import Tuple
 
+import torch
 from executorch.backends.arm.arm_backend import ArmCompileSpecBuilder
 from executorch.backends.arm.arm_partitioner import ArmPartitioner
 from executorch.backends.arm.quantizer.arm_quantizer import (
     ArmQuantizer,
     get_symmetric_quantization_config,
 )
+from executorch.backends.arm.scripts.arm_model_evaluator import GenericModelEvaluator
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
 from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 
@@ -147,6 +151,8 @@ def forward(self, x):
     "softmax": SoftmaxModule,
 }
 
+evaluators = {}
+
 targets = [
     "ethos-u55-32",
     "ethos-u55-64",
@@ -198,6 +204,37 @@ def get_compile_spec(target: str, intermediates: bool) -> ArmCompileSpecBuilder:
     return spec_builder.build()
 
 
+def get_evaluator(model_name: str) -> GenericModelEvaluator:
+    if model_name not in evaluators:
+        return GenericModelEvaluator
+    else:
+        return evaluators[model_name]
+
+
+def evaluate_model(
+    model_name: str,
+    intermediates: str,
+    model_fp32: torch.nn.Module,
+    model_int8: torch.nn.Module,
+    example_inputs: Tuple[torch.Tensor],
+):
+    evaluator = get_evaluator(model_name)
+
+    # Get the path of the TOSA flatbuffer that is dumped
+    intermediates_path = Path(intermediates)
+    tosa_paths = list(intermediates_path.glob("*.tosa"))
+
+    init_evaluator = evaluator(
+        model_name, model_fp32, model_int8, example_inputs, str(tosa_paths[0])
+    )
+
+    quant_metrics = init_evaluator.evaluate()
+    output_json_path = intermediates_path / "quant_metrics.json"
+
+    with output_json_path.open("w") as json_file:
+        json.dump(quant_metrics, json_file)
+
+
 def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -223,6 +260,14 @@ def get_args():
         choices=targets,
         help=f"For ArmBackend delegated models, pick the target, and therefore the instruction set generated. valid targets are {targets}",
     )
+    parser.add_argument(
+        "-e",
+        "--evaluate",
+        action="store_true",
+        required=False,
+        default=False,
+        help="Flag for running evaluation of the model.",
+    )
     parser.add_argument(
         "-q",
         "--quantize",
@@ -256,11 +301,11 @@ def get_args():
         help="Location for outputs, if not the default of cwd.",
     )
     args = parser.parse_args()
-    return args
 
-
-if __name__ == "__main__":
-    args = get_args()
+    if args.evaluate and (args.quantize is None or args.intermediates is None):
+        raise RuntimeError(
+            "--evaluate requires --quantize and --intermediates to be enabled."
+        )
 
     if args.debug:
         logging.basicConfig(level=logging.DEBUG, format=FORMAT, force=True)
@@ -283,16 +328,26 @@ def get_args():
     ):
         raise RuntimeError(f"Model {args.model_name} cannot be delegated.")
 
+    return args
+
+
+if __name__ == "__main__":
+    args = get_args()
+
     # Pick model from one of the supported lists
     model, example_inputs = get_model_and_inputs_from_name(args.model_name)
     model = model.eval()
 
+    model_fp32 = model
+
     # pre-autograd export. eventually this will become torch.export
     model = torch.export.export_for_training(model, example_inputs).module()
 
     # Quantize if required
+    model_int8 = None
     if args.quantize:
         model = quantize(model, example_inputs)
+        model_int8 = model
 
     edge = export_to_edge(
         model,
@@ -339,3 +394,8 @@ def get_args():
         output_name = os.path.join(args.output, output_name)
 
     save_pte_program(exec_prog, output_name)
+
+    if args.evaluate:
+        evaluate_model(
+            args.model_name, args.intermediates, model_fp32, model_int8, example_inputs
+        )