Bug-fixes and re-factoring for the rewrite module

- Fix input shape of rewrite replacement:
  During and after training of the replacement model for a rewrite the
  Keras model is converted and saved in TensorFlow Lite format. If the
  input shape does not match the teacher model exactly, e.g. if the
  batch size is undefined, the TFLiteConverter adds extra operators
  during conversion.
- Fix rewritten model output
  - Save the model output with the rewritten operator in the output dir
  - Log MAE and NRMSE of the rewrite
- Remove 'verbose' flag from rewrite module and rely on the logging
  mechanism to control verbose output.
- Re-factor utility classes for rewrites
  - Merge the two TFLiteModel classes
  - Move functionality to load/save TensorFlow Lite flatbuffers to
    nn/tensorflow/tflite_graph
- Fix issue with unknown shape in datasets
  After upgrading to TensorFlow 2.12 the unknown shape of the
  TFRecordDataset is causing problems when training the replacement models
  for rewrites. By explicitly setting the right shape of the tensors we
  can work around the issue.
- Adapt default parameters for rewrites. The training steps especially
  had to be increased significantly to be effective.

Resolves: MLIA-895, MLIA-907, MLIA-946, MLIA-979
Signed-off-by: Benjamin Klimczak <benjamin.klimczak@arm.com>
Change-Id: I887ad165aed0f2c6e5a0041f64cec5e6c5ab5c5c

1 files changed, 26 insertions, 1 deletions

diff --git a/src/mlia/nn/tensorflow/tflite_graph.py b/src/mlia/nn/tensorflow/tflite_graph.py
index 4f5e85f..7ca9337 100644
--- a/src/mlia/nn/tensorflow/tflite_graph.py
+++ b/src/mlia/nn/tensorflow/tflite_graph.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
 """Utilities for TensorFlow Lite graphs."""
 from __future__ import annotations
@@ -10,7 +10,10 @@ from pathlib import Path
 from typing import Any
 from typing import cast
 
+import flatbuffers
 from tensorflow.lite.python import schema_py_generated as schema_fb
+from tensorflow.lite.python.schema_py_generated import Model
+from tensorflow.lite.python.schema_py_generated import ModelT
 from tensorflow.lite.tools import visualize
 
 
@@ -137,3 +140,25 @@ def parse_subgraphs(tflite_file: Path) -> list[list[Op]]:
     ]
 
     return graphs
+
+
+def load_fb(input_tflite_file: str | Path) -> ModelT:
+    """Load a flatbuffer model from file."""
+    if not Path(input_tflite_file).exists():
+        raise FileNotFoundError(f"TFLite file not found at {input_tflite_file}\n")
+    with open(input_tflite_file, "rb") as file_handle:
+        file_data = bytearray(file_handle.read())
+    model_obj = Model.GetRootAsModel(file_data, 0)
+    model = ModelT.InitFromObj(model_obj)
+    return model
+
+
+def save_fb(model: ModelT, output_tflite_file: str | Path) -> None:
+    """Save a flatbuffer model to a given file."""
+    builder = flatbuffers.Builder(1024)  # Initial size of the buffer, which
+    # will grow automatically if needed
+    model_offset = model.Pack(builder)
+    builder.Finish(model_offset, file_identifier=b"TFL3")
+    model_data = builder.Output()
+    with open(output_tflite_file, "wb") as out_file:
+        out_file.write(model_data)