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 /tests/test_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 'tests/test_nn_tensorflow_utils.py')
-rw-r--r-- | tests/test_nn_tensorflow_utils.py | 31 |
1 files changed, 31 insertions, 0 deletions
diff --git a/tests/test_nn_tensorflow_utils.py b/tests/test_nn_tensorflow_utils.py index 14b06c4..dab8b4e 100644 --- a/tests/test_nn_tensorflow_utils.py +++ b/tests/test_nn_tensorflow_utils.py @@ -1,14 +1,17 @@ # SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates. # SPDX-License-Identifier: Apache-2.0 """Test for module utils/test_utils.""" +import re from pathlib import Path import numpy as np import pytest import tensorflow as tf +from mlia.nn.tensorflow.utils import check_tflite_datatypes from mlia.nn.tensorflow.utils import convert_to_tflite from mlia.nn.tensorflow.utils import get_tf_tensor_shape +from mlia.nn.tensorflow.utils import get_tflite_model_type_map from mlia.nn.tensorflow.utils import is_keras_model from mlia.nn.tensorflow.utils import is_tflite_model from mlia.nn.tensorflow.utils import representative_dataset @@ -109,3 +112,31 @@ def test_is_keras_model(model_path: Path, expected_result: bool) -> None: def test_get_tf_tensor_shape(test_tf_model: Path) -> None: """Test get_tf_tensor_shape with test model.""" assert get_tf_tensor_shape(str(test_tf_model)) == [1, 28, 28, 1] + + +def test_tflite_model_type_map( + test_tflite_model_fp32: Path, test_tflite_model: Path +) -> None: + """Test the model type map function.""" + assert get_tflite_model_type_map(test_tflite_model_fp32) == { + "serving_default_input:0": np.float32 + } + assert get_tflite_model_type_map(test_tflite_model) == { + "serving_default_input:0": np.int8 + } + + +def test_check_tflite_datatypes( + test_tflite_model_fp32: Path, test_tflite_model: Path +) -> None: + """Test the model type map function.""" + check_tflite_datatypes(test_tflite_model_fp32, np.float32) + check_tflite_datatypes(test_tflite_model, np.int8) + + with pytest.raises( + Exception, + match=re.escape( + "unexpected data types: ['float32']. Only ['int8'] are allowed" + ), + ): + check_tflite_datatypes(test_tflite_model_fp32, np.int8) |