diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 91cce64c27..d355fba2f3 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -54,9 +54,6 @@ class ONNXCalibrater: ''' model = onnx.load(self.model_path) - model = onnx.shape_inference.infer_shapes(model) - value_infos = {vi.name: vi for vi in model.graph.value_info} - value_infos.update({ot.name: ot for ot in model.graph.output}) added_nodes = [] added_outputs = [] @@ -66,18 +63,12 @@ class ONNXCalibrater: should_be_calibrate = ((node.op_type in self.calibrate_op_types) and (node.name not in self.black_nodes)) or (node.name in self.white_nodes) if should_be_calibrate: - for input_tensor_name in node.input: - if input_tensor_name in value_infos.keys(): - vi = value_infos[input_tensor_name] - if vi.type.HasField('tensor_type') and vi.type.tensor_type.elem_type == TensorProto.FLOAT and ( - input_tensor_name not in model.graph.initializer): - tensors_to_calibrate.add(input_tensor_name) + tensors_to_calibrate.update(node.input) + tensors_to_calibrate.update(node.output) - for output_tensor_name in node.output: - if output_tensor_name in value_infos.keys(): - vi = value_infos[output_tensor_name] - if vi.type.HasField('tensor_type') and vi.type.tensor_type.elem_type == TensorProto.FLOAT: - tensors_to_calibrate.add(output_tensor_name) + for tensor in tensors_to_calibrate: + if tensor in model.graph.initializer: + tensors_to_calibrate.remove(tensor) for tensor in tensors_to_calibrate: # Adding ReduceMin nodes @@ -269,4 +260,4 @@ def calibrate(model_path, quantization_params_dict = calibrater.calculate_quantization_params(dict_for_quantization) print("Calibrated,quantized parameters calculated and returned.") - return quantization_params_dict + return quantization_params_dict \ No newline at end of file diff --git a/onnxruntime/python/tools/quantization/onnx_model.py b/onnxruntime/python/tools/quantization/onnx_model.py index ae70492e4f..2d6b209326 100644 --- a/onnxruntime/python/tools/quantization/onnx_model.py +++ b/onnxruntime/python/tools/quantization/onnx_model.py @@ -1,5 +1,6 @@ import onnx from .quant_utils import find_by_name +from pathlib import Path class ONNXModel: @@ -125,3 +126,13 @@ class ONNXModel: if node_input == initializer.name: nodes.append(node) return nodes + + def save_model_to_file(self, output_path, use_external_data_format=False): + ''' + Save model to external data, which is needed for model size > 2GB + ''' + if use_external_data_format: + onnx.external_data_helper.convert_model_to_external_data(self.model, + all_tensors_to_one_file=True, + location=Path(output_path).name + ".data") + onnx.save_model(self.model, output_path) \ No newline at end of file diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index 8b1d012d86..06dbce2b29 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -12,7 +12,6 @@ from pathlib import Path import numpy as np from onnx import onnx_pb as onnx_proto -from onnx import shape_inference from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes @@ -77,9 +76,7 @@ def _get_qrange_for_qType(qType, reduce_range=False): class ONNXQuantizer: def __init__(self, model, per_channel, reduce_range, mode, static, weight_qType, input_qType, quantization_params, nodes_to_quantize, nodes_to_exclude, op_types_to_quantize): - onnx_model = shape_inference.infer_shapes(model) - self.model = ONNXModel(onnx_model) - self.value_infos = {vi.name: vi for vi in onnx_model.graph.value_info} + self.model = ONNXModel(model) self.per_channel = per_channel # weight-pack per channel self.reduce_range = reduce_range self.mode = mode # QuantizationMode.Value @@ -294,17 +291,6 @@ class ONNXQuantizer: initializer = find_by_name(input_name, self.model.initializer()) return initializer is not None - def _is_valid_quantize_value(self, value_name): - if value_name in self.value_infos: - value_info = self.value_infos[value_name] - return value_info.type.HasField( - 'tensor_type') and value_info.type.tensor_type.elem_type == onnx_proto.TensorProto.FLOAT - return self._is_valid_initializer_value(value_name) - - def _is_valid_initializer_value(self, value_name): - weight = find_by_name(value_name, self.model.initializer()) - return weight is not None and weight.data_type == onnx_proto.TensorProto.FLOAT - def _is_valid_quantize_weight(self, weight_name): weight = find_by_name(weight_name, self.model.initializer()) return weight is not None and weight.data_type == onnx_proto.TensorProto.FLOAT @@ -912,4 +898,4 @@ class ONNXQuantizer: for output in self.model.graph().output: dequantize_node = self._dequantize_value(output.name) if dequantize_node is not None: - self.new_nodes.append(dequantize_node) + self.new_nodes.append(dequantize_node) \ No newline at end of file diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index a5533b4fba..3f108051b6 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -12,7 +12,6 @@ from pathlib import Path import numpy as np from onnx import onnx_pb as onnx_proto -from onnx import shape_inference from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes @@ -131,7 +130,8 @@ def quantize_static(model_input, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], - nodes_to_exclude=[]): + nodes_to_exclude=[], + use_external_data_format=False): ''' Given an onnx model and calibration data reader, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize @@ -154,6 +154,7 @@ def quantize_static(model_input, :param nodes_to_exclude: List of nodes names to exclude. The nodes in this list will be excluded from quantization when it is not None. + :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default. ''' if activation_type != QuantType.QUInt8: @@ -183,7 +184,7 @@ def quantize_static(model_input, op_types_to_quantize) quantizer.quantize_model() - onnx.save_model(quantizer.model.model, model_output) + quantizer.model.save_model_to_file(model_output, use_external_data_format) def quantize_dynamic(model_input: Path, @@ -194,7 +195,8 @@ def quantize_dynamic(model_input: Path, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], - nodes_to_exclude=[]): + nodes_to_exclude=[], + use_external_data_format=False): ''' Given an onnx model, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize @@ -216,20 +218,18 @@ def quantize_dynamic(model_input: Path, :param nodes_to_exclude: List of nodes names to exclude. The nodes in this list will be excluded from quantization when it is not None. + :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default. ''' input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 weight_qType = onnx_proto.TensorProto.INT8 if weight_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 mode = QuantizationMode.IntegerOps - #optimize the original model - optimized_model = optimize_model(Path(model_input)) - if not op_types_to_quantize or len(op_types_to_quantize) == 0: op_types_to_quantize = list(IntegerOpsRegistry.keys()) quantizer = ONNXQuantizer( - optimized_model, + onnx.load(model_input), per_channel, reduce_range, mode, @@ -242,7 +242,7 @@ def quantize_dynamic(model_input: Path, op_types_to_quantize) quantizer.quantize_model() - onnx.save_model(quantizer.model.model, model_output) + quantizer.model.save_model_to_file(model_output, use_external_data_format) def quantize_qat(model_input: Path, @@ -253,7 +253,8 @@ def quantize_qat(model_input: Path, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], - nodes_to_exclude=[]): + nodes_to_exclude=[], + use_external_data_format=False): ''' Given a quantize-aware traning onnx model, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize @@ -273,6 +274,7 @@ def quantize_qat(model_input: Path, :param nodes_to_exclude: List of nodes names to exclude. The nodes in this list will be excluded from quantization when it is not None. + :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default. ''' input_qType = onnx_proto.TensorProto.INT8 if activation_type == QuantType.QInt8 else onnx_proto.TensorProto.UINT8 @@ -299,4 +301,4 @@ def quantize_qat(model_input: Path, op_types_to_quantize) quantizer.quantize_model() - onnx.save_model(quantizer.model.model, model_output) + quantizer.model.save_model_to_file(model_output, use_external_data_format) \ No newline at end of file