aboutsummaryrefslogtreecommitdiff
path: root/tests
diff options
context:
space:
mode:
Diffstat (limited to 'tests')
-rw-r--r--tests/conftest.py19
-rw-r--r--tests/test_nn_rewrite_core_train.py15
-rw-r--r--tests/test_nn_select.py4
-rw-r--r--tests/test_nn_tensorflow_optimizations_clustering.py8
-rw-r--r--tests/test_nn_tensorflow_optimizations_pruning.py8
-rw-r--r--tests/test_nn_tensorflow_tflite_compat.py12
-rw-r--r--tests/test_nn_tensorflow_tflite_convert.py7
-rw-r--r--tests/test_nn_tensorflow_tflite_metrics.py21
-rw-r--r--tests/test_nn_tensorflow_utils.py9
-rw-r--r--tests/test_target_cortex_a_operators.py12
-rw-r--r--tests/utils/common.py10
11 files changed, 65 insertions, 60 deletions
diff --git a/tests/conftest.py b/tests/conftest.py
index 53bfb0c..3d0b832 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -11,6 +11,7 @@ import _pytest
import numpy as np
import pytest
import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.backend.vela.compiler import compile_model
from mlia.core.context import ExecutionContext
@@ -103,21 +104,21 @@ def fixture_vela_ini_file(
return test_vela_path / "vela.ini"
-def get_test_keras_model() -> tf.keras.Model:
+def get_test_keras_model() -> keras.Model:
"""Return test Keras model."""
- model = tf.keras.Sequential(
+ model = keras.Sequential(
[
- tf.keras.Input(shape=(28, 28, 1), batch_size=1, name="input"),
- tf.keras.layers.Reshape((28, 28, 1)),
- tf.keras.layers.Conv2D(
+ keras.Input(shape=(28, 28, 1), batch_size=1, name="input"),
+ keras.layers.Reshape((28, 28, 1)),
+ keras.layers.Conv2D(
filters=12, kernel_size=(3, 3), activation="relu", name="conv1"
),
- tf.keras.layers.Conv2D(
+ keras.layers.Conv2D(
filters=12, kernel_size=(3, 3), activation="relu", name="conv2"
),
- tf.keras.layers.MaxPool2D(2, 2),
- tf.keras.layers.Flatten(),
- tf.keras.layers.Dense(10, name="output"),
+ keras.layers.MaxPool2D(2, 2),
+ keras.layers.Flatten(),
+ keras.layers.Dense(10, name="output"),
]
)
diff --git a/tests/test_nn_rewrite_core_train.py b/tests/test_nn_rewrite_core_train.py
index 7fb6f85..6d24133 100644
--- a/tests/test_nn_rewrite_core_train.py
+++ b/tests/test_nn_rewrite_core_train.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2023-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Tests for module mlia.nn.rewrite.core.train."""
# pylint: disable=too-many-arguments
@@ -12,6 +12,7 @@ from typing import Any
import numpy as np
import pytest
import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.rewrite.core.train import augment_fn_twins
from mlia.nn.rewrite.core.train import AUGMENTATION_PRESETS
@@ -24,7 +25,7 @@ from tests.utils.rewrite import MockTrainingParameters
def replace_fully_connected_with_conv(
input_shape: Any, output_shape: Any
-) -> tf.keras.Model:
+) -> keras.Model:
"""Get a replacement model for the fully connected layer."""
for name, shape in {
"Input": input_shape,
@@ -33,11 +34,11 @@ def replace_fully_connected_with_conv(
if len(shape) != 1:
raise RuntimeError(f"{name}: shape (N,) expected, but it is {input_shape}.")
- model = tf.keras.Sequential(name="RewriteModel")
- model.add(tf.keras.Input(input_shape))
- model.add(tf.keras.layers.Reshape((1, 1, input_shape[0])))
- model.add(tf.keras.layers.Conv2D(filters=output_shape[0], kernel_size=(1, 1)))
- model.add(tf.keras.layers.Reshape(output_shape))
+ model = keras.Sequential(name="RewriteModel")
+ model.add(keras.Input(input_shape))
+ model.add(keras.layers.Reshape((1, 1, input_shape[0])))
+ model.add(keras.layers.Conv2D(filters=output_shape[0], kernel_size=(1, 1)))
+ model.add(keras.layers.Reshape(output_shape))
return model
diff --git a/tests/test_nn_select.py b/tests/test_nn_select.py
index 15abf2d..aac07b4 100644
--- a/tests/test_nn_select.py
+++ b/tests/test_nn_select.py
@@ -10,7 +10,7 @@ from typing import Any
from typing import cast
import pytest
-import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.core.errors import ConfigurationError
from mlia.nn.rewrite.core.rewrite import RewriteConfiguration
@@ -175,7 +175,7 @@ def test_get_optimizer(
) or isinstance(config, RewriteConfiguration):
model = test_tflite_model
else:
- model = tf.keras.models.load_model(str(test_keras_model))
+ model = keras.models.load_model(str(test_keras_model))
optimizer = get_optimizer(model, config)
assert isinstance(optimizer, expected_type)
assert optimizer.optimization_config() == expected_config
diff --git a/tests/test_nn_tensorflow_optimizations_clustering.py b/tests/test_nn_tensorflow_optimizations_clustering.py
index 72ade58..11036ad 100644
--- a/tests/test_nn_tensorflow_optimizations_clustering.py
+++ b/tests/test_nn_tensorflow_optimizations_clustering.py
@@ -7,8 +7,8 @@ import math
from pathlib import Path
import pytest
-import tensorflow as tf
from flaky import flaky
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.tensorflow.optimizations.clustering import Clusterer
from mlia.nn.tensorflow.optimizations.clustering import ClusteringConfiguration
@@ -22,8 +22,8 @@ from tests.utils.common import train_model
def _prune_model(
- model: tf.keras.Model, target_sparsity: float, layers_to_prune: list[str] | None
-) -> tf.keras.Model:
+ model: keras.Model, target_sparsity: float, layers_to_prune: list[str] | None
+) -> keras.Model:
x_train, y_train = get_dataset()
batch_size = 1
num_epochs = 1
@@ -100,7 +100,7 @@ def test_cluster_simple_model_fully(
"""Simple MNIST test to see if clustering works correctly."""
target_sparsity = 0.5
- base_model = tf.keras.models.load_model(str(test_keras_model))
+ base_model = keras.models.load_model(str(test_keras_model))
train_model(base_model)
if sparsity_aware:
diff --git a/tests/test_nn_tensorflow_optimizations_pruning.py b/tests/test_nn_tensorflow_optimizations_pruning.py
index 9afc3ff..c942d83 100644
--- a/tests/test_nn_tensorflow_optimizations_pruning.py
+++ b/tests/test_nn_tensorflow_optimizations_pruning.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Test for module optimizations/pruning."""
from __future__ import annotations
@@ -6,7 +6,7 @@ from __future__ import annotations
from pathlib import Path
import pytest
-import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from numpy.core.numeric import isclose
from mlia.nn.tensorflow.optimizations.pruning import Pruner
@@ -47,7 +47,7 @@ def _test_check_sparsity(base_tflite_metrics: TFLiteMetrics) -> None:
def _get_tflite_metrics(
- path: Path, tflite_fn: str, model: tf.keras.Model
+ path: Path, tflite_fn: str, model: keras.Model
) -> TFLiteMetrics:
"""Save model as TFLiteModel and return metrics."""
temp_file = path / tflite_fn
@@ -70,7 +70,7 @@ def test_prune_simple_model_fully(
batch_size = 1
num_epochs = 1
- base_model = tf.keras.models.load_model(str(test_keras_model))
+ base_model = keras.models.load_model(str(test_keras_model))
train_model(base_model)
base_tflite_metrics = _get_tflite_metrics(
diff --git a/tests/test_nn_tensorflow_tflite_compat.py b/tests/test_nn_tensorflow_tflite_compat.py
index 4ca387c..ee60ff7 100644
--- a/tests/test_nn_tensorflow_tflite_compat.py
+++ b/tests/test_nn_tensorflow_tflite_compat.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Tests for tflite_compat module."""
from __future__ import annotations
@@ -6,7 +6,7 @@ from __future__ import annotations
from unittest.mock import MagicMock
import pytest
-import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from tensorflow.lite.python import convert
from mlia.nn.tensorflow.tflite_compat import converter_error_data_pb2
@@ -19,11 +19,11 @@ from mlia.nn.tensorflow.tflite_compat import TFLiteConversionErrorCode
def test_not_fully_compatible_model_flex_ops() -> None:
"""Test models that requires TF_SELECT_OPS."""
- model = tf.keras.models.Sequential(
+ model = keras.models.Sequential(
[
- tf.keras.layers.Dense(units=1, input_shape=[1], batch_size=1),
- tf.keras.layers.Dense(units=16, activation="softsign"),
- tf.keras.layers.Dense(units=1),
+ keras.layers.Dense(units=1, input_shape=[1], batch_size=1),
+ keras.layers.Dense(units=16, activation="softsign"),
+ keras.layers.Dense(units=1),
]
)
diff --git a/tests/test_nn_tensorflow_tflite_convert.py b/tests/test_nn_tensorflow_tflite_convert.py
index 3125c04..81655b1 100644
--- a/tests/test_nn_tensorflow_tflite_convert.py
+++ b/tests/test_nn_tensorflow_tflite_convert.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Test for module utils/test_utils."""
import os
@@ -9,6 +9,7 @@ from unittest.mock import MagicMock
import numpy as np
import pytest
import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.tensorflow import tflite_convert
from mlia.nn.tensorflow.tflite_convert import convert_to_tflite
@@ -40,14 +41,14 @@ def test_convert_saved_model_to_tflite(test_tf_model: Path) -> None:
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))
+ keras_model = keras.models.load_model(str(test_keras_model))
result = convert_to_tflite_bytes(keras_model)
assert isinstance(result, bytes)
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))
+ keras_model = keras.models.load_model(str(test_keras_model))
temp_file = tmp_path / "test_model_saving.tflite"
convert_to_tflite(keras_model, output_path=temp_file)
diff --git a/tests/test_nn_tensorflow_tflite_metrics.py b/tests/test_nn_tensorflow_tflite_metrics.py
index e8d7c09..cbb1b63 100644
--- a/tests/test_nn_tensorflow_tflite_metrics.py
+++ b/tests/test_nn_tensorflow_tflite_metrics.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Test for module utils/tflite_metrics."""
from __future__ import annotations
@@ -12,26 +12,27 @@ from typing import Generator
import numpy as np
import pytest
import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.tensorflow.tflite_metrics import ReportClusterMode
from mlia.nn.tensorflow.tflite_metrics import TFLiteMetrics
-def _sample_keras_model() -> tf.keras.Model:
+def _sample_keras_model() -> keras.Model:
# Create a sample model
- keras_model = tf.keras.Sequential(
+ keras_model = keras.Sequential(
[
- tf.keras.Input(shape=(8, 8, 3)),
- tf.keras.layers.Conv2D(4, 3),
- tf.keras.layers.DepthwiseConv2D(3),
- tf.keras.layers.Flatten(),
- tf.keras.layers.Dense(8),
+ keras.Input(shape=(8, 8, 3)),
+ keras.layers.Conv2D(4, 3),
+ keras.layers.DepthwiseConv2D(3),
+ keras.layers.Flatten(),
+ keras.layers.Dense(8),
]
)
return keras_model
-def _sparse_binary_keras_model() -> tf.keras.Model:
+def _sparse_binary_keras_model() -> keras.Model:
def get_sparse_weights(shape: list[int]) -> np.ndarray:
weights = np.zeros(shape)
with np.nditer(weights, op_flags=[["writeonly"]]) as weight_it:
@@ -43,7 +44,7 @@ def _sparse_binary_keras_model() -> tf.keras.Model:
keras_model = _sample_keras_model()
# Assign weights to have 0.5 sparsity
for layer in keras_model.layers:
- if not isinstance(layer, tf.keras.layers.Flatten):
+ if not isinstance(layer, keras.layers.Flatten):
weight = layer.weights[0]
weight.assign(get_sparse_weights(weight.shape))
print(layer)
diff --git a/tests/test_nn_tensorflow_utils.py b/tests/test_nn_tensorflow_utils.py
index e356a49..4e0c1e1 100644
--- a/tests/test_nn_tensorflow_utils.py
+++ b/tests/test_nn_tensorflow_utils.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Test for module utils/test_utils."""
import re
@@ -7,6 +7,7 @@ from pathlib import Path
import numpy as np
import pytest
import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.tensorflow.tflite_convert import convert_to_tflite
from mlia.nn.tensorflow.utils import check_tflite_datatypes
@@ -19,18 +20,18 @@ from mlia.nn.tensorflow.utils import save_keras_model
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))
+ keras_model = 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)
+ loaded_model = 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))
+ keras_model = keras.models.load_model(str(test_keras_model))
temp_file = tmp_path / "test_model_saving.tflite"
convert_to_tflite(keras_model, output_path=temp_file)
diff --git a/tests/test_target_cortex_a_operators.py b/tests/test_target_cortex_a_operators.py
index 16cdca5..a4cfb2e 100644
--- a/tests/test_target_cortex_a_operators.py
+++ b/tests/test_target_cortex_a_operators.py
@@ -1,10 +1,10 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Tests for Cortex-A operator compatibility."""
from pathlib import Path
import pytest
-import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
from mlia.nn.tensorflow.tflite_convert import convert_to_tflite_bytes
from mlia.target.cortex_a.config import CortexAConfiguration
@@ -42,13 +42,13 @@ def test_get_cortex_a_compatibility_info_not_compatible(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Construct and test a NOT fully compatible TensorFlow Lite model."""
- keras_model = tf.keras.Sequential(
+ keras_model = keras.Sequential(
[
- tf.keras.Input(shape=(28, 28, 1), batch_size=1, name="input"),
- tf.keras.layers.Conv2D(
+ keras.Input(shape=(28, 28, 1), batch_size=1, name="input"),
+ keras.layers.Conv2D(
filters=12, kernel_size=(3, 3), activation="softmax", name="conv1"
),
- tf.keras.layers.LeakyReLU(),
+ keras.layers.LeakyReLU(),
]
)
keras_model.compile(optimizer="sgd", loss="mean_squared_error")
diff --git a/tests/utils/common.py b/tests/utils/common.py
index c29b47c..eafa31b 100644
--- a/tests/utils/common.py
+++ b/tests/utils/common.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2022-2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Common test utils module."""
from __future__ import annotations
@@ -6,12 +6,12 @@ from __future__ import annotations
from pathlib import Path
import numpy as np
-import tensorflow as tf
+from keras.api._v2 import keras # Temporary workaround for now: MLIA-1107
def get_dataset() -> tuple[np.ndarray, np.ndarray]:
"""Return sample dataset."""
- mnist = tf.keras.datasets.mnist
+ mnist = keras.datasets.mnist
(x_train, y_train), _ = mnist.load_data()
x_train = x_train / 255.0
@@ -22,11 +22,11 @@ def get_dataset() -> tuple[np.ndarray, np.ndarray]:
return x_train, y_train
-def train_model(model: tf.keras.Model) -> None:
+def train_model(model: keras.Model) -> None:
"""Train model using sample dataset."""
num_epochs = 1
- loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+ loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])
x_train, y_train = get_dataset()