diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py b/TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py index 44e9c699250..a187ef17d3e 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py @@ -43,6 +43,7 @@ import os from typing import Dict, List, Union, Tuple, Optional import json +import warnings import numpy as np import onnx @@ -58,7 +59,7 @@ from aimet_common import libquant_info from aimet_common.defs import QuantScheme, QuantizationDataType from aimet_common.quantsim import extract_global_quantizer_args, VALID_ENCODING_VERSIONS -from aimet_common.utils import save_json_yaml, AimetLogger +from aimet_common.utils import save_json_yaml, AimetLogger, _red from aimet_common.connected_graph.product import Product from aimet_onnx import utils from aimet_onnx.meta.operations import Op @@ -771,6 +772,15 @@ def export(self, path: str, filename_prefix: str): :param path: dir to save encoding files :param filename_prefix: filename to save encoding files """ + if quantsim.encoding_version == '0.6.1': + msg = _red("Encoding version 0.6.1 will be deprecated in a future release, with version 1.0.0 becoming " + "the default. If your code depends on parsing the exported encodings file, ensure that it is " + "updated to be able to parse 1.0.0 format.\n" + "To swap the encoding version to 1.0.0, run the following lines prior to calling quantsim " + "export:\n\n" + "from aimet_common import quantsim\n" + "quantsim.encoding_version = '1.0.0'") + warnings.warn(msg, DeprecationWarning, stacklevel=2) self._export_encodings(os.path.join(path, filename_prefix) + '.encodings') self.remove_quantization_nodes() if self.model.model.ByteSize() >= onnx.checker.MAXIMUM_PROTOBUF: diff --git a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quantsim.py b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quantsim.py index b395f925639..c7b0ce00419 100644 --- a/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quantsim.py +++ b/TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/quantsim.py @@ -37,6 +37,7 @@ """ Quantsim for Keras """ from __future__ import annotations +import contextlib from dataclasses import dataclass import json import os @@ -47,7 +48,8 @@ from aimet_common.defs import QuantScheme, QuantizationDataType from aimet_common.utils import AimetLogger, save_json_yaml -from aimet_common.quantsim import encoding_version, extract_global_quantizer_args +from aimet_common import quantsim +from aimet_common.quantsim import extract_global_quantizer_args from aimet_tensorflow.keras.connectedgraph import ConnectedGraph from aimet_tensorflow.keras.graphsearchtuils import GraphSearchUtils from aimet_tensorflow.keras.quant_sim.qc_quantize_wrapper import QcQuantizeWrapper, QuantizerSettings @@ -424,7 +426,7 @@ def get_encodings_dict(self) -> Dict[str, Union[str, Dict]]: encoding_dict = self._get_encoding_dict_for_quantizer(output_quantizer) activation_encodings[tensor_name] = encoding_dict return { - 'version': encoding_version, + 'version': quantsim.encoding_version, 'activation_encodings': activation_encodings, 'param_encodings': param_encodings, 'quantizer_args': self.quant_args if hasattr(self, "quant_args") else {} @@ -477,6 +479,20 @@ def _set_op_mode_parameters(self, op_mode: libpymo.TensorQuantizerOpMode): if param_quantizer.is_enabled(): param_quantizer.quant_mode = op_mode + @staticmethod + @contextlib.contextmanager + def _set_encoding_version_to_0_6_1(): + assert quantsim.encoding_version in {'0.6.1', '1.0.0'} + if quantsim.encoding_version == '1.0.0': + _logger.info('Exporting to encoding version 1.0.0 is not yet supported. Exporting using version 0.6.1 ' + 'instead.') + old_encoding_version = quantsim.encoding_version + quantsim.encoding_version = '0.6.1' + + yield + + quantsim.encoding_version = old_encoding_version + def export(self, path, filename_prefix, custom_objects=None, convert_to_pb=True): """ This method exports out the quant-sim model so it is ready to be run on-target. @@ -488,29 +504,30 @@ def export(self, path, filename_prefix, custom_objects=None, convert_to_pb=True) :param filename_prefix: Prefix to use for filenames of the model pth and encodings files :param custom_objects: If there are custom objects to load, Keras needs a dict of them to map them """ - model_path = os.path.join(path, filename_prefix) + with self._set_encoding_version_to_0_6_1(): + model_path = os.path.join(path, filename_prefix) - #TF Version 2.4 has bug i.e. save() in tf format don't work for unrolled LSTM. - for layer in self._model_without_wrappers.layers: - if isinstance(layer, tf.keras.layers.LSTM): - break - else: - self._model_without_wrappers.save(model_path) - - self._model_without_wrappers.save(model_path + '.h5', save_format='h5') - - # Conversion of saved h5 model to pb model for consumption by SNPE/QNN - try: - if convert_to_pb: - convert_h5_model_to_pb_model(f'{model_path}.h5', custom_objects=custom_objects) - except ValueError: - _logger.error("Could not convert h5 to frozen pb. " - "Please call export() again with custom_objects defined.") - raise - finally: - encodings_dict = self.get_encodings_dict() - encoding_file_path = os.path.join(path, filename_prefix + '.encodings') - save_json_yaml(encoding_file_path, encodings_dict) + #TF Version 2.4 has bug i.e. save() in tf format don't work for unrolled LSTM. + for layer in self._model_without_wrappers.layers: + if isinstance(layer, tf.keras.layers.LSTM): + break + else: + self._model_without_wrappers.save(model_path) + + self._model_without_wrappers.save(model_path + '.h5', save_format='h5') + + # Conversion of saved h5 model to pb model for consumption by SNPE/QNN + try: + if convert_to_pb: + convert_h5_model_to_pb_model(f'{model_path}.h5', custom_objects=custom_objects) + except ValueError: + _logger.error("Could not convert h5 to frozen pb. " + "Please call export() again with custom_objects defined.") + raise + finally: + encodings_dict = self.get_encodings_dict() + encoding_file_path = os.path.join(path, filename_prefix + '.encodings') + save_json_yaml(encoding_file_path, encodings_dict) def _compute_and_set_parameter_encodings(self, ops_with_invalid_encodings: List): # pylint: disable=too-many-nested-blocks diff --git a/TrainingExtensions/tensorflow/test/python/test_quantsim_keras.py b/TrainingExtensions/tensorflow/test/python/test_quantsim_keras.py index ed26b1be761..ce4a85d915a 100644 --- a/TrainingExtensions/tensorflow/test/python/test_quantsim_keras.py +++ b/TrainingExtensions/tensorflow/test/python/test_quantsim_keras.py @@ -33,7 +33,7 @@ # # @@-COPYRIGHT-END-@@ # ============================================================================= - +import contextlib import json import os import tempfile @@ -50,6 +50,7 @@ import aimet_common.utils from aimet_common.defs import QuantScheme, RANGE_LEARNING_SCHEMES +from aimet_common import quantsim from aimet_tensorflow.examples.test_models import keras_model from aimet_tensorflow.keras.utils.quantizer_utils import SaveModelWithoutQuantsimWrappersCallback from aimet_tensorflow.keras.cross_layer_equalization import equalize_model @@ -1632,3 +1633,28 @@ def test_quantizable_lstm_export_encodings(): assert param_name in encodings['param_encodings'] assert encodings['param_encodings'][param_name] == encoding_dict +def test_quantsim_export_to_1_0_0(): + @contextlib.contextmanager + def _swap_encoding_version(): + old_version = quantsim.encoding_version + quantsim.encoding_version = '1.0.0' + + yield + + quantsim.encoding_version = old_version + + model = dense_functional() + rand_inp = np.random.randn(100, 5) + + + qsim = QuantizationSimModel(model, quant_scheme='tf') + qsim.compute_encodings(lambda m, _: m(rand_inp), None) + + with tempfile.TemporaryDirectory() as temp_dir, _swap_encoding_version(): + assert quantsim.encoding_version == '1.0.0' + qsim.export(temp_dir, 'test_export') + assert quantsim.encoding_version == '1.0.0' + + with open(os.path.join(temp_dir, 'test_export.encodings'), 'r') as encoding_file: + encodings = json.load(encoding_file) + assert encodings['version'] == '0.6.1' diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v1/quantsim.py b/TrainingExtensions/torch/src/python/aimet_torch/v1/quantsim.py index 0e85d9ab551..487fb23d52e 100644 --- a/TrainingExtensions/torch/src/python/aimet_torch/v1/quantsim.py +++ b/TrainingExtensions/torch/src/python/aimet_torch/v1/quantsim.py @@ -46,6 +46,7 @@ from typing import Tuple, List, Union, Dict, Callable, Optional, Any, runtime_checkable, Protocol, Mapping from collections import OrderedDict, defaultdict import json +import warnings import torch import onnx from packaging import version # pylint: disable=wrong-import-order @@ -59,6 +60,7 @@ from aimet_common.defs import QuantScheme, QuantizationDataType, SupportedKernelsAction, QuantDtypeBwInfo from aimet_common.quantsim import validate_quantsim_inputs, extract_global_quantizer_args, VALID_ENCODING_VERSIONS from aimet_common.quant_utils import get_conv_accum_bounds +from aimet_common.utils import deprecated, _red from aimet_torch.v1.nn.modules.custom import MatMul from aimet_torch.quantsim_config.quantsim_config import QuantSimConfigurator @@ -67,7 +69,6 @@ from aimet_torch.tensor_quantizer import initialize_learned_grid_quantizer_attributes, TensorQuantizer from aimet_torch.qc_quantize_op import get_encoding_by_quantizer as _get_encoding_by_quantizer from aimet_torch import torchscript_utils, utils, onnx_utils -from aimet_torch.utils import deprecated from aimet_torch.onnx_utils import ( OnnxSaver, OnnxExportApiArgs, @@ -520,7 +521,15 @@ def export(self, path: str, filename_prefix: str, dummy_input: Union[torch.Tenso :param filename_prefix_encodings: File name prefix to be used when saving encodings. If None, then user defaults to filename_prefix value """ - + if quantsim.encoding_version == '0.6.1': + msg = _red("Encoding version 0.6.1 will be deprecated in a future release, with version 1.0.0 becoming " + "the default. If your code depends on parsing the exported encodings file, ensure that it is " + "updated to be able to parse 1.0.0 format.\n" + "To swap the encoding version to 1.0.0, run the following lines prior to calling quantsim " + "export:\n\n" + "from aimet_common import quantsim\n" + "quantsim.encoding_version = '1.0.0'") + warnings.warn(msg, DeprecationWarning, stacklevel=2) warning_str = 'Exporting encodings to yaml will be deprecated in a future release. Ensure that your ' \ 'code can work with the exported files ending in ".encodings" which are saved using json ' \ 'format. For the time being, if yaml export is needed, set aimet_common.utils.SAVE_TO_YAML to ' \ diff --git a/TrainingExtensions/torch/test/python/test_quantizer.py b/TrainingExtensions/torch/test/python/test_quantizer.py index a80328b0949..e57494f2c47 100644 --- a/TrainingExtensions/torch/test/python/test_quantizer.py +++ b/TrainingExtensions/torch/test/python/test_quantizer.py @@ -788,7 +788,7 @@ def test_add_quantization_wrappers_with_modulelist_with_layers_to_ignore(self): assert 'layers_deep.5.0' in sim._excluded_layer_names assert 'layers_deep.5.1' in sim._excluded_layer_names - with tempfile.TemporaryDirectory() as tmpdir: + with tempfile.TemporaryDirectory() as tmpdir, pytest.warns(DeprecationWarning): sim.export(tmpdir, 'modulelist_with_layers_to_ignore', dummy_input=torch.rand(1, 3, 12, 12)) with open(os.path.join(tmpdir, "modulelist_with_layers_to_ignore.encodings"), "r") as encodings_file: encodings = json.load(encodings_file)