diff options
author | Annie Tallund <annie.tallund@arm.com> | 2023-09-19 16:34:14 +0200 |
---|---|---|
committer | Annie Tallund <annie.tallund@arm.com> | 2023-09-26 14:24:39 +0200 |
commit | 8dec455c467b8019223a40e107378845e1419f5d (patch) | |
tree | 823051d364c78e3ce7f70023eb94021d8b051274 /tests/test_nn_tensorflow_utils.py | |
parent | ba251631768f25b840e93ece6a4af3db119e6dd1 (diff) | |
download | mlia-8dec455c467b8019223a40e107378845e1419f5d.tar.gz |
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 <annie.tallund@arm.com>
Change-Id: I198dae2b53860f449f2fdbc71575babceed1ffcf
Diffstat (limited to 'tests/test_nn_tensorflow_utils.py')
-rw-r--r-- | tests/test_nn_tensorflow_utils.py | 8 |
1 files changed, 1 insertions, 7 deletions
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()) |