diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index 38b64a171c..38e3c5e30e 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -42,6 +42,7 @@ class ONNXQuantizer: self.static = static # use static quantization for inputs. self.fuse_dynamic_quant = False 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'] is_weight_int8 = weight_qType == QuantType.QInt8 self.is_weight_symmetric = is_weight_int8 if 'WeightSymmetric' not in self.extra_options else self.extra_options['WeightSymmetric'] @@ -171,7 +172,7 @@ class ONNXQuantizer: def remove_fake_quantized_nodes(self): ''' - Detect and remove the quantize/dequantizelinear node pairs(fake quantized nodes in Quantization-Aware training) + Detect and remove the quantize/dequantizelinear node pairs(fake quantized nodes in Quantization-Aware training) and reconnect and update the nodes. ''' nodes_to_remove = [] @@ -294,8 +295,11 @@ class ONNXQuantizer: self.model.graph().ClearField('node') self.model.graph().node.extend(self.new_nodes) - # Remove ununsed weights from graph. - self.remove_quantized_weights() + # Remove ununsed initializers from graph, starting from the top level graph. + if self.parent is None: + _, initializers_not_found = ONNXQuantizer.CleanGraphInitializers(self.model.graph(), self.model.model) + if len(initializers_not_found) > 0: + raise RuntimeError("Invalid model with unknown initializers/tensors." + str(initializers_not_found)) self.model.model.producer_name = __producer__ self.model.model.producer_version = __version__ @@ -542,6 +546,13 @@ class ONNXQuantizer: self.quantized_value_map[input_name] = QuantizedValue(input_name, output_name, scale_name, zp_name, qType) return nodes + [qlinear_node] + def find_quantized_value(self, input_name): + if input_name in self.quantized_value_map: + return self.quantized_value_map[input_name] + if self.parent is not None: + return self.parent.find_quantized_value(input_name) + return None + def quantize_bias_static(self, bias_name, input_name, weight_name): ''' Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale @@ -699,7 +710,7 @@ class ONNXQuantizer: :param weight: TensorProto initializer :param qType: type to quantize to :param keep_float_weight: Whether to quantize the weight. In some cases, we only want to qunatize scale and zero point. - If keep_float_weight is False, quantize the weight, or don't quantize the weight. + If keep_float_weight is False, quantize the weight, or don't quantize the weight. :return: quantized weight name, zero point name, scale name ''' # Find if this input is already quantized @@ -733,7 +744,7 @@ class ONNXQuantizer: return q_weight_name, zp_name, scale_name - def quantize_weight_per_channel(self, weight_name, weight_qType, channel_axis, reduce_range=True, + def quantize_weight_per_channel(self, weight_name, weight_qType, channel_axis, reduce_range=True, keep_float_weight=False): # Find if this input is already quantized if weight_name in self.quantized_value_map: @@ -857,23 +868,74 @@ class ONNXQuantizer: return quantization_params - def remove_quantized_weights(self): - ''' Remove the weights which are already quantized from graph initializer list. - This function assumes that after quantization, all nodes that previously use a weight: - - use output from DequantizeLinear as input if they do not support quantization. - - use quantized weight if they support quantization. + + # static method + def CleanGraphInitializers(graph, model): ''' - for tensor_name, quant_value in self.quantized_value_map.items(): - if quant_value.value_type == QuantizedValueType.Initializer: - weight = self.model.get_initializer(tensor_name) + Clean unused initializers including which is caused by quantizing the model. + return cleaned graph, and list of tensor names from this graph and all its subgraphes + that can not be found in this graph and its subgraphes + ''' + requesting_tensor_names = {} + requesting_tensor_names.update({input_name: 1 for node in graph.node for input_name in node.input if input_name}) + requesting_tensor_names.update({g_out.name: 1 for g_out in graph.output if g_out.name}) - if weight is not None: - self.model.initializer().remove(weight) + new_nodes = [] + for node in graph.node: + node_2_add = node + graph_attrs = [attr for attr in node.attribute if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS] + if len(graph_attrs) > 0: + kwargs = {} + for attr in node.attribute: + kv = {} + if attr.type == onnx.AttributeProto.GRAPH: + cleaned_sub_graph, sub_requesting_tensor_names = ONNXQuantizer.CleanGraphInitializers(attr.g, model) + kv = {attr.name: cleaned_sub_graph} + requesting_tensor_names.update({gn: 1 for gn in sub_requesting_tensor_names}) + elif attr.type == onnx.AttributeProto.GRAPHS: + cleaned_graphes = [] + for subgraph in attr.graphs: + cleaned_sub_graph, sub_requesting_tensor_names = ONNXQuantizer.CleanGraphInitializers(subgraph, model) + cleaned_graphes.extend([cleaned_sub_graph]) + requesting_tensor_names.update({gn: 1 for gn in sub_requesting_tensor_names}) + kv = {attr.name: cleaned_graphes} + else: + kv = attribute_to_kwarg(attr) + kwargs.update(kv) + node_2_add = onnx.helper.make_node(node.op_type, node.input, node.output, name=node.name, **kwargs) + new_nodes.extend([node_2_add]) - # Remove from graph.input - try: - weight_input = next(val for val in self.model.graph().input if val.name == tensor_name) - self.model.graph().input.remove(weight_input) - except StopIteration: - if self.model.ir_version() < 4: - print("Warning: invalid weight name {} found in the graph (not a graph input)".format(tensor_name)) + graph.ClearField('node') + graph.node.extend(new_nodes) + + generated_names = {} + generated_names.update({output_name: 1 for node in graph.node for output_name in node.output if output_name}) + for gn in generated_names: + requesting_tensor_names.pop(gn, None) + + name_to_input = {} + for input in graph.input: + name_to_input[input.name] = input + + unused_ini_tensors = [] + for ini_tensor in graph.initializer: + if ini_tensor.name in requesting_tensor_names: + requesting_tensor_names.pop(ini_tensor.name, None) + else: + # mark it to remove, remove here directly will cause mis-behavier + unused_ini_tensors.append(ini_tensor) + + for ini_tensor in unused_ini_tensors: + graph.initializer.remove(ini_tensor) + if ini_tensor.name in name_to_input: + try: + graph.input.remove(name_to_input[ini_tensor.name]) + except StopIteration: + if model.ir_version < 4: + print("Warning: invalid weight name {} found in the graph (not a graph input)".format(ini_tensor.name)) + + for input in graph.input: + if input.name in requesting_tensor_names: + requesting_tensor_names.pop(input.name, None) + + return graph, requesting_tensor_names diff --git a/onnxruntime/python/tools/quantization/operators/direct_q8.py b/onnxruntime/python/tools/quantization/operators/direct_q8.py index 255aba738c..9ec70436c7 100644 --- a/onnxruntime/python/tools/quantization/operators/direct_q8.py +++ b/onnxruntime/python/tools/quantization/operators/direct_q8.py @@ -1,6 +1,6 @@ from .base_operator import QuantOperatorBase from .qdq_base_operator import QDQOperatorBase -from ..quant_utils import QuantizedValue +from ..quant_utils import QuantizedValue, QuantizedValueType # For operators that support 8bits operations directly, and output could # reuse input[0]'s type, zeropoint, scale; For example,Transpose, Reshape, etc. @@ -11,21 +11,46 @@ class Direct8BitOp(QuantOperatorBase): def quantize(self): node = self.node - # Quantize when input[0] is quantized already. Otherwise keep it. - if node.input[0] not in self.quantizer.quantized_value_map: + if not self.quantizer.force_quantize_no_input_check: + # Keep backward compatiblity + # Quantize when input[0] is quantized already. Otherwise keep it. + quantized_input_value = self.quantizer.find_quantized_value(node.input[0]) + if quantized_input_value is None: + self.quantizer.new_nodes += [node] + return + + quantized_output_value = QuantizedValue(node.output[0], node.output[0] + "_quantized", + quantized_input_value.scale_name, quantized_input_value.zp_name, + quantized_input_value.value_type) + self.quantizer.quantized_value_map[node.output[0]] = quantized_output_value + + node.input[0] = quantized_input_value.q_name + node.output[0] = quantized_output_value.q_name self.quantizer.new_nodes += [node] - return - # Create an entry for output quantized value - quantized_input_value = self.quantizer.quantized_value_map[node.input[0]] - quantized_output_value = QuantizedValue(node.output[0], node.output[0] + "_quantized", - quantized_input_value.scale_name, quantized_input_value.zp_name, - quantized_input_value.value_type) - self.quantizer.quantized_value_map[node.output[0]] = quantized_output_value + else: + # Force quantize those ops if possible, use black list on node if this is not you want + if (not self.quantizer.is_valid_quantize_weight(node.input[0])): + super().quantize() + return + + (quantized_input_names, zero_point_names, scale_names, nodes) = \ + self.quantizer.quantize_inputs(node, [0]) + if quantized_input_names is None: + return super().quantize() + + # Create an entry for output quantized value + quantized_output_value = QuantizedValue(node.output[0], node.output[0] + "_quantized", + scale_names[0], zero_point_names[0], + QuantizedValueType.Input) + self.quantizer.quantized_value_map[node.output[0]] = quantized_output_value + + node.input[0] = quantized_input_names[0] + node.output[0] = quantized_output_value.q_name + nodes.append(node) + + self.quantizer.new_nodes += nodes - node.input[0] = quantized_input_value.q_name - node.output[0] = quantized_output_value.q_name - self.quantizer.new_nodes += [node] class QDQDirect8BitOp(QDQOperatorBase): diff --git a/onnxruntime/python/tools/quantization/qdq_quantizer.py b/onnxruntime/python/tools/quantization/qdq_quantizer.py index 839ee60b09..a14e220d3c 100644 --- a/onnxruntime/python/tools/quantization/qdq_quantizer.py +++ b/onnxruntime/python/tools/quantization/qdq_quantizer.py @@ -101,7 +101,7 @@ class QDQQuantizer(ONNXQuantizer): self.quantize_bias_tensors() self.remove_nodes() if not self.add_qdq_pair_to_weight: - self.remove_quantized_weights() + ONNXQuantizer.CleanGraphInitializers(self.model.graph(), self.model.model) self.model.model.producer_name = __producer__ self.model.model.producer_version = __version__ diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index e70a84c23b..d72e372bbf 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -188,6 +188,10 @@ def quantize_static(model_input, 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 + could be disabled per node using the nodes_to_exclude. MatMulConstBOnly = True/False: Default is False. If enabled, only MatMul with const B will be quantized. AddQDQPairToWeight = True/False : Default is False which quantizes floating-point weight and feeds it to soley inserted DeQuantizeLinear node. If True, it remains floating-point weight and @@ -283,6 +287,10 @@ def quantize_dynamic(model_input: Path, 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 + could be disabled per node using the nodes_to_exclude. MatMulConstBOnly = True/False: Default is False. If enabled, only MatMul with const B will be quantized. ''' diff --git a/onnxruntime/python/tools/quantization/registry.py b/onnxruntime/python/tools/quantization/registry.py index c51cd65151..3628bd2ec9 100644 --- a/onnxruntime/python/tools/quantization/registry.py +++ b/onnxruntime/python/tools/quantization/registry.py @@ -20,6 +20,7 @@ from .operators.concat import QLinearConcat, QDQConcat CommonOpsRegistry = { "Gather": GatherQuant, + "Transpose" : Direct8BitOp, "EmbedLayerNormalization": EmbedLayerNormalizationQuant, } @@ -45,7 +46,6 @@ QLinearOpsRegistry = { "Split": QSplit, "Pad": QPad, "Reshape": Direct8BitOp, - "Transpose" : Direct8BitOp, "Squeeze" : Direct8BitOp, "Unsqueeze" : Direct8BitOp, "Resize": QResize,