From c6ef6b5bc8eeb49a6b8f6fadfd212c9e07ddde31 Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Wed, 11 Aug 2021 16:35:52 -0700 Subject: [PATCH] Subgraph support for quantization tools (#8012) By default, not do enable subgraph quantization to make it consistent with existing behavior. It should be OK to enable it at quantize_dynamic mode with extra_options. --- .../python/tools/quantization/onnx_model.py | 64 ++++++++-- .../tools/quantization/onnx_quantizer.py | 114 ++++++++++++++++-- .../tools/quantization/operators/pad.py | 1 + .../python/tools/quantization/quantize.py | 5 + 4 files changed, 164 insertions(+), 20 deletions(-) diff --git a/onnxruntime/python/tools/quantization/onnx_model.py b/onnxruntime/python/tools/quantization/onnx_model.py index 0c888abb76..dfa2e274d3 100644 --- a/onnxruntime/python/tools/quantization/onnx_model.py +++ b/onnxruntime/python/tools/quantization/onnx_model.py @@ -1,9 +1,8 @@ import onnx import itertools -from .quant_utils import find_by_name +from .quant_utils import find_by_name, attribute_to_kwarg from pathlib import Path - class ONNXModel: def __init__(self, model): self.model = model @@ -142,10 +141,39 @@ class ONNXModel: nodes.append(node) return nodes - def replace_gemm_with_matmul(self): - new_nodes = [] + @staticmethod + def __get_initializer(name, graph_path): + for gid in range(len(graph_path) - 1, -1, -1): + graph = graph_path[gid] + for tensor in graph.initializer: + if tensor.name == name: + return tensor, graph + return None, None + + @staticmethod + def __replace_gemm_with_matmul(graph_path): + new_nodes = [] + graph = graph_path[-1] + for node in graph.node: + graph_attrs = [attr for attr in node.attribute if attr.type == 5 or attr.type == 10] + if len(graph_attrs): + node_name = node.name + kwargs = {} + for attr in node.attribute: + if attr.type == 5: + graph_path.append(attr.g) + kv = {attr.name: ONNXModel.__replace_gemm_with_matmul(graph_path)} + elif attr.type == 10: + value = [] + for subgraph in attr.graphs: + graph_path.append(subgraph) + value.extend([ONNXModel.__replace_gemm_with_matmul(graph_path)]) + kv = {attr.name: value} + else: + kv = attribute_to_kwarg(attr) + kwargs.update(kv) + node = onnx.helper.make_node(node.op_type, node.input, node.output, name=node.name, **kwargs) - for node in self.nodes(): if node.op_type == 'Gemm': alpha = 1.0 beta = 1.0 @@ -163,34 +191,38 @@ class ONNXModel: if alpha == 1.0 and beta == 1.0 and transA == 0: inputB = node.input[1] if transB == 1: - B = self.get_initializer(node.input[1]) + B, Bs_graph = ONNXModel.__get_initializer(node.input[1], graph_path) if B: # assume B is not used by any other node B_array = onnx.numpy_helper.to_array(B) B_trans = onnx.numpy_helper.from_array(B_array.T) B_trans.name = B.name - self.remove_initializer(B) - self.add_initializer(B_trans) + Bs_graph.initializer.remove(B) + for input in Bs_graph.input: + if input.name == inputB: + Bs_graph.input.remove(input) + break + Bs_graph.initializer.extend([B_trans]) else: inputB += '_Transposed' transpose_node = onnx.helper.make_node('Transpose', inputs=[node.input[1]], outputs=[inputB], - name=node.name + '_Transpose') + name=node.name + '_Transpose' if node.name != "" else "") new_nodes.append(transpose_node) matmul_node = onnx.helper.make_node( 'MatMul', inputs=[node.input[0], inputB], outputs=[node.output[0] + ('_MatMul' if len(node.input) > 2 else '')], - name=node.name + '_MatMul' if node.name else "") + name=node.name + '_MatMul' if node.name != "" else "") new_nodes.append(matmul_node) if len(node.input) > 2: add_node = onnx.helper.make_node('Add', inputs=[node.output[0] + '_MatMul', node.input[2]], outputs=node.output, - name=node.name + '_Add' if node.name else "") + name=node.name + '_Add' if node.name != "" else "") new_nodes.append(add_node) # unsupported @@ -201,8 +233,14 @@ class ONNXModel: else: new_nodes.append(node) - self.graph().ClearField('node') - self.graph().node.extend(new_nodes) + graph.ClearField('node') + graph.node.extend(new_nodes) + graph_path.pop() + return graph + + def replace_gemm_with_matmul(self): + graph_path = [self.graph()] + ONNXModel.__replace_gemm_with_matmul(graph_path) def save_model_to_file(self, output_path, use_external_data_format=False): ''' diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index 98116c748d..0692e26a7c 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -27,8 +27,10 @@ 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 - model = onnx.shape_inference.infer_shapes(model) + # 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) 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}) @@ -39,7 +41,7 @@ class ONNXQuantizer: 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.q_matmul_const_b_only = 'MatMulConstBOnly' in self.extra_options and self.extra_options['MatMulConstBOnly'] self.is_weight_symmetric = True if 'WeightSymmetric' not in self.extra_options else self.extra_options['WeightSymmetric'] self.is_activation_symmetric = False if 'ActivationSymmetric' not in self.extra_options else self.extra_options['ActivationSymmetric'] @@ -62,6 +64,13 @@ class ONNXQuantizer: self.nodes_to_exclude = nodes_to_exclude # specific nodes to exclude self.op_types_to_quantize = op_types_to_quantize self.new_nodes = [] + self.parent = None + self.graph_scope = "/" # for human readable debug information + self.tensor_names = { } # in case the shape inference not totally working + self.tensor_names.update({ot.name: 1 for ot in model.graph.output}) + self.tensor_names.update({it.name: 1 for it in model.graph.input}) + for node in self.model.model.graph.node: + self.tensor_names.update({output_name: 1 for output_name in node.output}) self.opset_version = self.check_opset_version() @@ -85,6 +94,55 @@ class ONNXQuantizer: # no dequantized will be applied when needed later self.generated_value_names = self.model.get_non_initializer_inputs() + # routines for subgraph support + def quantize_subgraph(self, subgraph, graph_key): + ''' + generate submodel for the subgraph, so that we re-utilize current quantization implementation. + quantize the submodel + update subgraph and set it back to node + ''' + warped_model = onnx.helper.make_model(subgraph, producer_name='onnx-quantizer', + opset_imports=self.model.model.opset_import) + sub_quanitzer = ONNXQuantizer(warped_model, + self.per_channel, + self.reduce_range, + self.mode, + self.static, + self.weight_qType, + self.input_qType, + self.tensors_range, + self.nodes_to_quantize, + self.nodes_to_exclude, + self.op_types_to_quantize, + self.extra_options) + sub_quanitzer.parent = self + sub_quanitzer.graph_scope = "{}{}/".format(self.graph_scope, graph_key) + sub_quanitzer.quantize_model() + return sub_quanitzer.model.model.graph + + def quantize_node_with_sub_graph(self, node): + ''' + Check subgraph, if any, quantize it and replace it. + return new_nodes added for quantizing subgraph + ''' + 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: + return node + node_name = node.name if node.name != "" else "{}_node_count_{}".format(node.op_type, len(self.new_nodes)) + kwargs = {} + for attr in node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + kv = {attr.name: self.quantize_subgraph(attr.g, "{}:{}".format(node_name, attr.name))} + elif attr.type == onnx.AttributeProto.GRAPHS: + value = [] + for subgraph in attr.graphs: + value.extend([self.quantize_subgraph(subgraph, "{}:{}:{}".format(node_name, attr.name, len(value)))]) + kv = {attr.name: value} + else: + kv = attribute_to_kwarg(attr) + kwargs.update(kv) + return onnx.helper.make_node(node.op_type, node.input, node.output, name=node.name, **kwargs) + def check_opset_version(self): ai_onnx_domain = [ opset for opset in self.model.model.opset_import if not opset.domain or opset.domain == "ai.onnx" @@ -177,6 +235,13 @@ class ONNXQuantizer: return self.model.model + def find_initializer_in_path(self, initializer_name): + if find_by_name(initializer_name, self.model.initializer()) is not None: + return True + if self.parent is not None: + return self.parent.find_initializer_in_path(initializer_name) + return False + def should_quantize(self, node): if self.nodes_to_quantize is not None and len( self.nodes_to_quantize) != 0 and node.name not in self.nodes_to_quantize: @@ -190,15 +255,26 @@ class ONNXQuantizer: # do not quantize non-constant B matrices for matmul if self.q_matmul_const_b_only: - if node.op_type == "MatMul" and find_by_name(node.input[1], self.model.initializer()) is None: + if node.op_type == "MatMul" and (not self.find_initializer_in_path(node.input[1])): + print("Ignore MatMul due to non constant B: {}[{}]".format(self.graph_scope, node.name)) return False return True + def add_new_nodes(self, nodes): + self.new_nodes.extend(nodes) + for node in nodes: + for output_name in node.output: + self.generated_value_names.add(output_name) + def quantize_model(self): self.remove_fake_quantized_nodes() for node in self.model.nodes(): + # quantize subgraphes if have + if self.enable_subgraph_quantization: + node = self.quantize_node_with_sub_graph(node) + number_of_existing_new_nodes = len(self.new_nodes) if self.should_quantize(node): op_quantizer = CreateOpQuantizer(self, node) @@ -522,7 +598,13 @@ class ONNXQuantizer: return quantized_bias_name - def quantize_inputs(self, node, indices, initializer_use_weight_qType=True, reduce_range=False, op_level_per_channel=False, axis=-1): + def contains_tensor(self, tensor_name): + ''' + only check for value info and newly generated tensor names, initializers are checked seperately + ''' + return (tensor_name in self.value_infos) or (tensor_name in self.tensor_names) or (tensor_name in self.generated_value_names) + + def quantize_inputs(self, node, indices, initializer_use_weight_qType=True, reduce_range=False, op_level_per_channel=False, axis=-1, from_subgraph=False): ''' Given a node, this function quantizes the inputs as follows: - If input is an initializer, quantize the initializer data, replace old initializer @@ -567,13 +649,16 @@ class ONNXQuantizer: quantized_input_names.append(q_weight_name) zero_point_names.append(zp_name) scale_names.append(scale_name) - else: + elif self.contains_tensor(node_input): # Add QuantizeLinear node. qlinear_node = self.model.find_node_by_name(node_input + "_QuantizeLinear", self.new_nodes, self.model.graph()) if qlinear_node is None: quantize_input_nodes = self._get_quantize_input_nodes(node, input_index, self.input_qType) - nodes.extend(quantize_input_nodes) + if from_subgraph: + self.add_new_nodes(quantize_input_nodes) + else: + nodes.extend(quantize_input_nodes) qlinear_node = quantize_input_nodes[-1] if qlinear_node.op_type == "QuantizeLinear": @@ -584,6 +669,21 @@ class ONNXQuantizer: quantized_input_names.append(qlinear_node.output[0]) scale_names.append(qlinear_node.output[1]) zero_point_names.append(qlinear_node.output[2]) + elif self.parent is not None: + (parent_quantized_input_names, parent_zero_point_names, parent_scale_names, _) = self.parent.quantize_inputs( + node, + [input_index], + initializer_use_weight_qType=initializer_use_weight_qType, + reduce_range=reduce_range, + op_level_per_channel=op_level_per_channel, + axis=axis, + from_subgraph=True) + quantized_input_names.append(parent_quantized_input_names[0]) + scale_names.append(parent_scale_names[0]) + zero_point_names.append(parent_zero_point_names[0]) + # node should not be add this child level here + else: + raise ValueError('Invalid tensor name to quantize: {} @graph scope{}'.format(node_input, self.graph_scope)) return (quantized_input_names, zero_point_names, scale_names, nodes) diff --git a/onnxruntime/python/tools/quantization/operators/pad.py b/onnxruntime/python/tools/quantization/operators/pad.py index ed3d001a35..75a0cad10f 100644 --- a/onnxruntime/python/tools/quantization/operators/pad.py +++ b/onnxruntime/python/tools/quantization/operators/pad.py @@ -49,6 +49,7 @@ class QPad(QuantOperatorBase): self.quantizer.model.add_initializer(quantized_padding_constant_initializer) node.input[2] = quantized_padding_constant_name else: + # TODO: check quantize_inputs after sub graph is supported pad_value_qnodes = self.quantizer._get_quantize_input_nodes(node, 2, self.quantizer.input_qType, quantized_input_value.scale_name, quantized_input_value.zp_name) diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 5b88be1317..c14c46104a 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -184,6 +184,11 @@ def quantize_static(model_input, extra.Sigmoid.nnapi = True/False (Default is False) ActivationSymmetric = True/False: symmetrize calibration data for activations (default is False). 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. + ''' mode = QuantizationMode.QLinearOps