1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
|
# 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
from mlia.nn.tensorflow.utils import save_keras_model
from mlia.nn.tensorflow.utils import save_tflite_model
def test_generate_representative_dataset() -> None:
"""Test function for generating representative dataset."""
dataset = representative_dataset([1, 3, 3], 5)
data = list(dataset())
assert len(data) == 5
for elem in data:
assert isinstance(elem, list)
assert len(elem) == 1
ndarray = elem[0]
assert ndarray.dtype == np.float32
assert isinstance(ndarray, np.ndarray)
def test_convert_saved_model_to_tflite(test_tf_model: Path) -> None:
"""Test converting SavedModel to TensorFlow Lite."""
result = convert_to_tflite(test_tf_model.as_posix())
assert isinstance(result, bytes)
def test_convert_keras_to_tflite(test_keras_model: Path) -> None:
"""Test converting Keras model to TensorFlow Lite."""
keras_model = tf.keras.models.load_model(str(test_keras_model))
result = convert_to_tflite(keras_model)
assert isinstance(result, bytes)
def test_save_keras_model(tmp_path: Path, test_keras_model: Path) -> None:
"""Test saving Keras model."""
keras_model = tf.keras.models.load_model(str(test_keras_model))
temp_file = tmp_path / "test_model_saving.h5"
save_keras_model(keras_model, temp_file)
loaded_model = tf.keras.models.load_model(temp_file)
assert loaded_model.summary() == keras_model.summary()
def test_save_tflite_model(tmp_path: Path, test_keras_model: Path) -> None:
"""Test saving TensorFlow Lite model."""
keras_model = tf.keras.models.load_model(str(test_keras_model))
tflite_model = convert_to_tflite(keras_model)
temp_file = tmp_path / "test_model_saving.tflite"
save_tflite_model(tflite_model, temp_file)
interpreter = tf.lite.Interpreter(model_path=str(temp_file))
assert interpreter
def test_convert_unknown_model_to_tflite() -> None:
"""Test that unknown model type cannot be converted to TensorFlow Lite."""
with pytest.raises(
ValueError, match="Unable to create TensorFlow Lite converter for 123"
):
convert_to_tflite(123)
@pytest.mark.parametrize(
"model_path, expected_result",
[
[Path("sample_model.tflite"), True],
[Path("strange_model.tflite.tfl"), False],
[Path("sample_model.h5"), False],
[Path("sample_model"), False],
],
)
def test_is_tflite_model(model_path: Path, expected_result: bool) -> None:
"""Test function is_tflite_model."""
result = is_tflite_model(model_path)
assert result == expected_result
@pytest.mark.parametrize(
"model_path, expected_result",
[
[Path("sample_model.h5"), True],
[Path("strange_model.h5.keras"), False],
[Path("sample_model.tflite"), False],
[Path("sample_model"), False],
],
)
def test_is_keras_model(model_path: Path, expected_result: bool) -> None:
"""Test function is_keras_model."""
result = is_keras_model(model_path)
assert result == expected_result
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)
|
