From 8dec455c467b8019223a40e107378845e1419f5d Mon Sep 17 00:00:00 2001 From: Annie Tallund Date: Tue, 19 Sep 2023 16:34:14 +0200 Subject: MLIA-469 Support batch size > 1 for optimizations - Add a PruningPolicy to skip layers that are not supported by the Keras pruning API - Make dataset generation more generic to support use-cases beyond classification Signed-off-by: Annie Tallund Change-Id: I198dae2b53860f449f2fdbc71575babceed1ffcf --- tests/test_nn_tensorflow_utils.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) (limited to 'tests') diff --git a/tests/test_nn_tensorflow_utils.py b/tests/test_nn_tensorflow_utils.py index 5131171..14b06c4 100644 --- a/tests/test_nn_tensorflow_utils.py +++ b/tests/test_nn_tensorflow_utils.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 """Test for module utils/test_utils.""" from pathlib import Path @@ -31,12 +31,6 @@ def test_generate_representative_dataset() -> None: assert isinstance(ndarray, np.ndarray) -def test_generate_representative_dataset_wrong_shape() -> None: - """Test that only shape with batch size=1 is supported.""" - with pytest.raises(Exception, match="Only the input batch_size=1 is supported!"): - representative_dataset([2, 3, 3], 5) - - 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()) -- cgit v1.2.1