diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 5b81de6b33..c02234ee7d 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -56,6 +56,7 @@ 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 = [] @@ -167,7 +168,7 @@ class ONNXCalibrater: else: self.input_name_to_nodes[input_name].append(node) - def calculate_scale_zeropoint(self, node, next_node, rmin, rmax): + def calculate_scale_zeropoint(self, next_node, rmin, rmax): zp_and_scale = [] # adjust rmin and rmax such that 0 is included in the range. This is required @@ -178,16 +179,17 @@ class ONNXCalibrater: # We update the output range min and max when next node is clip or relu # With this technique we can remove these 2 ops and # reduce the output range which in turn helps to improve accuracy - if next_node.op_type == 'Clip': - clip_min = next_node.attribute[0].f - clip_max = next_node.attribute[1].f - if rmin < clip_min: - rmin = clip_min - if rmax > clip_max: - rmax = clip_max - if next_node.op_type == 'Relu': - if rmin < 0: - rmin = 0 + if next_node: + if next_node.op_type == 'Clip': + clip_min = next_node.attribute[0].f + clip_max = next_node.attribute[1].f + if rmin < clip_min: + rmin = clip_min + if rmax > clip_max: + rmax = clip_max + elif next_node.op_type == 'Relu': + if rmin < 0: + rmin = 0 scale = np.float32((rmax - rmin) / 255 if rmin != rmax else 1) initial_zero_point = (0 - rmin) / scale @@ -227,16 +229,15 @@ class ONNXCalibrater: self._get_input_name_to_nodes(model) - for node in model.graph.node: - for node_output_name in node.output: - if node_output_name in self.input_name_to_nodes: - children = self.input_name_to_nodes[node_output_name] - for child in children: - if node_output_name in quantization_thresholds: - node_thresholds = quantization_thresholds[node_output_name] - node_params = self.calculate_scale_zeropoint(node, child, node_thresholds[0], - node_thresholds[1]) - quantization_params[node_output_name] = node_params + for tensor_name in quantization_thresholds.keys(): + child = None + if tensor_name in self.input_name_to_nodes: + children = self.input_name_to_nodes[tensor_name] + if (len(children) == 1): + child = children[0] + node_thresholds = quantization_thresholds[tensor_name] + node_params = self.calculate_scale_zeropoint(child, node_thresholds[0], node_thresholds[1]) + quantization_params[tensor_name] = node_params return quantization_params diff --git a/onnxruntime/python/tools/quantization/onnx_model.py b/onnxruntime/python/tools/quantization/onnx_model.py index 7e73196da6..ae70492e4f 100644 --- a/onnxruntime/python/tools/quantization/onnx_model.py +++ b/onnxruntime/python/tools/quantization/onnx_model.py @@ -1,5 +1,5 @@ import onnx -from .quant_utils import _find_by_name +from .quant_utils import find_by_name class ONNXModel: @@ -37,7 +37,7 @@ class ONNXModel: self.model.graph.node.extend(nodes_to_add) def add_initializer(self, tensor): - if _find_by_name(tensor.name, self.model.graph.initializer) is None: + if find_by_name(tensor.name, self.model.graph.initializer) is None: self.model.graph.initializer.extend([tensor]) def get_initializer(self, name): @@ -112,7 +112,7 @@ class ONNXModel: ''' graph_nodes_list = list(graph.node) #deep copy graph_nodes_list.extend(new_nodes_list) - node = _find_by_name(node_name, graph_nodes_list) + node = find_by_name(node_name, graph_nodes_list) return node def find_nodes_by_initializer(self, graph, initializer): diff --git a/onnxruntime/python/tools/quantization/onnx_quantizer.py b/onnxruntime/python/tools/quantization/onnx_quantizer.py index 88ce7a43a6..27ab05b3f4 100644 --- a/onnxruntime/python/tools/quantization/onnx_quantizer.py +++ b/onnxruntime/python/tools/quantization/onnx_quantizer.py @@ -16,7 +16,7 @@ from onnx import shape_inference from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes -from .quant_utils import _find_by_name, _get_elem_index, _get_mul_node, _generate_identified_filename, _attribute_to_kwarg +from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg from .quant_utils import QuantType, onnx_domain, __producer__, __version__ from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer @@ -46,7 +46,7 @@ def quantize_data(data, quantize_range, qType): if qType == onnx_proto.TensorProto.INT8: max_range = max(abs(rmin), abs(rmax)) - scale = (float(max_range) * 2) / quantize_range + scale = (float(max_range) * 2) / quantize_range if max_range > 0 else 1 zero_point = 0 # signed byte type quantized_data = (np.asarray(data) / scale).round().astype('b') @@ -60,27 +60,28 @@ def quantize_data(data, quantize_range, qType): return rmin, rmax, zero_point, scale, quantized_data -def _get_qrange_for_qType(qType): +def _get_qrange_for_qType(qType, reduce_range=False): ''' Helper function to get the quantization range for a type. parameter qType: quantization type. return: quantization range. ''' if qType == onnx_proto.TensorProto.UINT8: - return 255 # 2^b - 1 + return 127 if reduce_range else 255 elif qType == onnx_proto.TensorProto.INT8: - return 254 # [-(2^{b-1}-1), 2^{b-1}-1]: [-127, 127] for 8 bits. + return 128 if reduce_range else 254 # [-64, 64] for reduce_range, and [-127, 127] full_range. else: raise ValueError('unsupported quantization data type') class ONNXQuantizer: - def __init__(self, model, per_channel, mode, static, weight_qType, input_qType, quantization_params, + 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.per_channel = per_channel # weight-pack per channel + self.reduce_range = reduce_range self.mode = mode # QuantizationMode.Value self.static = static # use static quantization for inputs. self.fuse_dynamic_quant = False @@ -119,15 +120,18 @@ class ONNXQuantizer: raise ValueError('Failed to find proper ai.onnx domain') opset_version = ai_onnx_domain[0].version - if opset_version < 10: - raise ValueError("The original model opset version is {}, which does not support quantized operators.\n\ - The opset version of quantized model will be set to 10. Use onnx model checker to verify model after quantization." - .format(opset_version)) - if opset_version == 10: - self.fuse_dynamic_quant = False - else: - self.fuse_dynamic_quant = True + print( + "Warning: The original model opset version is {}, which does not support node fusions. Please update the model to opset >= 11 for better performance." + .format(opset_version)) + return + + if opset_version < 10: + print( + "Warning: The original model opset version is {}, which does not support quantization. Please update the model to opset >= 11. Updating the model automatically to opset 11. Please verify the quantized model." + .format(opset_version)) + self.model.model.opset_import.remove(ai_onnx_domain[0]) + self.model.model.opset_import.extend([onnx.helper.make_opsetid("", 11)]) def replace_gemm_with_matmul(self): nodes_to_remove = [] @@ -196,8 +200,8 @@ class ONNXQuantizer: # TODO: convert it to the specified input_type scale_tensor_name = curr_node.input[1] zp_tensor_name = curr_node.input[2] - initializer_scale = _find_by_name(scale_tensor_name, self.model.initializer()) - initializer_zp = _find_by_name(zp_tensor_name, self.model.initializer()) + initializer_scale = find_by_name(scale_tensor_name, self.model.initializer()) + initializer_zp = find_by_name(zp_tensor_name, self.model.initializer()) zp_and_scale = [ onnx.numpy_helper.to_array(initializer_zp), onnx.numpy_helper.to_array(initializer_scale) @@ -205,7 +209,7 @@ class ONNXQuantizer: #connect the previous and successive node input and output for succ_node in succ_nodes: - succ_idx = _get_elem_index(next_node.output[0], succ_node.input) + succ_idx = get_elem_index(next_node.output[0], succ_node.input) if succ_idx != -1: succ_node.input[succ_idx] = curr_node.input[0] else: @@ -273,11 +277,8 @@ class ONNXQuantizer: return self.model.model - def find_weight_data(self, initializer): - ''' - :param initializer: TensorProto initializer object from a graph - :return: a list of initialized data in a given initializer object - ''' + @staticmethod + def tensor_proto_to_array(initializer): if initializer.data_type == onnx_proto.TensorProto.FLOAT: weights = onnx.numpy_helper.to_array(initializer) else: @@ -285,6 +286,10 @@ class ONNXQuantizer: initializer.name, type_to_name[initializer.data_type])) return weights + def is_input_a_weight(self, input_name): + 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] @@ -293,11 +298,11 @@ class ONNXQuantizer: 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()) + 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()) + weight = find_by_name(weight_name, self.model.initializer()) return weight is not None and weight.data_type == onnx_proto.TensorProto.FLOAT def _remove_quantized_weights(self): @@ -318,7 +323,7 @@ class ONNXQuantizer: if self.model.ir_version() < 4: print("Warning: invalid weight name {} found in the graph (not a graph input)".format(weight.name)) - def _update_graph(self, weight): + def _update_weight(self, weight): ''' Given a weight object, update the graph by doing the following: - remove old initializer, update new initializers for quantized weight, zero point, and scale @@ -357,9 +362,9 @@ class ONNXQuantizer: :param qType: type to quantize to :return: Weight class with quantization information ''' - weights_data = self.find_weight_data(initializer) - rmin, rmax, zero_point, scale, quantized_weights_data = quantize_data(weights_data.flatten().tolist(), - _get_qrange_for_qType(qType), qType) + weights_data = self.tensor_proto_to_array(initializer) + rmin, rmax, zero_point, scale, quantized_weights_data = quantize_data( + weights_data.flatten().tolist(), _get_qrange_for_qType(qType, self.reduce_range), qType) weight = QuantizedInitializer(initializer.name, initializer, [rmin], [rmax], [zero_point], [scale], weights_data, @@ -375,7 +380,7 @@ class ONNXQuantizer: return weight - def _get_quantized_weight_convolution(self, initializer, qType): + def _get_quantized_weight_per_channel(self, initializer, qType, channel_axis): ''' :param initializer: initializer TypeProto to quantize :param qType: type to quantize to @@ -384,38 +389,34 @@ class ONNXQuantizer: if not self.per_channel: return self._get_quantized_weight(initializer, qType) - weights = self.find_weight_data(initializer) - # Quantize per output channel - # Assuming (M x C/group x kH x kW) format where M is number of output channels. - channel_count = initializer.dims[0] - np_data = np.reshape(weights, initializer.dims) + weights = self.tensor_proto_to_array(initializer) + channel_count = weights.shape[channel_axis] rmin_list = [] rmax_list = [] zero_point_list = [] scale_list = [] quantized_per_channel_data_list = [] for i in range(channel_count): - # for each channel, compute quantization data. Assuming (M x C/group x kH x kW) - per_channel_data = np_data[i, :, :, :].flatten() + per_channel_data = weights.take(i, channel_axis) rmin, rmax, zero_point, scale, quantized_per_channel_data = quantize_data( - per_channel_data.flatten().tolist(), _get_qrange_for_qType(qType), qType) + per_channel_data.flatten().tolist(), _get_qrange_for_qType(qType, self.reduce_range), qType) rmin_list.append(rmin) rmax_list.append(rmax) zero_point_list.append(zero_point) scale_list.append(scale) quantized_per_channel_data_list.append(quantized_per_channel_data) - channel_index = 0 # (M x C/group x kH x kW) + # combine per_channel_data into one - reshape_dims = list(initializer.dims) # deep copy - reshape_dims[channel_index] = 1 # only one per channel for reshape + reshape_dims = list(weights.shape) # deep copy + reshape_dims[channel_axis] = 1 # only one per channel for reshape quantized_weights = np.asarray(quantized_per_channel_data_list[0]).reshape(reshape_dims) for i in range(1, len(quantized_per_channel_data_list)): channel_weights = np.asarray(quantized_per_channel_data_list[i]).reshape(reshape_dims) - quantized_weights = np.concatenate((quantized_weights, channel_weights), axis=0) + quantized_weights = np.concatenate((quantized_weights, channel_weights), channel_axis) weight = QuantizedInitializer(initializer.name, initializer, rmin_list, rmax_list, zero_point_list, scale_list, weights, - quantized_weights.flatten().tolist(), channel_index, qType) + quantized_weights.flatten().tolist(), channel_axis, qType) # Make entry for this quantized weight assert (weight.name not in self.quantized_value_map) @@ -642,7 +643,7 @@ class ONNXQuantizer: return nodes + [qlinear_node] - def _get_bias_add_nodes(self, nodes, node, last_output, quantized_bias_name): + def get_bias_add_nodes(self, nodes, node, last_output, quantized_bias_name): ''' Given a node, this function handles bias add by adding a "reshape" node on bias and an "add" node parameter nodes: new nodes would be appended into nodes @@ -738,13 +739,13 @@ class ONNXQuantizer: # get scale for weight weight_scale_name = self.quantized_value_map[node.input[1]].scale_name - weight_initializer = _find_by_name(weight_scale_name, self.model.initializer()) - weight_scale = self.find_weight_data(weight_initializer) + weight_initializer = find_by_name(weight_scale_name, self.model.initializer()) + weight_scale = self.tensor_proto_to_array(weight_initializer) # get bias bias_name = node.input[2] - bias_initializer = _find_by_name(bias_name, self.model.initializer()) - bias_data = self.find_weight_data(bias_initializer) + bias_initializer = find_by_name(bias_name, self.model.initializer()) + bias_data = self.tensor_proto_to_array(bias_initializer) quantized_bias_name = bias_name + "_quantized" # input scale is not provided and this input is dynamically quantized so it is not pre-computed at this point @@ -762,8 +763,8 @@ class ONNXQuantizer: raise ValueError("Expected {} to be in quantized value map for static quantization".format( node.input[0])) - inputscale_initializer = _find_by_name(input_scale_name, self.model.initializer()) - input_scale = self.find_weight_data(inputscale_initializer) + inputscale_initializer = find_by_name(input_scale_name, self.model.initializer()) + input_scale = self.tensor_proto_to_array(inputscale_initializer) # calcuate scale for bias @@ -792,7 +793,7 @@ class ONNXQuantizer: return quantized_bias_name - def _quantize_inputs(self, node, indices): + def quantize_inputs(self, node, indices): ''' Given a node, this function quantizes the inputs as follows: - If input is an initializer, quantize the initializer data, replace old initializer @@ -800,17 +801,15 @@ class ONNXQuantizer: - Else, add QuantizeLinear nodes to perform quantization parameter node: node being quantized in NodeProto format. parameter indices: input indices to quantize. - parameter new_nodes_list: List of new nodes created before processing this node. This is used to - check that two QuantizeLinear nodes are not being added for same input. return: (List of quantized input names, List of zero point names used for input quantization, List of scale names used for input quantization, List of new QuantizeLinear nodes created) ''' - quantized_input_names = [] - zero_point_names = [] scale_names = [] + zero_point_names = [] + quantized_input_names = [] nodes = [] for input_index in indices: @@ -819,27 +818,18 @@ class ONNXQuantizer: # Find if this input is already quantized if node_input in self.quantized_value_map: quantized_value = self.quantized_value_map[node_input] - qType = self.weight_qType if quantized_value.value_type == QuantizedValueType.Initializer else self.input_qType - if quantized_value.qType != qType: - raise ValueError( - "{} is being used by multiple nodes which are being quantized to different types. " - "This is not suported.", node_input) - - quantized_input_names.append(quantized_value.q_name) scale_names.append(quantized_value.scale_name) zero_point_names.append(quantized_value.zp_name) + quantized_input_names.append(quantized_value.q_name) continue # Quantize the input - initializer = _find_by_name(node_input, self.model.initializer()) + initializer = find_by_name(node_input, self.model.initializer()) if initializer is not None: - if node.op_type == "Conv": - weight = self._get_quantized_weight_convolution(initializer, self.weight_qType) - else: - weight = self._get_quantized_weight(initializer, self.weight_qType) + weight = self._get_quantized_weight(initializer, self.weight_qType) # Update graph - self._update_graph(weight) + self._update_weight(weight) quantized_input_names.append(weight.name + "_quantized") zero_point_names.append(weight.name + "_zero_point") @@ -864,6 +854,20 @@ class ONNXQuantizer: return (quantized_input_names, zero_point_names, scale_names, nodes) + def quantize_weight_per_channel(self, weight_name, axis): + # Find if this input is already quantized + if weight_name in self.quantized_value_map: + quantized_value = self.quantized_value_map[weight_name] + return (quantized_value.zp_name, quantized_value.q_name, quantized_value.scale_name) + else: + initializer = find_by_name(weight_name, self.model.initializer()) + if initializer is None: + raise ValueError("{} is not an initializer", weight_name) + + weight = self._get_quantized_weight_per_channel(initializer, self.weight_qType, axis) + self._update_weight(weight) + return (weight.name + "_quantized", weight.name + "_zero_point", weight.name + "_scale") + def _dequantize_value(self, value_name): ''' Given a value (input/output) which is quantized, add a DequantizeLinear node to dequantize diff --git a/onnxruntime/python/tools/quantization/operators/attention.py b/onnxruntime/python/tools/quantization/operators/attention.py index 6392d0a054..b8382f0337 100644 --- a/onnxruntime/python/tools/quantization/operators/attention.py +++ b/onnxruntime/python/tools/quantization/operators/attention.py @@ -1,6 +1,6 @@ import onnx from .base_operator import QuantOperatorBase -from ..quant_utils import _attribute_to_kwarg, ms_domain +from ..quant_utils import attribute_to_kwarg, ms_domain from onnx import onnx_pb as onnx_proto ''' Quantize Attention @@ -21,7 +21,7 @@ class AttentionQuant(QuantOperatorBase): assert (node.op_type == "Attention") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + self.quantizer.quantize_inputs(node, [0, 1]) qattention_name = "" if node.name == "" else node.name + "_quant" @@ -35,7 +35,7 @@ class AttentionQuant(QuantOperatorBase): kwargs = {} for attribute in node.attribute: - kwargs.update(_attribute_to_kwarg(attribute)) + kwargs.update(attribute_to_kwarg(attribute)) kwargs["domain"] = ms_domain qattention_node = onnx.helper.make_node("QAttention", inputs, node.output, qattention_name, **kwargs) nodes.append(qattention_node) diff --git a/onnxruntime/python/tools/quantization/operators/binary_op.py b/onnxruntime/python/tools/quantization/operators/binary_op.py index 171e724a3e..8782b302b4 100644 --- a/onnxruntime/python/tools/quantization/operators/binary_op.py +++ b/onnxruntime/python/tools/quantization/operators/binary_op.py @@ -1,6 +1,6 @@ import onnx from .base_operator import QuantOperatorBase -from ..quant_utils import _attribute_to_kwarg, ms_domain, QuantizedValue, QuantizedValueType +from ..quant_utils import attribute_to_kwarg, ms_domain, QuantizedValue, QuantizedValueType from onnx import onnx_pb as onnx_proto @@ -17,14 +17,14 @@ class QLinearBinaryOp(QuantOperatorBase): return super().quantize() (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + self.quantizer.quantize_inputs(node, [0, 1]) qlinear_binary_math_output = node.output[0] + "_quantized" qlinear_binary_math_name = node.name + "_quant" if node.name != "" else "" kwargs = {} for attribute in node.attribute: - kwargs.update(_attribute_to_kwarg(attribute)) + kwargs.update(attribute_to_kwarg(attribute)) kwargs["domain"] = ms_domain qlinear_binary_math_inputs = [] diff --git a/onnxruntime/python/tools/quantization/operators/conv.py b/onnxruntime/python/tools/quantization/operators/conv.py index c1b7cc1c04..a4c555eb36 100644 --- a/onnxruntime/python/tools/quantization/operators/conv.py +++ b/onnxruntime/python/tools/quantization/operators/conv.py @@ -1,6 +1,6 @@ import onnx from .base_operator import QuantOperatorBase -from ..quant_utils import _find_by_name, _get_mul_node, QuantizedValue, QuantizedValueType, _attribute_to_kwarg +from ..quant_utils import find_by_name, get_mul_node, QuantizedValue, QuantizedValueType, attribute_to_kwarg from onnx import onnx_pb as onnx_proto @@ -13,7 +13,7 @@ class ConInteger(QuantOperatorBase): assert (node.op_type == "Conv") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + self.quantizer.quantize_inputs(node, [0, 1]) # quantize bias if exist quantized_bias_name = "" @@ -27,7 +27,7 @@ class ConInteger(QuantOperatorBase): kwargs = {} for attribute in node.attribute: - kwargs.update(_attribute_to_kwarg(attribute)) + kwargs.update(attribute_to_kwarg(attribute)) conv_integer_node = onnx.helper.make_node("ConvInteger", quantized_input_names + zero_point_names, [conv_integer_output], conv_integer_name, **kwargs) nodes.append(conv_integer_node) @@ -51,9 +51,9 @@ class ConInteger(QuantOperatorBase): else: scales_mul_op = scale_names[0] + "_" + scale_names[1] + "_mul" - scales_mul_node = _find_by_name(scales_mul_op, self.nodes) + scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes) if scales_mul_node is None: - scales_mul_node = _get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op) + scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op) nodes.append(scales_mul_node) scales_mul_op_output = scales_mul_node.output[0] @@ -61,9 +61,9 @@ class ConInteger(QuantOperatorBase): # Add mul operation to multiply mul_scales_op result with output of ConvInteger # and make the output of this node the same as output of original conv node. output_scale_mul_op = conv_integer_name + "_output_scale_mul" if conv_integer_name != "" else "" - nodes.append(_get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op)) + nodes.append(get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op)) - self.new_nodes += nodes + self.quantizer.new_nodes += nodes class QLinearCov(QuantOperatorBase): @@ -74,8 +74,16 @@ class QLinearCov(QuantOperatorBase): node = self.node assert (node.op_type == "Conv") - (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + if self.quantizer.is_input_a_weight(node.input[1]): + (quantized_input_names, zero_point_names, scale_names, nodes) = \ + self.quantizer.quantize_inputs(node, [0]) + quant_weight_tuple = self.quantizer.quantize_weight_per_channel(node.input[1], 0) + quantized_input_names.append(quant_weight_tuple[0]) + zero_point_names.append(quant_weight_tuple[1]) + scale_names.append(quant_weight_tuple[2]) + else: + (quantized_input_names, zero_point_names, scale_names, nodes) = \ + self.quantizer.quantize_inputs(node, [0, 1]) quantized_bias_name = "" bias_present = False @@ -94,7 +102,7 @@ class QLinearCov(QuantOperatorBase): kwargs = {} for attribute in node.attribute: - kwargs.update(_attribute_to_kwarg(attribute)) + kwargs.update(attribute_to_kwarg(attribute)) qlinear_conv_inputs = [] # Input 0 qlinear_conv_inputs.append(quantized_input_names[0]) diff --git a/onnxruntime/python/tools/quantization/operators/embed_layernorm.py b/onnxruntime/python/tools/quantization/operators/embed_layernorm.py index 8d9759f287..b110e67115 100644 --- a/onnxruntime/python/tools/quantization/operators/embed_layernorm.py +++ b/onnxruntime/python/tools/quantization/operators/embed_layernorm.py @@ -15,7 +15,7 @@ class EmbedLayerNormalizationQuant(QuantOperatorBase): assert (node.op_type == "EmbedLayerNormalization") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [2, 3, 4]) + self.quantizer.quantize_inputs(node, [2, 3, 4]) nodes.append(node) diff --git a/onnxruntime/python/tools/quantization/operators/gather.py b/onnxruntime/python/tools/quantization/operators/gather.py index 4d257ed448..9f9595775b 100644 --- a/onnxruntime/python/tools/quantization/operators/gather.py +++ b/onnxruntime/python/tools/quantization/operators/gather.py @@ -19,7 +19,7 @@ class GatherQuant(QuantOperatorBase): return (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0]) + self.quantizer.quantize_inputs(node, [0]) gather_new_output = node.output[0] + "_quantized" diff --git a/onnxruntime/python/tools/quantization/operators/matmul.py b/onnxruntime/python/tools/quantization/operators/matmul.py index b9358c9562..d765a043e1 100644 --- a/onnxruntime/python/tools/quantization/operators/matmul.py +++ b/onnxruntime/python/tools/quantization/operators/matmul.py @@ -1,6 +1,6 @@ import onnx from .base_operator import QuantOperatorBase -from ..quant_utils import _find_by_name, _get_mul_node, QuantizedValue, QuantizedValueType +from ..quant_utils import find_by_name, get_mul_node, QuantizedValue, QuantizedValueType from onnx import onnx_pb as onnx_proto ''' Used when quantize mode is QuantizationMode.IntegerOps. @@ -16,7 +16,7 @@ class MatMulInteger(QuantOperatorBase): assert (node.op_type == "MatMul") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + self.quantizer.quantize_inputs(node, [0, 1]) matmul_integer_output = node.output[0] + "_quantized" matmul_integer_name = node.name + "_quant" if node.name != "" else "" @@ -36,9 +36,9 @@ class MatMulInteger(QuantOperatorBase): scales_mul_op = matmul_integer_name + "_scales_mul" if matmul_integer_name != "" else scale_names[ 0] + "_" + scale_names[1] + "_mul" - scales_mul_node = _find_by_name(scales_mul_op, self.quantizer.new_nodes) + scales_mul_node = find_by_name(scales_mul_op, self.quantizer.new_nodes) if scales_mul_node is None: - scales_mul_node = _get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op) + scales_mul_node = get_mul_node(scale_names, scales_mul_op + ":0", scales_mul_op) nodes.append(scales_mul_node) scales_mul_op_output = scales_mul_node.output[0] @@ -48,7 +48,7 @@ class MatMulInteger(QuantOperatorBase): output_scale_mul_op = "" if matmul_integer_name != "": output_scale_mul_op = matmul_integer_name + "_output_scale_mul" - nodes.append(_get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op)) + nodes.append(get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op)) self.quantizer.new_nodes += nodes @@ -66,7 +66,7 @@ class QLinearMatMul(QuantOperatorBase): assert (node.op_type == "MatMul") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self.quantizer._quantize_inputs(node, [0, 1]) + self.quantizer.quantize_inputs(node, [0, 1]) data_found, output_scale_name, output_zp_name, _, _ = \ self.quantizer._get_quantization_params(node.output[0]) diff --git a/onnxruntime/python/tools/quantization/quant_utils.py b/onnxruntime/python/tools/quantization/quant_utils.py index d4e1f1b7c8..1d164dd874 100644 --- a/onnxruntime/python/tools/quantization/quant_utils.py +++ b/onnxruntime/python/tools/quantization/quant_utils.py @@ -104,7 +104,7 @@ class QuantizedValue: self.qType = qType -def _attribute_to_kwarg(attribute): +def attribute_to_kwarg(attribute): ''' Convert attribute to kwarg format for use with onnx.helper.make_node. :parameter attribute: attribute in AttributeProto format. @@ -141,7 +141,7 @@ def _attribute_to_kwarg(attribute): return {attribute.name: value} -def _find_by_name(item_name, item_list): +def find_by_name(item_name, item_list): ''' Helper function to find item by name in a list. parameter item_name: name of the item. @@ -152,7 +152,7 @@ def _find_by_name(item_name, item_list): return items[0] if len(items) > 0 else None -def _get_elem_index(elem_name, elem_list): +def get_elem_index(elem_name, elem_list): ''' Helper function to return index of an item in a node list ''' @@ -163,7 +163,7 @@ def _get_elem_index(elem_name, elem_list): return elem_idx -def _get_mul_node(inputs, output, name): +def get_mul_node(inputs, output, name): ''' Helper function to create a Mul node. parameter inputs: list of input names. @@ -174,7 +174,7 @@ def _get_mul_node(inputs, output, name): return onnx.helper.make_node("Mul", inputs, [output], name) -def _generate_identified_filename(filename: Path, identifier: str) -> Path: +def generate_identified_filename(filename: Path, identifier: str) -> Path: ''' Helper function to generate a identifiable filepath by concatenating the given identifier as a suffix. ''' diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 37bbdf4178..631b38a597 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -16,7 +16,7 @@ from onnx import shape_inference from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from .quant_utils import QuantizationMode, QuantizedValueType, QuantizedInitializer, QuantizedValue, quantization_modes -from .quant_utils import _find_by_name, _get_elem_index, _get_mul_node, _generate_identified_filename, _attribute_to_kwarg +from .quant_utils import find_by_name, get_elem_index, get_mul_node, generate_identified_filename, attribute_to_kwarg from .quant_utils import QuantType from .registry import CreateOpQuantizer, CreateDefaultOpQuantizer, QLinearOpsRegistry, IntegerOpsRegistry @@ -32,7 +32,7 @@ def optimize_model(model_path: Path): parameter model_path: path to the original onnx model return: optimized onnx model ''' - opt_model_path = _generate_identified_filename(model_path, "-opt") + opt_model_path = generate_identified_filename(model_path, "-opt") sess_option = SessionOptions() sess_option.optimized_model_filepath = opt_model_path.as_posix() sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC @@ -126,6 +126,7 @@ def quantize_static(model_input, calibration_data_reader: CalibrationDataReader, op_types_to_quantize=[], per_channel=False, + reduce_range=False, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], @@ -138,6 +139,7 @@ def quantize_static(model_input, :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default. :param op_types: operators to quantize :param per_channel: quantize weights per channel + :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode :param activation_type: quantization data type of activation :param weight_type: quantization data type of weight :param nodes_to_quantize: @@ -153,8 +155,8 @@ def quantize_static(model_input, when it is not None. ''' - if activation_type != QuantType.QUInt8 or weight_type != QuantType.QUInt8: - raise ValueError("Static quantization only support uint8 now.") + if activation_type != QuantType.QUInt8: + raise ValueError("Static quantization only support uint8 for activation now.") 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 @@ -169,6 +171,7 @@ def quantize_static(model_input, quantizer = ONNXQuantizer( onnx.load(model_input), per_channel, + reduce_range, mode, True, # static weight_qType, @@ -186,6 +189,7 @@ def quantize_dynamic(model_input: Path, model_output: Path, op_types_to_quantize=[], per_channel=False, + reduce_range=False, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], @@ -194,8 +198,9 @@ def quantize_dynamic(model_input: Path, Given an onnx model, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize :param model_output: file path of quantized model - :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default. + :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default :param per_channel: quantize weights per channel + :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode :param nbits: number of bits to represent quantized data. Currently only supporting 8-bit types :param activation_type: quantization data type of activation :param weight_type: quantization data type of weight @@ -225,6 +230,7 @@ def quantize_dynamic(model_input: Path, quantizer = ONNXQuantizer( optimized_model, per_channel, + reduce_range, mode, False, #static weight_qType, @@ -242,6 +248,7 @@ def quantize_qat(model_input: Path, model_output: Path, op_types_to_quantize=[], per_channel=False, + reduce_range=False, activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8, nodes_to_quantize=[], @@ -250,8 +257,9 @@ def quantize_qat(model_input: Path, 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 :param model_output: file path of quantized model - :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default. + :param op_types_to_quantize: specify the types of operators to quantize, like ['Conv'] to quantize Conv only. It quantizes all supported operators by default :param per_channel: quantize weights per channel + :param reduce_range: quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode :param activation_type: quantization data type of activation :param nodes_to_quantize: List of nodes names to quantize. When this list is not None only the nodes in this list @@ -279,6 +287,7 @@ def quantize_qat(model_input: Path, quantizer = ONNXQuantizer( optimized_model, per_channel, + reduce_range, mode, False, #static weight_qType,