Add MobileNetV2Evaluator to Arm Backend

* Add MobileNetV2Evaluator which can measure top-1 and top-5 accuracy of MobileNet V2
* Extend aot_arm_compiler.py to accept more Evaluators
* Improve handling of calibration datasets aot_arm_compiler.py
* Improve typing

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

Author: tom-arm

backends/arm/test/misc/test_model_evaluator.py

@@ -4,17 +4,14 @@
 # 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.util.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))
+COMPRESSION_RATIO_TEST = torch.rand([1024, 1024])
 
 
 def mocked_model_1(input: torch.Tensor) -> torch.Tensor:
@@ -47,20 +44,16 @@ def test_get_model_error(self):
 
     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
+            torch.save(COMPRESSION_RATIO_TEST, temp_bin)
 
             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,
+                temp_bin.name,
             )
 
             ratio = evaluator.get_compression_ratio()
-            self.assertAlmostEqual(ratio, 1.0, places=2)
+            self.assertAlmostEqual(ratio, 1.1, places=1)

backends/arm/util/arm_model_evaluator.py

@@ -4,13 +4,25 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 import os
+import random
 import tempfile
 import zipfile
+
 from collections import defaultdict
-from typing import Optional, Tuple
+from pathlib import Path
+from typing import Any, Optional, Tuple
 
 import torch
+from torch.nn.modules import Module
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+
+# Logger for outputting progress for longer running evaluation
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
 
 def flatten_args(args) -> tuple | list:
@@ -28,6 +40,8 @@ def flatten_args(args) -> tuple | list:
 
 
 class GenericModelEvaluator:
+    REQUIRES_CONFIG = False
+
     def __init__(
         self,
         model_name: str,
@@ -90,7 +104,7 @@ def get_compression_ratio(self) -> float:
 
         return compression_ratio
 
-    def evaluate(self) -> dict[any]:
+    def evaluate(self) -> dict[Any]:
         model_error_dict = self.get_model_error()
 
         output_metrics = {"name": self.model_name, "metrics": dict(model_error_dict)}
@@ -103,3 +117,93 @@ def evaluate(self) -> dict[any]:
             ] = self.get_compression_ratio()
 
         return output_metrics
+
+
+class MobileNetV2Evaluator(GenericModelEvaluator):
+    REQUIRES_CONFIG = True
+
+    def __init__(
+        self,
+        model_name: str,
+        fp32_model: Module,
+        int8_model: Module,
+        example_input: Tuple[torch.Tensor],
+        tosa_output_path: str | None,
+        batch_size: int,
+        validation_dataset_path: str,
+    ) -> None:
+        super().__init__(
+            model_name, fp32_model, int8_model, example_input, tosa_output_path
+        )
+
+        self.__batch_size = batch_size
+        self.__validation_set_path = validation_dataset_path
+
+    @staticmethod
+    def __load_dataset(directory: str) -> datasets.ImageFolder:
+        directory_path = Path(directory)
+        if not directory_path.exists():
+            raise FileNotFoundError(f"Directory: {directory} does not exist.")
+
+        transform = transforms.Compose(
+            [
+                transforms.Resize(256),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.484, 0.454, 0.403], std=[0.225, 0.220, 0.220]
+                ),
+            ]
+        )
+        return datasets.ImageFolder(directory_path, transform=transform)
+
+    @staticmethod
+    def get_calibrator(training_dataset_path: str) -> DataLoader:
+        dataset = MobileNetV2Evaluator.__load_dataset(training_dataset_path)
+        rand_indices = random.sample(range(len(dataset)), k=1000)
+
+        # Return a subset of the dataset to be used for calibration
+        return torch.utils.data.DataLoader(
+            torch.utils.data.Subset(dataset, rand_indices),
+            batch_size=1,
+            shuffle=False,
+        )
+
+    def __evaluate_mobilenet(self) -> Tuple[float, float]:
+        dataset = MobileNetV2Evaluator.__load_dataset(self.__validation_set_path)
+        loaded_dataset = DataLoader(
+            dataset,
+            batch_size=self.__batch_size,
+            shuffle=False,
+        )
+
+        top1_correct = 0
+        top5_correct = 0
+
+        for i, (image, target) in enumerate(loaded_dataset):
+            prediction = self.int8_model(image)
+            top1_prediction = torch.topk(prediction, k=1, dim=1).indices
+            top5_prediction = torch.topk(prediction, k=5, dim=1).indices
+
+            top1_correct += (top1_prediction == target.view(-1, 1)).sum().item()
+            top5_correct += (top5_prediction == target.view(-1, 1)).sum().item()
+
+            logger.info("Iteration: {}".format((i + 1) * self.__batch_size))
+            logger.info(
+                "Top 1: {}".format(top1_correct / ((i + 1) * self.__batch_size))
+            )
+            logger.info(
+                "Top 5: {}".format(top5_correct / ((i + 1) * self.__batch_size))
+            )
+
+        top1_accuracy = top1_correct / len(dataset)
+        top5_accuracy = top5_correct / len(dataset)
+
+        return top1_accuracy, top5_accuracy
+
+    def evaluate(self) -> dict[str, Any]:
+        top1_correct, top5_correct = self.__evaluate_mobilenet()
+        output = super().evaluate()
+
+        output["metrics"]["accuracy"] = {"top-1": top1_correct, "top-5": top5_correct}
+        return output