diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 3165c5be8e..9d2f1ec3e5 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -5,22 +5,20 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- - -import os +import abc +import itertools import numpy as np -import onnx import onnxruntime + +import onnx from onnx import helper, TensorProto, ModelProto from onnx import onnx_pb as onnx_proto from enum import Enum +from pathlib import Path -from .quant_utils import QuantType, smooth_distribution, apply_plot +from .quant_utils import QuantType, model_has_infer_metadata, smooth_distribution, apply_plot, load_model, clone_model_with_shape_infer from .registry import QLinearOpsRegistry -import abc -import itertools - - class CalibrationMethod(Enum): MinMax = 0 Entropy = 1 @@ -47,7 +45,9 @@ class CalibraterBase: :param use_external_data_format: use external data format to store model which size is >= 2Gb ''' if isinstance(model, str): - self.model = onnx.load(model) + self.model = load_model(Path(model), False) + elif isinstance(model, Path): + self.model = load_model(model, False) elif isinstance(model, ModelProto): self.model = model else: @@ -172,9 +172,7 @@ class MinMaxCalibrater(CalibraterBase): model and ensures their outputs are stored as part of the graph output :return: augmented ONNX model ''' - model = onnx_proto.ModelProto() - model.CopyFrom(self.model) - model = onnx.shape_inference.infer_shapes(model) + model = clone_model_with_shape_infer(self.model) added_nodes = [] added_outputs = [] @@ -338,9 +336,7 @@ class HistogramCalibrater(CalibraterBase): make all quantization_candidates op type nodes as part of the graph output. :return: augmented ONNX model ''' - model = onnx_proto.ModelProto() - model.CopyFrom(self.model) - model = onnx.shape_inference.infer_shapes(model) + model = clone_model_with_shape_infer(self.model) added_nodes = [] added_outputs = [] diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index 2b2a97971e..000d6e7112 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -17,6 +17,7 @@ from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg, type_to_name from .quant_utils import quantize_nparray, quantize_data, compute_scale_zp, get_qrange_for_qType, get_qmin_qmax_for_qType +from .quant_utils import save_and_reload_model, model_has_infer_metadata, add_infer_metadata from .quant_utils import QuantType, onnx_domain, __producer__, __version__ from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer @@ -27,20 +28,22 @@ class ONNXQuantizer: def __init__(self, model, per_channel, reduce_range, mode, static, weight_qType, input_qType, tensors_range, nodes_to_quantize, nodes_to_exclude, op_types_to_quantize, extra_options={}): - # run shape inference on the model (enabled by default) - self.extra_options = extra_options if extra_options is not None else {} - if not ('DisableShapeInference' in self.extra_options and self.extra_options['DisableShapeInference']): - model = onnx.shape_inference.infer_shapes(model) + if not model_has_infer_metadata(model): + model = save_and_reload_model(model) self.value_infos = {vi.name: vi for vi in model.graph.value_info} self.value_infos.update({ot.name: ot for ot in model.graph.output}) self.value_infos.update({it.name: it for it in model.graph.input}) self.model = ONNXModel(model) + if not static: self.model.replace_gemm_with_matmul() + self.per_channel = per_channel # weight-pack per channel self.reduce_range = reduce_range self.mode = mode # QuantizationMode.Value self.static = static # use static quantization for inputs. self.fuse_dynamic_quant = False + + self.extra_options = extra_options if extra_options is not None else {} self.enable_subgraph_quantization = 'EnableSubgraph' in self.extra_options and self.extra_options['EnableSubgraph'] self.force_quantize_no_input_check = 'ForceQuantizeNoInputCheck' in self.extra_options and self.extra_options['ForceQuantizeNoInputCheck'] self.q_matmul_const_b_only = 'MatMulConstBOnly' in self.extra_options and self.extra_options['MatMulConstBOnly'] @@ -105,6 +108,7 @@ class ONNXQuantizer: ''' warped_model = onnx.helper.make_model(subgraph, producer_name='onnx-quantizer', opset_imports=self.model.model.opset_import) + add_infer_metadata(warped_model) sub_quanitzer = ONNXQuantizer(warped_model, self.per_channel, self.reduce_range, diff --git a/onnxruntime/python/tools/quantization/quant_utils.py b/onnxruntime/python/tools/quantization/quant_utils.py index 8cd0814dd3..01a3b7197b 100644 --- a/onnxruntime/python/tools/quantization/quant_utils.py +++ b/onnxruntime/python/tools/quantization/quant_utils.py @@ -1,11 +1,15 @@ import logging import numpy import onnx +import tempfile from enum import Enum from onnx import onnx_pb as onnx_proto +from onnx import external_data_helper from pathlib import Path +from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel + __producer__ = "onnx.quantize" __version__ = "0.1.0" onnx_domain = "ai.onnx" @@ -447,3 +451,71 @@ def smooth_distribution(p, eps=0.0001): assert (hist <= 0).sum() == 0 return hist + +def model_has_external_data(model_path : Path): + model = onnx.load(model_path.as_posix(), load_external_data=False) + for intializer in model.graph.initializer: + if external_data_helper.uses_external_data(intializer): + return True + return False + +def optimize_model(model_path : Path, opt_model_path : Path): + ''' + Generate model that applies graph optimization (constant folding, etc.) + parameter model_path: path to the original onnx model + parameter opt_model_path: path to the optimized onnx model + :return: optimized onnx model + ''' + sess_option = SessionOptions() + sess_option.optimized_model_filepath = opt_model_path.as_posix() + sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC + _ = InferenceSession(model_path.as_posix(), sess_option, providers=['CPUExecutionProvider']) + +def add_infer_metadata(model): + metadata_props = {"onnx.infer": "onnxruntime.quant"} + if model.metadata_props: + for p in model.metadata_props: + metadata_props.update({p.key : p.value}) + onnx.helper.set_model_props(model, metadata_props) + +def model_has_infer_metadata(model): + if model.metadata_props: + for p in model.metadata_props: + if p.key == "onnx.infer" and p.value == "onnxruntime.quant": + return True + return False + +def load_model_with_shape_infer(model_path : Path): + inferred_model_path = generate_identified_filename(model_path, "-inferred") + onnx.shape_inference.infer_shapes_path(str(model_path), str(inferred_model_path)) + model = onnx.load(inferred_model_path.as_posix()) + inferred_model_path.unlink() + return model + + +def load_model(model_path : Path, need_optimize : bool): + with tempfile.TemporaryDirectory(prefix='ort.quant.') as quant_tmp_dir: + if need_optimize and not model_has_external_data(model_path): + opt_model_path = Path(quant_tmp_dir).joinpath("model.onnx") + optimize_model(model_path, opt_model_path) + model_path = opt_model_path + + model = load_model_with_shape_infer(model_path) + add_infer_metadata(model) + return model + +def save_and_reload_model(model): + with tempfile.TemporaryDirectory(prefix='ort.quant.') as quant_tmp_dir: + model_path = Path(quant_tmp_dir).joinpath("model.onnx") + onnx.external_data_helper.convert_model_to_external_data(model, + all_tensors_to_one_file=True) + onnx.save_model(model, model_path.as_posix()) + return load_model(model_path, False) + +def clone_model_with_shape_infer(model): + if model_has_infer_metadata(model): + cloned_model = onnx_proto.ModelProto() + cloned_model.CopyFrom(model) + else: + cloned_model = save_and_reload_model(model) + return cloned_model \ No newline at end of file diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 4d2de4de76..0439e40cf4 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -3,21 +3,14 @@ # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- -import os -import onnx -import onnx.numpy_helper -import struct import logging -import numpy as np - from pathlib import Path - from onnx import onnx_pb as onnx_proto -from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg from .quant_utils import QuantType, QuantFormat +from .quant_utils import load_model from .registry import QLinearOpsRegistry, IntegerOpsRegistry @@ -27,33 +20,6 @@ from .qdq_quantizer import QDQQuantizer from .calibrate import CalibrationDataReader, create_calibrator, CalibrationMethod -def optimize_model(model_path : Path): - ''' - Generate model that applies graph optimization (constant folding, etc.) - parameter model_path: path to the original onnx model - :return: optimized onnx model - ''' - opt_model_path = generate_identified_filename(model_path, "-opt") - sess_option = SessionOptions() - sess_option.optimized_model_filepath = opt_model_path.as_posix() - sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC - _ = InferenceSession(model_path.as_posix(), sess_option, providers=['CPUExecutionProvider']) - optimized_model = onnx.load(opt_model_path.as_posix()) - return optimized_model - - -def load_model(model_path : Path, optimize=True, handle_gemm_with_matmul=True): - - model = optimize_model(Path(model_path)) if optimize else onnx.load(Path(model_path)) - - if handle_gemm_with_matmul: - onnx_model = ONNXModel(model) - onnx_model.replace_gemm_with_matmul() - return onnx_model.model - - return model - - def check_static_quant_arguments(quant_format : QuantFormat, activation_type : QuantType, weight_type : QuantType): @@ -205,8 +171,6 @@ def quantize_static(model_input, WeightSymmetric = True/False: symmetrize calibration data for weights (default is True). EnableSubgraph = True/False : Default is False. If enabled, subgraph will be quantized. Dyanmic mode currently is supported. Will support more in future. - DisableShapeInference = True/False : in dynamic quantize mode, shape inference is not must have - and if it cause some issue, you could disable it. ForceQuantizeNoInputCheck = True/False : By default, some latent operators like maxpool, transpose, do not quantize if their input is not quantized already. Setting to True to force such operator always quantize input and so generate quantized output. Also the True behavior @@ -236,7 +200,7 @@ def quantize_static(model_input, if not op_types_to_quantize or len(op_types_to_quantize) == 0: op_types_to_quantize = list(QLinearOpsRegistry.keys()) - model = load_model(Path(model_input), optimize_model, False) + model = load_model(Path(model_input), optimize_model) calib_extra_options_keys = [ ('CalibTensorRangeSymmetric', 'symmetric'), @@ -328,8 +292,6 @@ def quantize_dynamic(model_input: Path, WeightSymmetric = True/False: symmetrize calibration data for weights (default is True). EnableSubgraph = True/False : Default is False. If enabled, subgraph will be quantized. Dyanmic mode currently is supported. Will support more in future. - DisableShapeInference = True/False : in dynamic quantize mode, shape inference is not must have - and if it cause some issue, you could disable it. ForceQuantizeNoInputCheck = True/False : By default, some latent operators like maxpool, transpose, do not quantize if their input is not quantized already. Setting to True to force such operator always quantize input and so generate quantized output. Also the True behavior diff --git a/onnxruntime/python/tools/transformers/quantize_helper.py b/onnxruntime/python/tools/transformers/quantize_helper.py index b84fda10d8..f7227a990b 100644 --- a/onnxruntime/python/tools/transformers/quantize_helper.py +++ b/onnxruntime/python/tools/transformers/quantize_helper.py @@ -58,22 +58,13 @@ class QuantizeHelper: @staticmethod def quantize_onnx_model(onnx_model_path, quantized_model_path, use_external_data_format=False): - from onnxruntime.quantization import quantize, QuantizationMode + from onnxruntime.quantization import quantize_dynamic + from pathlib import Path + Path(quantized_model_path).parent.mkdir(parents=True, exist_ok=True) logger.info(f'Size of full precision ONNX model(MB):{os.path.getsize(onnx_model_path)/(1024*1024)}') - onnx_opt_model = onnx.load_model(onnx_model_path) - quantized_onnx_model = quantize(onnx_opt_model, - quantization_mode=QuantizationMode.IntegerOps, - symmetric_weight=True, - force_fusions=True) - - if use_external_data_format: - from pathlib import Path - Path(quantized_model_path).parent.mkdir(parents=True, exist_ok=True) - onnx.external_data_helper.convert_model_to_external_data(quantized_onnx_model, - all_tensors_to_one_file=True, - location=Path(quantized_model_path).name + ".data") - onnx.save_model(quantized_onnx_model, quantized_model_path) - + quantize_dynamic(onnx_model_path, + quantized_model_path, + use_external_data_format = use_external_data_format) logger.info(f"quantized model saved to:{quantized_model_path}") #TODO: inlcude external data in total model size. logger.info(f'Size of quantized ONNX model(MB):{os.path.getsize(quantized_model_path)/(1024*1024)}')