saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-05 04:17:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Subgraph support for quantization tools (#8012)

Browse source 
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.
This commit is contained in:
Zhang Lei 2021-08-11 16:35:52 -07:00 committed by GitHub
parent c5c5d3499b
commit c6ef6b5bc8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 164 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
onnxruntime/python/tools/quantization/onnx_model.py
View file

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

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

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

1
onnxruntime/python/tools/quantization/operators/pad.py
View file

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

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

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