From baaf4de286762c1955c874f78cd802d4703a8ba5 Mon Sep 17 00:00:00 2001
From: Gergely Nagy <gergely.nagy@arm.com>
Date: Thu, 22 Jun 2023 14:35:21 +0100
Subject: Re-factoring of rewrite management & added metrics

- List available rewrites
  - Refactor/rename 'Rewrite' class to 'RewritingOptimizer'
  - Introduce a registry for rewrite functions
  - Refactor 'Rewriter' to use the registry to look up rewrite functions
  - Remove mentions of hardcoded "fully_connected" from CLI help and
    error messages, using the registry instead
  - Add unit tests
- Enable rewrites for all targets:
  Extract optimization (including rewrite specific code) from the
  Ethos-U-specific data collector into OptimizingDataCollector.
  This is reused in other targets' collectors, such as TOSA
  and Cortex-A.
- Add more logging for rewrite
  - add display of MAE and NRMSE values for the trained result
  - add total model MAE and NRMSE metric

Resolves: MLIA-891, MLIA-899, MLIA-906
Change-Id: Ie798749e1ed60cab14fdb6d9c2271c833960e93f
Signed-off-by: Benjamin Klimczak <benjamin.klimczak@arm.com>
---
 tests/test_common_optimization.py | 67 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)
 create mode 100644 tests/test_common_optimization.py

(limited to 'tests/test_common_optimization.py')

diff --git a/tests/test_common_optimization.py b/tests/test_common_optimization.py
new file mode 100644
index 0000000..599610d
--- /dev/null
+++ b/tests/test_common_optimization.py
@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright 2023, Arm Limited and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the common optimization module."""
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from mlia.core.context import ExecutionContext
+from mlia.nn.common import Optimizer
+from mlia.nn.tensorflow.config import TFLiteModel
+from mlia.target.common.optimization import OptimizingDataCollector
+from mlia.target.config import TargetProfile
+
+
+class FakeOptimizer(Optimizer):
+    """Optimizer for testing purposes."""
+
+    def __init__(self, optimized_model_path: Path) -> None:
+        """Initialize."""
+        super().__init__()
+        self.optimized_model_path = optimized_model_path
+        self.invocation_count = 0
+
+    def apply_optimization(self) -> None:
+        """Count the invocations."""
+        self.invocation_count += 1
+
+    def get_model(self) -> TFLiteModel:
+        """Return optimized model."""
+        return TFLiteModel(self.optimized_model_path)
+
+    def optimization_config(self) -> str:
+        """Return something: doesn't matter, not used."""
+        return ""
+
+
+def test_optimizing_data_collector(
+    test_keras_model: Path,
+    test_tflite_model: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test OptimizingDataCollector, base support for various targets."""
+    optimizations = [
+        [
+            {"optimization_type": "fake", "optimization_target": 42},
+        ]
+    ]
+    context = ExecutionContext(
+        config_parameters={"common_optimizations": {"optimizations": optimizations}}
+    )
+
+    target_profile = MagicMock(spec=TargetProfile)
+
+    fake_optimizer = FakeOptimizer(test_tflite_model)
+
+    monkeypatch.setattr(
+        "mlia.target.common.optimization.get_optimizer",
+        MagicMock(return_value=fake_optimizer),
+    )
+
+    collector = OptimizingDataCollector(test_keras_model, target_profile)
+
+    collector.set_context(context)
+    collector.collect_data()
+
+    assert fake_optimizer.invocation_count == 1
-- 
cgit v1.2.1