Remove shape inference and fix save large model(>2g) issue (#5210)

* remove shape inference and fix save large model problem

* remove unnecessary import

* refine code and add external format for quantize_qat

* remove initializers in tensors_to_calibrate

* small refine

Co-authored-by: t-yguo <t-yguo@microsoft.com>

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

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)

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)

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)