diff options
author | Benjamin Klimczak <benjamin.klimczak@arm.com> | 2023-07-12 15:18:26 +0100 |
---|---|---|
committer | Benjamin Klimczak <benjamin.klimczak@arm.com> | 2023-10-11 16:16:32 +0100 |
commit | ecc4264b93d4a89fa2cb40518b225d8371b7ffad (patch) | |
tree | 47244d2d67ab6c50bfc15eab768252359eae0df6 /src/mlia/nn/tensorflow/utils.py | |
parent | baaf4de286762c1955c874f78cd802d4703a8ba5 (diff) | |
download | mlia-ecc4264b93d4a89fa2cb40518b225d8371b7ffad.tar.gz |
Enable rewrites for quantized input models
If the input model for rewriting is quantized:
- Record de-quantized TFRecords
- enable writing de-quantized calibration data for the training
- re-generate augmented training data, if needed
- Use quantization-aware training (QAT) to train the replacement models
- Check if replacement model is quantized:
If source model is quantized, we make sure rewrite's output model
is quantized too. Right now, only int8 is supported so raising
an error if any other datatype is present in the output.
Resolves: MLIA-907, MLIA-908, MLIA-927
Signed-off-by: Benjamin Klimczak <benjamin.klimczak@arm.com>
Change-Id: Icb4070a9e6f1fdb5ce36120d73823986e89ac955
Diffstat (limited to 'src/mlia/nn/tensorflow/utils.py')
-rw-r--r-- | src/mlia/nn/tensorflow/utils.py | 30 |
1 files changed, 30 insertions, 0 deletions
diff --git a/src/mlia/nn/tensorflow/utils.py b/src/mlia/nn/tensorflow/utils.py index 77ac529..b8d45c6 100644 --- a/src/mlia/nn/tensorflow/utils.py +++ b/src/mlia/nn/tensorflow/utils.py @@ -123,3 +123,33 @@ def get_tflite_converter( converter.inference_output_type = tf.int8 return converter + + +def get_tflite_model_type_map(model_filename: str | Path) -> dict[str, type]: + """Get type map from tflite model.""" + model_type_map: dict[str, Any] = {} + interpreter = tf.lite.Interpreter(str(model_filename)) + interpreter.allocate_tensors() + input_details = interpreter.get_input_details() + model_type_map = { + input_detail["name"]: input_detail["dtype"] for input_detail in input_details + } + return model_type_map + + +def check_tflite_datatypes(model_filename: str | Path, *allowed_types: type) -> None: + """Check if the model only has the given allowed datatypes.""" + type_map = get_tflite_model_type_map(model_filename) + types = set(type_map.values()) + allowed = set(allowed_types) + unexpected = types - allowed + + def cls_to_str(types: set[type]) -> list[str]: + return [t.__name__ for t in types] + + if len(unexpected) > 0: + raise TypeError( + f"Model {model_filename} has " + f"unexpected data types: {cls_to_str(unexpected)}. " + f"Only {cls_to_str(allowed)} are allowed." + ) |