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, 18 insertions, 3 deletions

diff --git a/src/mlia/nn/rewrite/core/train.py b/src/mlia/nn/rewrite/core/train.py
index 42bf653..82af747 100644
--- a/src/mlia/nn/rewrite/core/train.py
+++ b/src/mlia/nn/rewrite/core/train.py
@@ -136,12 +136,27 @@ def train(
                     output_filename = output_model
                 join_in_dir(train_dir, filename, output_filename)
 
+        # Assess the output diff between the parts after the rewrite subgraph
+        # in original and optimized model
+        optimized_end_path = Path(train_dir, "optimized_end.tfrec")
+        end_path = Path(train_dir, "end.tfrec")
+
+        record_model(
+            str(input_tfrec),
+            output_filename,
+            optimized_end_path,
+            num_procs=train_params.num_procs,
+            num_threads=train_params.num_threads,
+        )
+        mae, nrmse = diff_stats(end_path, str(optimized_end_path))
+
     if unmodified_model_dir:
         cast(tempfile.TemporaryDirectory, unmodified_model_dir).cleanup()
 
-    return (
-        results if train_params.checkpoint_at else results[0]
-    )  # only return a list if multiple checkpoints are asked for
+    return (results if train_params.checkpoint_at else results[0]), [
+        mae,
+        nrmse,
+    ]  # only return a list if multiple checkpoints are asked for
 
 
 def eval_in_dir(