aboutsummaryrefslogtreecommitdiff
path: root/tests/test_nn_tensorflow_optimizations_pruning.py
diff options
context:
space:
mode:
authorGergely Nagy <gergely.nagy@arm.com>2023-11-21 12:29:38 +0000
committerGergely Nagy <gergely.nagy@arm.com>2023-12-07 17:09:31 +0000
commit54eec806272b7574a0757c77a913a369a9ecdc70 (patch)
tree2e6484b857b2a68279a2707dbb21e5c26685f4e2 /tests/test_nn_tensorflow_optimizations_pruning.py
parent7c50f1d6367186c03a282ac7ecb8fca0f905ba30 (diff)
downloadmlia-54eec806272b7574a0757c77a913a369a9ecdc70.tar.gz
MLIA-835 Invalid JSON output
TFLiteConverter was producing log messages in the output that was not possible to capture and redirect to logging. The solution/workaround is to run it as a subprocess. This change required some refactoring around existing invocations of the converter. Change-Id: I394bd0d49d36e6686cfcb9d658e4aad05326cb87 Signed-off-by: Gergely Nagy <gergely.nagy@arm.com>
Diffstat (limited to 'tests/test_nn_tensorflow_optimizations_pruning.py')
-rw-r--r--tests/test_nn_tensorflow_optimizations_pruning.py7
1 files changed, 3 insertions, 4 deletions
diff --git a/tests/test_nn_tensorflow_optimizations_pruning.py b/tests/test_nn_tensorflow_optimizations_pruning.py
index d97b3d3..9afc3ff 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, 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 optimizations/pruning."""
from __future__ import annotations
@@ -11,9 +11,8 @@ from numpy.core.numeric import isclose
from mlia.nn.tensorflow.optimizations.pruning import Pruner
from mlia.nn.tensorflow.optimizations.pruning import PruningConfiguration
+from mlia.nn.tensorflow.tflite_convert import convert_to_tflite
from mlia.nn.tensorflow.tflite_metrics import TFLiteMetrics
-from mlia.nn.tensorflow.utils import convert_to_tflite
-from mlia.nn.tensorflow.utils import save_tflite_model
from tests.utils.common import get_dataset
from tests.utils.common import train_model
@@ -52,7 +51,7 @@ def _get_tflite_metrics(
) -> TFLiteMetrics:
"""Save model as TFLiteModel and return metrics."""
temp_file = path / tflite_fn
- save_tflite_model(convert_to_tflite(model), temp_file)
+ convert_to_tflite(model, output_path=temp_file)
return TFLiteMetrics(str(temp_file))