diff options
| -rw-r--r-- | README.md | 23 | ||||
| -rw-r--r-- | src/mlia/nn/rewrite/core/graph_edit/record.py | 10 | ||||
| -rw-r--r-- | src/mlia/nn/rewrite/core/train.py | 33 | ||||
| -rw-r--r-- | src/mlia/resources/optimization_profiles/optimization.toml | 1 | ||||
| -rw-r--r-- | src/mlia/resources/optimization_profiles/optimization_custom_augmentation.toml | 13 | ||||
| -rw-r--r-- | src/mlia/target/common/optimization.py | 55 | ||||
| -rw-r--r-- | tests/conftest.py | 10 | ||||
| -rw-r--r-- | tests/test_common_optimization.py | 58 | ||||
| -rw-r--r-- | tox.ini | 2 |
9 files changed, 181 insertions, 24 deletions
@@ -215,9 +215,26 @@ Training parameters for rewrites can be specified. There are a number of predefined profiles: -| Name | Batch Size | LR | Show Progress | Steps | LR Schedule | Num Procs | Num Threads | Checkpoints | -| :----------: | :--------: | :--: | :-----------: | :---: | :---------: | :-------: | :---------: | :---------: | -| optimization | 32 | 1e-3 | True | 48000 | "cosine" | 1 | 0 | None | +| Name | Batch Size | LR | Show Progress | Steps | LR Schedule | Num Procs | Num Threads | Checkpoints | Augmentations | +| :----------: | :--------: | :--: | :-----------: | :---: | :---------: | :-------: | :---------: | :---------: | :-------------: | +| optimization | 32 | 1e-3 | True | 48000 | "cosine" | 1 | 0 | None | "gaussian" | + +| Name | Batch Size | LR | Show Progress | Steps | LR Schedule | Num Procs | Num Threads | Checkpoints | Augmentations - gaussian_strength | Augmentations - mixup_strength | +| :------------------------------: | :--------: | :--: | :-----------: | :---: | :---------: | :-------: | :---------: | :---------: | :-------------------------------: | :----------------------------: | +| optimization_custom_augmentation | 32 | 1e-3 | True | 48000 | "cosine" | 1 | 0 | None | 0.1 | 0.1 | + +The augmentations consist of 2 parameters: mixup strength and gaussian strength. + +Augmenations can be selected from a number of pre-defined profiles (see the table below) or each individual parameter can be chosen (see optimization_custom_augmentation above for an example): + +| Name | MixUp Strength | Gaussian Strength | +| :------------------: | :------------: | :---------------: | +| "none" | None | None | +| "gaussian" | None | 1.0 | +| "mixup" | 1.0 | None | +| "mixout" | 1.6 | None | +| "mix_gaussian_large" | 2.0 | 1.0 | +| "mix_gaussian_small" | 1.6 | 0.3 | ```bash ##### An example for using optimization Profiles diff --git a/src/mlia/nn/rewrite/core/graph_edit/record.py b/src/mlia/nn/rewrite/core/graph_edit/record.py index f85433d..7d9f219 100644 --- a/src/mlia/nn/rewrite/core/graph_edit/record.py +++ b/src/mlia/nn/rewrite/core/graph_edit/record.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 """Save subgraph data.""" # pylint: disable=too-many-locals @@ -32,7 +32,7 @@ def dequantized_path(filename: str | Path) -> Path: return path -def record_model( +def record_model( # pylint: disable=too-many-arguments input_filename: str | Path, model_filename: str | Path, output_filename: str | Path, @@ -41,6 +41,7 @@ def record_model( num_procs: int = 1, num_threads: int = 0, dequantize_output: bool = False, + quantize_input: bool = False, ) -> None: """Model recorder. @@ -92,7 +93,10 @@ def record_model( for _, named_x in enumerate( track(dataset.as_numpy_iterator(), total=total, disable=not show_progress) ): - named_y = model(named_x) + if quantize_input: + named_y = model(model.quantize_inputs(named_x)) + else: + named_y = model(named_x) write(writer, named_y) if dequantize_output: diff --git a/src/mlia/nn/rewrite/core/train.py b/src/mlia/nn/rewrite/core/train.py index 88efa23..4204978 100644 --- a/src/mlia/nn/rewrite/core/train.py +++ b/src/mlia/nn/rewrite/core/train.py @@ -62,7 +62,7 @@ LEARNING_RATE_SCHEDULES = get_args(LearningRateSchedule) class TrainingParameters: """Define default parameters for the training.""" - augmentations: tuple[float | None, float | None] = AUGMENTATION_PRESETS["gaussian"] + augmentations: tuple[float | None, float | None] = AUGMENTATION_PRESETS["none"] batch_size: int = 32 steps: int = 48000 learning_rate: float = 1e-3 @@ -147,7 +147,8 @@ def train( # pylint: disable=too-many-arguments # Assess the output diff between the parts after the rewrite subgraph # in original and optimized model optimized_end_path = Path(train_dir, "optimized_end.tfrec") - end_path = Path(train_dir, "end.tfrec") + optimized_end_path_dequant = Path(train_dir, "optimized_end_dequant.tfrec") + end_path = Path(train_dir, "end_dequant.tfrec") record_model( str(input_tfrec), @@ -155,8 +156,10 @@ def train( # pylint: disable=too-many-arguments optimized_end_path, num_procs=train_params.num_procs, num_threads=train_params.num_threads, + dequantize_output=True, ) - mae, nrmse = diff_stats(end_path, str(optimized_end_path)) + + mae, nrmse = diff_stats(end_path, optimized_end_path_dequant) if unmodified_model_dir: cast(tempfile.TemporaryDirectory, unmodified_model_dir).cleanup() @@ -179,24 +182,27 @@ def eval_in_dir( model_input = ( model_input_path if model_input_path.exists() - else ExtractPaths.tfrec.input(target_dir, False) + else ExtractPaths.tfrec.input(target_dir, True) ) output = ( model_output_path if model_output_path.exists() - else ExtractPaths.tfrec.output(target_dir, False) + else ExtractPaths.tfrec.output(target_dir, True) ) with tempfile.TemporaryDirectory() as tmp_dir: predict = Path(tmp_dir, "predict.tfrec") + predict_dequant = Path(tmp_dir, "predict_dequant.tfrec") record_model( str(model_input), new_part, str(predict), num_procs=num_procs, num_threads=num_threads, + dequantize_output=True, + quantize_input=True, ) - mae, nrmse = diff_stats(str(output), str(predict)) + mae, nrmse = diff_stats(str(output), predict_dequant) return mae, nrmse @@ -249,7 +255,7 @@ def set_up_data_pipeline( augmentations: tuple[float | None, float | None], steps: int, batch_size: int = 32, -) -> tf.data.Dataset: +) -> tuple[tf.data.Dataset, int]: """Create a data pipeline for training of the replacement model.""" _check_model_compatibility(teacher, replace) @@ -340,7 +346,7 @@ def set_up_data_pipeline( dataset = dataset.map(restore_shapes) dataset = dataset.prefetch(tf.data.AUTOTUNE) - return dataset + return dataset, steps_per_epoch def train_in_dir( @@ -373,7 +379,7 @@ def train_in_dir( if model_is_quantized: replace.check_datatypes(np.int8) - dataset = set_up_data_pipeline( + dataset, steps_per_epoch = set_up_data_pipeline( teacher, replace, train_dir, @@ -453,6 +459,7 @@ def train_in_dir( input_shape, output_shape, loss_fn, + steps_per_epoch, post_process=True, ) @@ -492,6 +499,7 @@ def train_in_dir( input_shape, output_shape, loss_fn, + steps_per_epoch, ) # Placeholder for now, will be parametrized later (MLIA-1114) # rewrite.check_optimization( # type: ignore[attr-defined] @@ -548,6 +556,7 @@ def model_fit( # pylint: disable=too-many-arguments input_shape: int, output_shape: int, loss_fn: Callable, + steps_per_epoch: int, post_process: bool = False, ) -> keras.Model: """Train a tflite model.""" @@ -593,8 +602,12 @@ def model_fit( # pylint: disable=too-many-arguments model_to_save = model else: checkpoint_filename = str(output_filename) + logger.info("Evaluate final Keras Model using %d steps", steps_per_epoch) + model.evaluate( + dataset, + steps=steps_per_epoch, + ) model_to_save = model - with log_action( f"{steps_so_far}/{train_params.steps}: Saved as {checkpoint_filename}" ): diff --git a/src/mlia/resources/optimization_profiles/optimization.toml b/src/mlia/resources/optimization_profiles/optimization.toml index 623a763..42b64f0 100644 --- a/src/mlia/resources/optimization_profiles/optimization.toml +++ b/src/mlia/resources/optimization_profiles/optimization.toml @@ -7,5 +7,6 @@ learning_rate = 1e-3 show_progress = true steps = 48000 learning_rate_schedule = "cosine" +augmentations = "gaussian" num_procs = 1 num_threads = 0 diff --git a/src/mlia/resources/optimization_profiles/optimization_custom_augmentation.toml b/src/mlia/resources/optimization_profiles/optimization_custom_augmentation.toml new file mode 100644 index 0000000..5d1f917 --- /dev/null +++ b/src/mlia/resources/optimization_profiles/optimization_custom_augmentation.toml @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 + +[training] +batch_size = 32 +learning_rate = 1e-3 +show_progress = true +steps = 48000 +learning_rate_schedule = "cosine" +num_procs = 1 +num_threads = 0 +augmentations.gaussian_strength = 0.1 +augmentations.mixup_strength = 0.1 diff --git a/src/mlia/target/common/optimization.py b/src/mlia/target/common/optimization.py index 1423189..a139a7d 100644 --- a/src/mlia/target/common/optimization.py +++ b/src/mlia/target/common/optimization.py @@ -17,6 +17,7 @@ from mlia.core.errors import FunctionalityNotSupportedError from mlia.core.performance import estimate_performance from mlia.core.performance import P from mlia.core.performance import PerformanceEstimator +from mlia.nn.rewrite.core.train import AUGMENTATION_PRESETS from mlia.nn.select import get_optimizer from mlia.nn.select import OptimizationSettings from mlia.nn.tensorflow.config import get_keras_model @@ -218,7 +219,54 @@ _DEFAULT_OPTIMIZATION_TARGETS = [ ] -def add_common_optimization_params(advisor_parameters: dict, extra_args: dict) -> None: +def parse_augmentations( + augmentations: dict | str | None, +) -> tuple[float | None, float | None]: + """Parse augmentations from optimization-profile and return a valid tuple.""" + if isinstance(augmentations, str): + match_augmentation = AUGMENTATION_PRESETS.get(augmentations) + if not match_augmentation: + match_augmentation = AUGMENTATION_PRESETS["none"] + return match_augmentation + if isinstance(augmentations, dict): + augmentation_keys_test_for_valid = list(augmentations.keys()) + augmentation_keys_test_for_float = list(augmentations.keys()) + valid_keys = ["mixup_strength", "gaussian_strength"] + tuple_to_return = [] + for valid_key in valid_keys.copy(): + if augmentations.get(valid_key): + del augmentation_keys_test_for_valid[ + augmentation_keys_test_for_valid.index(valid_key) + ] + if isinstance(augmentations.get(valid_key), float): + tuple_to_return.append(augmentations[valid_key]) + del augmentation_keys_test_for_float[ + augmentation_keys_test_for_float.index(valid_key) + ] + else: + tuple_to_return.append(None) + else: + tuple_to_return.append(None) + + if len(augmentation_keys_test_for_valid) > 0: + logger.warning( + "Warning! Expected augmentation parameters to be 'gaussian_strength' " + "and/or 'mixup_strength' got %s. " + "Removing invalid augmentations", + str(list(augmentations.keys())), + ) + elif len(augmentation_keys_test_for_float) > 0: + logger.warning( + "Warning! Not all augmentation parameters were floats, " + "removing non-float augmentations" + ) + return (tuple_to_return[0], tuple_to_return[1]) + return AUGMENTATION_PRESETS["none"] + + +def add_common_optimization_params( # pylint: disable=too-many-branches + advisor_parameters: dict, extra_args: dict +) -> None: """Add common optimization parameters.""" optimization_targets = extra_args.get("optimization_targets") if not optimization_targets: @@ -234,6 +282,11 @@ def add_common_optimization_params(advisor_parameters: dict, extra_args: dict) - raise TypeError("Training Parameter values has wrong format.") training_parameters = extra_args["optimization_profile"].get("training") + if training_parameters: + training_parameters["augmentations"] = parse_augmentations( + training_parameters.get("augmentations") + ) + advisor_parameters.update( { "common_optimizations": { diff --git a/tests/conftest.py b/tests/conftest.py index 3d0b832..981bf3b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -257,17 +257,15 @@ def fixture_test_tfrecord_fp32( yield from create_tfrecord(tmp_path_factory, random_data) -@pytest.fixture(scope="session", autouse=True) +@pytest.fixture(scope="function", autouse=True) def set_training_steps( request: _pytest.fixtures.SubRequest, ) -> Generator[None, None, None]: """Speed up tests by using MockTrainingParameters.""" - if "set_training_steps" == request.fixturename: - yield - else: + if "skip_set_training_steps" not in request.keywords: with pytest.MonkeyPatch.context() as monkeypatch: monkeypatch.setattr( "mlia.nn.select._get_rewrite_params", - MagicMock(return_value=[MockTrainingParameters(), None, None]), + MagicMock(return_value=MockTrainingParameters()), ) - yield + yield diff --git a/tests/test_common_optimization.py b/tests/test_common_optimization.py index 341e0d2..58ea8af 100644 --- a/tests/test_common_optimization.py +++ b/tests/test_common_optimization.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates. # SPDX-License-Identifier: Apache-2.0 """Tests for the common optimization module.""" +from __future__ import annotations + from contextlib import ExitStack as does_not_raises from pathlib import Path from typing import Any @@ -15,6 +17,7 @@ from mlia.nn.tensorflow.config import TFLiteModel from mlia.target.common.optimization import _DEFAULT_OPTIMIZATION_TARGETS from mlia.target.common.optimization import add_common_optimization_params from mlia.target.common.optimization import OptimizingDataCollector +from mlia.target.common.optimization import parse_augmentations from mlia.target.config import load_profile from mlia.target.config import TargetProfile @@ -167,3 +170,58 @@ def test_add_common_optimization_params(extra_args: dict, error_to_raise: Any) - advisor_parameters["common_optimizations"]["training_parameters"] == extra_args["optimization_profile"]["training"] ) + + +@pytest.mark.parametrize( + "augmentations, expected_output", + [ + ( + {"gaussian_strength": 1.0, "mixup_strength": 1.0}, + (1.0, 1.0), + ), + ( + {"gaussian_strength": 1.0}, + (None, 1.0), + ), + ( + {"Wrong param": 1.0, "mixup_strength": 1.0}, + (1.0, None), + ), + ( + {"Wrong param1": 1.0, "Wrong param2": 1.0}, + (None, None), + ), + ( + "gaussian", + (None, 1.0), + ), + ( + "mix_gaussian_large", + (2.0, 1.0), + ), + ( + "not in presets", + (None, None), + ), + ( + {"gaussian_strength": 1.0, "mixup_strength": 1.0, "mix2": 1.0}, + (1.0, 1.0), + ), + ( + {"gaussian_strength": "not a float", "mixup_strength": 1.0}, + (1.0, None), + ), + ( + None, + (None, None), + ), + ], +) +def test_parse_augmentations( + augmentations: dict | str | None, expected_output: tuple +) -> None: + """Check that augmentation parameters in optimization_profiles are + correctly parsed.""" + + augmentation_output = parse_augmentations(augmentations) + assert augmentation_output == expected_output @@ -69,7 +69,7 @@ commands = description = Create the documentation. allowlist_externals = make deps = - Sphinx==4.5.0 + Sphinx==5.0.0 sphinx-rtd-theme==1.0.0 commands = sphinx-apidoc -f -o docs/source src/mlia |