saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-20 19:12:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve onnx shape inference in quant tool (#11106)

Browse source 
onnx.shape_inference.infer_shapes only works for model size < 2GB, while onnx.shape_inference.infer_shapes_path works for all models. This PR replaces infer_shapes with infer_shapes_path.
This commit is contained in:
Yufeng Li 2022-04-18 08:07:31 -07:00 committed by GitHub
parent 9765ef8b4e
commit dec99657a1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 99 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
onnxruntime/python/tools/quantization/calibrate.py
View file

@ -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 = []

12
onnxruntime/python/tools/quantization/onnx_quantizer.py
View file

@ -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,

72
onnxruntime/python/tools/quantization/quant_utils.py
View file

@ -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

42
onnxruntime/python/tools/quantization/quantize.py
View file

@ -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

21
onnxruntime/python/tools/transformers/quantize_helper.py
View file

@ -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)}')