From 446c379c92e15ad8f24ed0db853dd0fc9c271151 Mon Sep 17 00:00:00 2001 From: Ruomei Yan Date: Mon, 20 Feb 2023 15:32:54 +0000 Subject: Add a CLI component to enable rewrites * Add flags for rewrite (--rewrite, --rewrite-start, --rewrite-end, --rewrite-target) * Refactor CLI interfaces to accept tflite models with optimize for rewrite, keras models with optimize for clustering and pruning * Refactor and move common.py and select.py out of the folder nn/tensorflow/optimizations * Add file nn/rewrite/core/rewrite.py as placeholder * Update/add unit tests * Refactor OptimizeModel in ethos_u/data_collection.py for accepting tflite model case * Extend the logic so that if "--rewrite" is specified, we don't add pruning to also accept TFLite models. * Update README.md Resolves: MLIA-750, MLIA-854, MLIA-865 Signed-off-by: Benjamin Klimczak Change-Id: I67d85f71fa253d2bad4efe304ad8225970b9622c --- tests/test_cli_commands.py | 144 ++++++++++++-- tests/test_cli_helpers.py | 2 +- tests/test_cli_main.py | 15 ++ tests/test_cli_options.py | 7 + tests/test_nn_select.py | 241 +++++++++++++++++++++++ tests/test_nn_tensorflow_optimizations_select.py | 241 ----------------------- tests/test_target_ethos_u_advice_generation.py | 2 +- tests/test_target_ethos_u_advisor.py | 51 ++++- tests/test_target_ethos_u_data_analysis.py | 2 +- tests/test_target_ethos_u_data_collection.py | 2 +- 10 files changed, 441 insertions(+), 266 deletions(-) create mode 100644 tests/test_nn_select.py delete mode 100644 tests/test_nn_tensorflow_optimizations_select.py (limited to 'tests') diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py index f3213c4..6765a53 100644 --- a/tests/test_cli_commands.py +++ b/tests/test_cli_commands.py @@ -3,6 +3,7 @@ """Tests for cli.commands module.""" from __future__ import annotations +from contextlib import ExitStack as does_not_raise from pathlib import Path from typing import Any from unittest.mock import call @@ -49,37 +50,148 @@ def test_performance_unknown_target( @pytest.mark.parametrize( - "target_profile, pruning, clustering, pruning_target, clustering_target", + "target_profile, pruning, clustering, pruning_target, clustering_target, " + "rewrite, rewrite_target, rewrite_start, rewrite_end, expected_error", [ - ["ethos-u55-256", True, False, 0.5, None], - ["ethos-u65-512", False, True, 0.5, 32], - ["ethos-u55-256", True, True, 0.5, None], - ["ethos-u55-256", False, False, 0.5, None], - ["ethos-u55-256", False, True, "invalid", 32], + [ + "ethos-u55-256", + True, + False, + 0.5, + None, + False, + None, + "node_a", + "node_b", + does_not_raise(), + ], + [ + "ethos-u55-256", + False, + False, + None, + None, + True, + "fully_connected", + "node_a", + "node_b", + does_not_raise(), + ], + [ + "ethos-u55-256", + True, + False, + 0.5, + None, + True, + "fully_connected", + "node_a", + "node_b", + pytest.raises( + Exception, + match=(r"Only 'rewrite' is supported for TensorFlow Lite files."), + ), + ], + [ + "ethos-u65-512", + False, + True, + 0.5, + 32, + False, + None, + None, + None, + does_not_raise(), + ], + [ + "ethos-u55-256", + False, + False, + 0.5, + None, + True, + "random", + "node_x", + "node_y", + pytest.raises( + Exception, + match=(r"Currently only remove and fully_connected are supported."), + ), + ], + [ + "ethos-u55-256", + False, + False, + 0.5, + None, + True, + None, + "node_m", + "node_n", + pytest.raises( + Exception, + match=( + r"To perform rewrite, rewrite-target, " + r"rewrite-start and rewrite-end must be set." + ), + ), + ], + [ + "ethos-u55-256", + False, + False, + "invalid", + None, + True, + "remove", + None, + "node_end", + pytest.raises( + Exception, + match=( + r"To perform rewrite, rewrite-target, " + r"rewrite-start and rewrite-end must be set." + ), + ), + ], ], ) -def test_opt_valid_optimization_target( +def test_opt_valid_optimization_target( # pylint: disable=too-many-arguments target_profile: str, sample_context: ExecutionContext, pruning: bool, clustering: bool, pruning_target: float | None, clustering_target: int | None, + rewrite: bool, + rewrite_target: str | None, + rewrite_start: str | None, + rewrite_end: str | None, + expected_error: Any, monkeypatch: pytest.MonkeyPatch, test_keras_model: Path, + test_tflite_model: Path, ) -> None: """Test that command should not fail with valid optimization targets.""" mock_performance_estimation(monkeypatch) - optimize( - ctx=sample_context, - target_profile=target_profile, - model=str(test_keras_model), - pruning=pruning, - clustering=clustering, - pruning_target=pruning_target, - clustering_target=clustering_target, - ) + model_type = test_tflite_model if rewrite else test_keras_model + + with expected_error: + optimize( + ctx=sample_context, + target_profile=target_profile, + model=str(model_type), + pruning=pruning, + clustering=clustering, + pruning_target=pruning_target, + clustering_target=clustering_target, + rewrite=rewrite, + rewrite_target=rewrite_target, + rewrite_start=rewrite_start, + rewrite_end=rewrite_end, + ) def mock_performance_estimation(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/tests/test_cli_helpers.py b/tests/test_cli_helpers.py index 6d19207..494ed89 100644 --- a/tests/test_cli_helpers.py +++ b/tests/test_cli_helpers.py @@ -10,7 +10,7 @@ import pytest from mlia.cli.helpers import CLIActionResolver from mlia.cli.helpers import copy_profile_file_to_output_dir -from mlia.nn.tensorflow.optimizations.select import OptimizationSettings +from mlia.nn.select import OptimizationSettings class TestCliActionResolver: diff --git a/tests/test_cli_main.py b/tests/test_cli_main.py index 2f89268..e415284 100644 --- a/tests/test_cli_main.py +++ b/tests/test_cli_main.py @@ -165,6 +165,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable: pruning_target=None, clustering_target=None, backend=None, + rewrite=False, + rewrite_target=None, + rewrite_start=None, + rewrite_end=None, + dataset=None, ), ], [ @@ -189,6 +194,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable: pruning_target=0.5, clustering_target=32, backend=None, + rewrite=False, + rewrite_target=None, + rewrite_start=None, + rewrite_end=None, + dataset=None, ), ], [ @@ -210,6 +220,11 @@ def wrap_mock_command(mock: MagicMock, command: Callable) -> Callable: pruning_target=None, clustering_target=None, backend=["some_backend"], + rewrite=False, + rewrite_target=None, + rewrite_start=None, + rewrite_end=None, + dataset=None, ), ], [ diff --git a/tests/test_cli_options.py b/tests/test_cli_options.py index c02ef89..75ace0b 100644 --- a/tests/test_cli_options.py +++ b/tests/test_cli_options.py @@ -30,6 +30,7 @@ from mlia.core.typing import OutputFormat "optimization_type": "pruning", "optimization_target": 0.5, "layers_to_optimize": None, + "dataset": None, } ], ], @@ -44,6 +45,7 @@ from mlia.core.typing import OutputFormat "optimization_type": "pruning", "optimization_target": 0.5, "layers_to_optimize": None, + "dataset": None, } ], ], @@ -58,6 +60,7 @@ from mlia.core.typing import OutputFormat "optimization_type": "clustering", "optimization_target": 32, "layers_to_optimize": None, + "dataset": None, } ], ], @@ -72,11 +75,13 @@ from mlia.core.typing import OutputFormat "optimization_type": "pruning", "optimization_target": 0.5, "layers_to_optimize": None, + "dataset": None, }, { "optimization_type": "clustering", "optimization_target": 32, "layers_to_optimize": None, + "dataset": None, }, ], ], @@ -91,6 +96,7 @@ from mlia.core.typing import OutputFormat "optimization_type": "pruning", "optimization_target": 0.4, "layers_to_optimize": None, + "dataset": None, } ], ], @@ -117,6 +123,7 @@ from mlia.core.typing import OutputFormat "optimization_type": "clustering", "optimization_target": 32.2, "layers_to_optimize": None, + "dataset": None, } ], ], diff --git a/tests/test_nn_select.py b/tests/test_nn_select.py new file mode 100644 index 0000000..31628d2 --- /dev/null +++ b/tests/test_nn_select.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for module select.""" +from __future__ import annotations + +from contextlib import ExitStack as does_not_raise +from pathlib import Path +from typing import Any + +import pytest +import tensorflow as tf + +from mlia.core.errors import ConfigurationError +from mlia.nn.select import get_optimizer +from mlia.nn.select import MultiStageOptimizer +from mlia.nn.select import OptimizationSettings +from mlia.nn.tensorflow.optimizations.clustering import Clusterer +from mlia.nn.tensorflow.optimizations.clustering import ClusteringConfiguration +from mlia.nn.tensorflow.optimizations.pruning import Pruner +from mlia.nn.tensorflow.optimizations.pruning import PruningConfiguration + + +@pytest.mark.parametrize( + "config, expected_error, expected_type, expected_config", + [ + ( + OptimizationSettings( + optimization_type="pruning", + optimization_target=0.5, + layers_to_optimize=None, + ), + does_not_raise(), + Pruner, + "pruning: 0.5", + ), + ( + PruningConfiguration(0.5), + does_not_raise(), + Pruner, + "pruning: 0.5", + ), + ( + OptimizationSettings( + optimization_type="clustering", + optimization_target=32, + layers_to_optimize=None, + ), + does_not_raise(), + Clusterer, + "clustering: 32", + ), + ( + OptimizationSettings( + optimization_type="clustering", + optimization_target=0.5, + layers_to_optimize=None, + ), + pytest.raises( + ConfigurationError, + match="Optimization target should be a " + "positive integer. " + "Optimization target provided: 0.5", + ), + None, + None, + ), + ( + ClusteringConfiguration(32), + does_not_raise(), + Clusterer, + "clustering: 32", + ), + ( + OptimizationSettings( + optimization_type="superoptimization", + optimization_target="supertarget", # type: ignore + layers_to_optimize="all", # type: ignore + ), + pytest.raises( + ConfigurationError, + match="Unsupported optimization type: superoptimization", + ), + None, + None, + ), + ( + OptimizationSettings( + optimization_type="", + optimization_target=0.5, + layers_to_optimize=None, + ), + pytest.raises( + ConfigurationError, + match="Optimization type is not provided", + ), + None, + None, + ), + ( + "wrong_config", + pytest.raises( + Exception, + match="Unknown optimization configuration wrong_config", + ), + None, + None, + ), + ( + OptimizationSettings( + optimization_type="pruning", + optimization_target=None, # type: ignore + layers_to_optimize=None, + ), + pytest.raises( + Exception, + match="Optimization target is not provided", + ), + None, + None, + ), + ( + [ + OptimizationSettings( + optimization_type="pruning", + optimization_target=0.5, + layers_to_optimize=None, + ), + OptimizationSettings( + optimization_type="clustering", + optimization_target=32, + layers_to_optimize=None, + ), + ], + does_not_raise(), + MultiStageOptimizer, + "pruning: 0.5 - clustering: 32", + ), + ], +) +def test_get_optimizer( + config: Any, + expected_error: Any, + expected_type: type, + expected_config: str, + test_keras_model: Path, +) -> None: + """Test function get_optimzer.""" + model = tf.keras.models.load_model(str(test_keras_model)) + + with expected_error: + optimizer = get_optimizer(model, config) + assert isinstance(optimizer, expected_type) + assert optimizer.optimization_config() == expected_config + + +@pytest.mark.parametrize( + "params, expected_result", + [ + ( + [], + [], + ), + ( + [("pruning", 0.5)], + [ + OptimizationSettings( + optimization_type="pruning", + optimization_target=0.5, + layers_to_optimize=None, + ) + ], + ), + ( + [("pruning", 0.5), ("clustering", 32)], + [ + OptimizationSettings( + optimization_type="pruning", + optimization_target=0.5, + layers_to_optimize=None, + ), + OptimizationSettings( + optimization_type="clustering", + optimization_target=32, + layers_to_optimize=None, + ), + ], + ), + ], +) +def test_optimization_settings_create_from( + params: list[tuple[str, float]], expected_result: list[OptimizationSettings] +) -> None: + """Test creating settings from parsed params.""" + assert OptimizationSettings.create_from(params) == expected_result + + +@pytest.mark.parametrize( + "settings, expected_next_target, expected_error", + [ + [ + OptimizationSettings("clustering", 32, None), + OptimizationSettings("clustering", 16, None), + does_not_raise(), + ], + [ + OptimizationSettings("clustering", 4, None), + OptimizationSettings("clustering", 4, None), + does_not_raise(), + ], + [ + OptimizationSettings("clustering", 10, None), + OptimizationSettings("clustering", 8, None), + does_not_raise(), + ], + [ + OptimizationSettings("pruning", 0.5, None), + OptimizationSettings("pruning", 0.6, None), + does_not_raise(), + ], + [ + OptimizationSettings("pruning", 0.9, None), + OptimizationSettings("pruning", 0.9, None), + does_not_raise(), + ], + [ + OptimizationSettings("super_optimization", 42, None), + None, + pytest.raises( + Exception, match="Optimization type super_optimization is unknown." + ), + ], + ], +) +def test_optimization_settings_next_target( + settings: OptimizationSettings, + expected_next_target: OptimizationSettings, + expected_error: Any, +) -> None: + """Test getting next optimization target.""" + with expected_error: + assert settings.next_target() == expected_next_target diff --git a/tests/test_nn_tensorflow_optimizations_select.py b/tests/test_nn_tensorflow_optimizations_select.py deleted file mode 100644 index f5ba6f0..0000000 --- a/tests/test_nn_tensorflow_optimizations_select.py +++ /dev/null @@ -1,241 +0,0 @@ -# SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates. -# SPDX-License-Identifier: Apache-2.0 -"""Tests for module select.""" -from __future__ import annotations - -from contextlib import ExitStack as does_not_raise -from pathlib import Path -from typing import Any - -import pytest -import tensorflow as tf - -from mlia.core.errors import ConfigurationError -from mlia.nn.tensorflow.optimizations.clustering import Clusterer -from mlia.nn.tensorflow.optimizations.clustering import ClusteringConfiguration -from mlia.nn.tensorflow.optimizations.pruning import Pruner -from mlia.nn.tensorflow.optimizations.pruning import PruningConfiguration -from mlia.nn.tensorflow.optimizations.select import get_optimizer -from mlia.nn.tensorflow.optimizations.select import MultiStageOptimizer -from mlia.nn.tensorflow.optimizations.select import OptimizationSettings - - -@pytest.mark.parametrize( - "config, expected_error, expected_type, expected_config", - [ - ( - OptimizationSettings( - optimization_type="pruning", - optimization_target=0.5, - layers_to_optimize=None, - ), - does_not_raise(), - Pruner, - "pruning: 0.5", - ), - ( - PruningConfiguration(0.5), - does_not_raise(), - Pruner, - "pruning: 0.5", - ), - ( - OptimizationSettings( - optimization_type="clustering", - optimization_target=32, - layers_to_optimize=None, - ), - does_not_raise(), - Clusterer, - "clustering: 32", - ), - ( - OptimizationSettings( - optimization_type="clustering", - optimization_target=0.5, - layers_to_optimize=None, - ), - pytest.raises( - ConfigurationError, - match="Optimization target should be a " - "positive integer. " - "Optimization target provided: 0.5", - ), - None, - None, - ), - ( - ClusteringConfiguration(32), - does_not_raise(), - Clusterer, - "clustering: 32", - ), - ( - OptimizationSettings( - optimization_type="superoptimization", - optimization_target="supertarget", # type: ignore - layers_to_optimize="all", # type: ignore - ), - pytest.raises( - ConfigurationError, - match="Unsupported optimization type: superoptimization", - ), - None, - None, - ), - ( - OptimizationSettings( - optimization_type="", - optimization_target=0.5, - layers_to_optimize=None, - ), - pytest.raises( - ConfigurationError, - match="Optimization type is not provided", - ), - None, - None, - ), - ( - "wrong_config", - pytest.raises( - Exception, - match="Unknown optimization configuration wrong_config", - ), - None, - None, - ), - ( - OptimizationSettings( - optimization_type="pruning", - optimization_target=None, # type: ignore - layers_to_optimize=None, - ), - pytest.raises( - Exception, - match="Optimization target is not provided", - ), - None, - None, - ), - ( - [ - OptimizationSettings( - optimization_type="pruning", - optimization_target=0.5, - layers_to_optimize=None, - ), - OptimizationSettings( - optimization_type="clustering", - optimization_target=32, - layers_to_optimize=None, - ), - ], - does_not_raise(), - MultiStageOptimizer, - "pruning: 0.5 - clustering: 32", - ), - ], -) -def test_get_optimizer( - config: Any, - expected_error: Any, - expected_type: type, - expected_config: str, - test_keras_model: Path, -) -> None: - """Test function get_optimzer.""" - model = tf.keras.models.load_model(str(test_keras_model)) - - with expected_error: - optimizer = get_optimizer(model, config) - assert isinstance(optimizer, expected_type) - assert optimizer.optimization_config() == expected_config - - -@pytest.mark.parametrize( - "params, expected_result", - [ - ( - [], - [], - ), - ( - [("pruning", 0.5)], - [ - OptimizationSettings( - optimization_type="pruning", - optimization_target=0.5, - layers_to_optimize=None, - ) - ], - ), - ( - [("pruning", 0.5), ("clustering", 32)], - [ - OptimizationSettings( - optimization_type="pruning", - optimization_target=0.5, - layers_to_optimize=None, - ), - OptimizationSettings( - optimization_type="clustering", - optimization_target=32, - layers_to_optimize=None, - ), - ], - ), - ], -) -def test_optimization_settings_create_from( - params: list[tuple[str, float]], expected_result: list[OptimizationSettings] -) -> None: - """Test creating settings from parsed params.""" - assert OptimizationSettings.create_from(params) == expected_result - - -@pytest.mark.parametrize( - "settings, expected_next_target, expected_error", - [ - [ - OptimizationSettings("clustering", 32, None), - OptimizationSettings("clustering", 16, None), - does_not_raise(), - ], - [ - OptimizationSettings("clustering", 4, None), - OptimizationSettings("clustering", 4, None), - does_not_raise(), - ], - [ - OptimizationSettings("clustering", 10, None), - OptimizationSettings("clustering", 8, None), - does_not_raise(), - ], - [ - OptimizationSettings("pruning", 0.5, None), - OptimizationSettings("pruning", 0.6, None), - does_not_raise(), - ], - [ - OptimizationSettings("pruning", 0.9, None), - OptimizationSettings("pruning", 0.9, None), - does_not_raise(), - ], - [ - OptimizationSettings("super_optimization", 42, None), - None, - pytest.raises( - Exception, match="Optimization type super_optimization is unknown." - ), - ], - ], -) -def test_optimization_settings_next_target( - settings: OptimizationSettings, - expected_next_target: OptimizationSettings, - expected_error: Any, -) -> None: - """Test getting next optimization target.""" - with expected_error: - assert settings.next_target() == expected_next_target diff --git a/tests/test_target_ethos_u_advice_generation.py b/tests/test_target_ethos_u_advice_generation.py index 772fc56..ac4e5e9 100644 --- a/tests/test_target_ethos_u_advice_generation.py +++ b/tests/test_target_ethos_u_advice_generation.py @@ -12,7 +12,7 @@ from mlia.core.common import DataItem from mlia.core.context import ExecutionContext from mlia.core.helpers import ActionResolver from mlia.core.helpers import APIActionResolver -from mlia.nn.tensorflow.optimizations.select import OptimizationSettings +from mlia.nn.select import OptimizationSettings from mlia.target.ethos_u.advice_generation import EthosUAdviceProducer from mlia.target.ethos_u.advice_generation import EthosUStaticAdviceProducer from mlia.target.ethos_u.data_analysis import AllOperatorsSupportedOnNPU diff --git a/tests/test_target_ethos_u_advisor.py b/tests/test_target_ethos_u_advisor.py index 11aefc7..20131d2 100644 --- a/tests/test_target_ethos_u_advisor.py +++ b/tests/test_target_ethos_u_advisor.py @@ -1,7 +1,11 @@ # SPDX-FileCopyrightText: Copyright 2022-2023, Arm Limited and/or its affiliates. # SPDX-License-Identifier: Apache-2.0 """Tests for Ethos-U MLIA module.""" +from __future__ import annotations + +from contextlib import ExitStack as does_not_raise from pathlib import Path +from typing import Any import pytest @@ -16,16 +20,53 @@ def test_advisor_metadata() -> None: assert EthosUInferenceAdvisor.name() == "ethos_u_inference_advisor" -def test_unsupported_advice_categories(tmp_path: Path, test_tflite_model: Path) -> None: +@pytest.mark.parametrize( + "optimization_targets, expected_error", + [ + [ + [ + { + "optimization_type": "pruning", + "optimization_target": 0.5, + "layers_to_optimize": None, + } + ], + pytest.raises( + Exception, + match="Only 'rewrite' is supported for TensorFlow Lite files.", + ), + ], + [ + [ + { + "optimization_type": "rewrite", + "optimization_target": "fully_connected", + "layers_to_optimize": [ + "MobileNet/avg_pool/AvgPool", + "MobileNet/fc1/BiasAdd", + ], + } + ], + does_not_raise(), + ], + ], +) +def test_unsupported_advice_categories( + tmp_path: Path, + test_tflite_model: Path, + optimization_targets: list[dict[str, Any]], + expected_error: Any, +) -> None: """Test that advisor should throw an exception for unsupported categories.""" - with pytest.raises( - Exception, match="Optimizations are not supported for TensorFlow Lite files." - ): + with expected_error: ctx = ExecutionContext( output_dir=tmp_path, advice_category={AdviceCategory.OPTIMIZATION} ) advisor = configure_and_get_ethosu_advisor( - ctx, "ethos-u55-256", str(test_tflite_model) + ctx, + "ethos-u55-256", + str(test_tflite_model), + optimization_targets=optimization_targets, ) advisor.configure(ctx) diff --git a/tests/test_target_ethos_u_data_analysis.py b/tests/test_target_ethos_u_data_analysis.py index 80f0603..713e8ef 100644 --- a/tests/test_target_ethos_u_data_analysis.py +++ b/tests/test_target_ethos_u_data_analysis.py @@ -12,7 +12,7 @@ from mlia.backend.vela.compat import Operator from mlia.backend.vela.compat import Operators from mlia.core.common import DataItem from mlia.core.data_analysis import Fact -from mlia.nn.tensorflow.optimizations.select import OptimizationSettings +from mlia.nn.select import OptimizationSettings from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityInfo from mlia.nn.tensorflow.tflite_compat import TFLiteCompatibilityStatus from mlia.nn.tensorflow.tflite_compat import TFLiteConversionError diff --git a/tests/test_target_ethos_u_data_collection.py b/tests/test_target_ethos_u_data_collection.py index fd824ae..6244f8b 100644 --- a/tests/test_target_ethos_u_data_collection.py +++ b/tests/test_target_ethos_u_data_collection.py @@ -10,7 +10,7 @@ from mlia.backend.vela.compat import Operators from mlia.core.context import Context from mlia.core.data_collection import DataCollector from mlia.core.errors import FunctionalityNotSupportedError -from mlia.nn.tensorflow.optimizations.select import OptimizationSettings +from mlia.nn.select import OptimizationSettings from mlia.target.ethos_u.config import EthosUConfiguration from mlia.target.ethos_u.data_collection import EthosUOperatorCompatibility from mlia.target.ethos_u.data_collection import EthosUOptimizationPerformance -- cgit v1.2.1