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>

1 files changed, 40 insertions, 22 deletions

diff --git a/tests/test_target_ethos_u_data_collection.py b/tests/test_target_ethos_u_data_collection.py
index 6244f8b..be93c26 100644
--- a/tests/test_target_ethos_u_data_collection.py
+++ b/tests/test_target_ethos_u_data_collection.py
@@ -8,9 +8,11 @@ import pytest
 
 from mlia.backend.vela.compat import Operators
 from mlia.core.context import Context
+from mlia.core.context import ExecutionContext
 from mlia.core.data_collection import DataCollector
 from mlia.core.errors import FunctionalityNotSupportedError
 from mlia.nn.select import OptimizationSettings
+from mlia.target.common.optimization import add_common_optimization_params
 from mlia.target.ethos_u.config import EthosUConfiguration
 from mlia.target.ethos_u.data_collection import EthosUOperatorCompatibility
 from mlia.target.ethos_u.data_collection import EthosUOptimizationPerformance
@@ -46,6 +48,20 @@ def test_collectors_metadata(
     assert collector.name() == expected_name
 
 
+def setup_optimization(optimizations: list) -> Context:
+    """Set up optimization params for the context."""
+    params: dict = {}
+    add_common_optimization_params(
+        params,
+        {
+            "optimization_targets": optimizations,
+        },
+    )
+
+    context = ExecutionContext(config_parameters=params)
+    return context
+
+
 def test_operator_compatibility_collector(
     sample_context: Context, test_tflite_model: Path
 ) -> None:
@@ -76,7 +92,6 @@ def test_performance_collector(
 
 def test_optimization_performance_collector(
     monkeypatch: pytest.MonkeyPatch,
-    sample_context: Context,
     test_keras_model: Path,
     test_tflite_model: Path,
 ) -> None:
@@ -84,16 +99,14 @@ def test_optimization_performance_collector(
     target = EthosUConfiguration.load_profile("ethos-u55-256")
 
     mock_performance_estimation(monkeypatch, target)
-    collector = EthosUOptimizationPerformance(
-        test_keras_model,
-        target,
+
+    context = setup_optimization(
         [
-            [
-                {"optimization_type": "pruning", "optimization_target": 0.5},
-            ]
+            {"optimization_type": "pruning", "optimization_target": 0.5},
         ],
     )
-    collector.set_context(sample_context)
+    collector = EthosUOptimizationPerformance(test_keras_model, target)
+    collector.set_context(context)
     result = collector.collect_data()
 
     assert isinstance(result, OptimizationPerformanceMetrics)
@@ -105,34 +118,39 @@ def test_optimization_performance_collector(
     assert opt == [OptimizationSettings("pruning", 0.5, None)]
     assert isinstance(metrics, PerformanceMetrics)
 
-    collector_no_optimizations = EthosUOptimizationPerformance(
-        test_keras_model,
-        target,
-        [],
+    context = ExecutionContext(
+        config_parameters={"common_optimizations": {"optimizations": [[]]}}
     )
+
+    collector_no_optimizations = EthosUOptimizationPerformance(test_keras_model, target)
+    collector_no_optimizations.set_context(context)
     with pytest.raises(FunctionalityNotSupportedError):
         collector_no_optimizations.collect_data()
 
-    collector_tflite = EthosUOptimizationPerformance(
-        test_tflite_model,
-        target,
+    context = setup_optimization(
         [
-            [
-                {"optimization_type": "pruning", "optimization_target": 0.5},
-            ]
+            {"optimization_type": "pruning", "optimization_target": 0.5},
         ],
     )
-    collector_tflite.set_context(sample_context)
+
+    collector_tflite = EthosUOptimizationPerformance(test_tflite_model, target)
+    collector_tflite.set_context(context)
     with pytest.raises(FunctionalityNotSupportedError):
         collector_tflite.collect_data()
 
     with pytest.raises(
         Exception, match="Optimization parameters expected to be a list"
     ):
-        collector_bad_config = EthosUOptimizationPerformance(
-            test_keras_model, target, {"optimization_type": "pruning"}  # type: ignore
+        context = ExecutionContext(
+            config_parameters={
+                "common_optimizations": {
+                    "optimizations": [{"optimization_type": "pruning"}]
+                }
+            }
         )
-        collector.set_context(sample_context)
+
+        collector_bad_config = EthosUOptimizationPerformance(test_keras_model, target)
+        collector_bad_config.set_context(context)
         collector_bad_config.collect_data()