diff --git a/onnxruntime/python/tools/quantization/README.md b/onnxruntime/python/tools/quantization/README.md index 6f0ccfe48e..abbd30fdd8 100644 --- a/onnxruntime/python/tools/quantization/README.md +++ b/onnxruntime/python/tools/quantization/README.md @@ -1,7 +1,7 @@ # Quantization tool Overview This tool supports quantization of an onnx model. quantize() takes a model in ModelProto format and returns the quantized model in ModelProto format. -## Quantize an onnx model +## Quantize an ONNX model ```python import onnx from quantize import quantize, QuantizationMode @@ -17,11 +17,11 @@ onnx.save(quantized_model, 'path/to/the/quantized_model.onnx') ## Examples of various quantization modes - **QuantizationMode.IntegerOps with static input quantization**: - Quantize using integer ops. Inputs/activations are quantized using static scale and zero point values which are specified through "input_quantization_params" option. + Quantize using integer ops. Inputs/activations are quantized using static scale and zero point values which are specified through "quantization_params" option. ```python quantized_model = quantize(model, quantization_mode=QuantizationMode.IntegerOps, static=True, - input_quantization_params={ + quantization_params={ 'input_1': [np.uint8(113), np.float32(0.05)] }) ``` @@ -38,21 +38,19 @@ onnx.save(quantized_model, 'path/to/the/quantized_model.onnx') ```python quantized_model = quantize(model, quantization_mode=QuantizationMode.QLinearOps, static=True, - input_quantization_params={ + quantization_params={ 'input_1': [np.uint8(113), np.float32(0.05)] - }, - output_quantization_params={ 'output_1': [np.uint8(113), np.float32(0.05)] }) ``` - **QuantizationMode.QLinearOps with dynamic input quantization**: Quantize using QLinear ops. Inputs/activations are quantized using dynamic scale and zero point values which are computed while running the model. - Output scale and zero point values have to be specified using "output_quantization_params" option. + Output scale and zero point values have to be specified using "quantization_params" option. ```python quantized_model = quantize(model, quantization_mode=QuantizationMode.QLinearOps, static=False, - output_quantization_params={ + quantization_params={ 'output_1': [np.uint8(113), np.float32(0.05)] }) ``` @@ -71,15 +69,19 @@ See below for a description of all the options to quantize(): *QuantizationMode.IntegerOps*: Quantize using integer ops. Only [ConvInteger](https://github.com/onnx/onnx/blob/master/docs/Operators.md#ConvInteger) and [MatMulInteger](https://github.com/onnx/onnx/blob/master/docs/Operators.md#MatMulInteger) ops are supported now. *QuantizationMode.QLinearOps*: Quantize using QLinear ops. Only [QLinearConv](https://github.com/onnx/onnx/blob/master/docs/Operators.md#qlinearconv) and [QLinearMatMul](https://github.com/onnx/onnx/blob/master/docs/Operators.md#QLinearMatMul) ops are supported now. - **static**: *default:False* -If True, the inputs/activations are quantized using static scale and zero point values specified through input_quantization_params. +If True, the inputs/activations are quantized using static scale and zero point values specified through quantization_params. If False, the inputs/activations are quantized using dynamic scale and zero point values computed while running the model. - **asymmetric_input_types**: *default: False* If True, weights are quantized into signed integers and inputs/activations into unsigned integers. If False, weights and inputs/activations are quantized into unsigned integers. -- **input_quantization_params**: *default: None* - Dictionary to specify the zero point and scale values for inputs to conv and matmul nodes. +- **force_fusions**: *default: False* + If True, nodes added for dynamic quantization are fused. + If False, no fusion is applied for nodes which are added for dynamic quantization. + This optimization is available from opset 11. +- **quantization_params**: *default: None* + Dictionary to specify the zero point and scale values for inputs to and outputs from conv and matmul nodes. Should be specified when static is set to True. - The input_quantization_params should be specified in the following format: + The quantization_params should be specified in the following format: { "input_name": [zero_point, scale] }. @@ -89,16 +91,11 @@ If False, the inputs/activations are quantized using dynamic scale and zero poin 'resnet_model/Relu_1:0': [np.uint8(0), np.float32(0.019539741799235344)], 'resnet_model/Relu_2:0': [np.uint8(0), np.float32(0.011359662748873234)] } -- **output_quantization_params**: *default: None* - Dictionary to specify the zero point and scale values for outputs of conv and matmul nodes. - Should be specified in QuantizationMode.QLinearOps mode. - The output_quantization_params should be specified in the following format: - { - "output_name": [zero_point, scale] - } - zero_point can be of type np.uint8/np.int8 and scale should be of type np.float32. - example: - { - 'resnet_model/Relu_3:0': [np.int8(0), np.float32(0.011359662748873234)], - 'resnet_model/Relu_4:0': [np.uint8(0), np.float32(0.011359662748873234)] - } \ No newline at end of file +- **nodes_to quantize**: *default: None* + List of nodes names to quantize. When this list is not None only the nodes in this list + are quantized. + exmaple: + [ + 'Cov__224', + 'Conv__252' + ] diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index ce09c8e18f..e886fd0df9 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -14,7 +14,7 @@ from onnx import onnx_pb as onnx_proto __producer__ = "onnx.quantize" __version__ = "0.1.0" onnx_domain = "ai.onnx" -onnx_op_set_version = 10 +onnx_op_set_version = 11 type_to_name = { 1: "FLOAT", @@ -41,26 +41,10 @@ class QuantizationMode(): IntegerOps = 0 QLinearOps = 1 -# Data Quantization mode -# Linear_NonScaled: Quantize data using linear, non scaled tranformation. -# Linear_Scaled: Quantize data using linear, scaled transformation. -class DataQuantizationMode(): - Linear_NonScaled = 0 - Linear_Scaled = 1 - - @staticmethod - def mode_for_data_type(data_type): - return DataQuantizationMode.Linear_Scaled if data_type == onnx_proto.TensorProto.INT8\ - else DataQuantizationMode.Linear_NonScaled - - quantization_modes = [getattr(QuantizationMode, attr) for attr in dir(QuantizationMode)\ if not callable(getattr(QuantizationMode, attr)) and not attr.startswith("__")] -data_quantization_modes = [getattr(DataQuantizationMode, attr) for attr in dir(DataQuantizationMode)\ - if not callable(getattr(DataQuantizationMode, attr)) and not attr.startswith("__")] - -class Weight: +class QuantizedInitializer: ''' Represents a linearly quantized weight input from ONNX operators ''' @@ -78,16 +62,33 @@ class Weight: # If empty, single zero point and scales computed from a single rmin and rmax self.qType = qType # type of quantized data. +class QuantizedValueType(): + Input = 0 + Initializer = 1 -def quantize_data(data, quantize_range, mode=DataQuantizationMode.Linear_NonScaled): +class QuantizedValue: + ''' + Represents a linearly quantized value (input\output\intializer) + ''' + def __init__(self, name, new_quantized_name, scale_name, zero_point_name, quantized_value_type, axis=None, + qType=onnx_proto.TensorProto.UINT8): + self.original_name = name + self.q_name = new_quantized_name + self.scale_name = scale_name + self.zp_name = zero_point_name + self.value_type = quantized_value_type + self.axis = axis + self.qType = qType + +def quantize_data(data, quantize_range, qType): ''' :parameter quantize_range: list of data to weight pack. - :parameter mode: mode to quantize data of type DataQuantizationMode + :parameter qType: data type to quantize to. Supported types UINT8 and INT8 :return: minimum, maximum, zero point, scale, and quantized weights To pack weights, we compute a linear transformation - - in non-scaled mode, from [rmin, rmax] -> [0, 2^{b-1}] and - - in scaled mode, from [-m , m] -> [-(2^{b-1}-1), 2^{b-1}-1] where + - when data type == uint8 mode, from [rmin, rmax] -> [0, 2^{b-1}] and + - when data type == int8, from [-m , m] -> [-(2^{b-1}-1), 2^{b-1}-1] where m = max(abs(rmin), abs(rmax)) and add necessary intermediate nodes to trasnform quantized weight to full weight using the equation @@ -100,15 +101,18 @@ def quantize_data(data, quantize_range, mode=DataQuantizationMode.Linear_NonScal rmin = min(min(data), 0) rmax = max(max(data), 0) - if mode == DataQuantizationMode.Linear_Scaled: + if qType == onnx_proto.TensorProto.INT8: max_range = max(abs(rmin), abs(rmax)) scale = (float(max_range)*2) / quantize_range zero_point = 0 quantized_data = (np.asarray(data) / scale).round().astype('b') #signed byte type - else: + elif qType == onnx_proto.TensorProto.UINT8: scale = (float(rmax) - rmin) / quantize_range if rmin != rmax else 1 zero_point = round((0 - rmin) / scale) # round to nearest integer quantized_data = ((np.asarray(data) / scale).round() + zero_point).astype('B') # unsigned byte type + else: + raise ValueError("Unexpected data type {} requested. Only INT8 and UINT8 are supported.") + return rmin, rmax, zero_point, scale, quantized_data @@ -225,16 +229,16 @@ def _find_nodes_using_initializer(graph, initializer): return result class ONNXQuantizer: - def __init__(self, model, per_channel, mode, static, weight_qType, input_qType, - input_quantization_params, output_quantization_params, nodes_to_quantize): + def __init__(self, model, per_channel, mode, static, fuse_dynamic_quant, weight_qType, input_qType, + quantization_params, nodes_to_quantize): self.model = model - self.per_channel = per_channel # weight-pack per channel - self.weight_qType = weight_qType # quantize data type + self.per_channel = per_channel # weight-pack per channel self.mode = mode # QuantizationMode.Value self.static = static # use static quantization for inputs. + self.fuse_dynamic_quant = fuse_dynamic_quant self.input_qType = input_qType # quantize input type - self.input_quantization_params = input_quantization_params # zero point and scale values for node inputs. - self.output_quantization_params = output_quantization_params # zero point and scale values for node outputs. + self.weight_qType = weight_qType # quantize data type + self.quantization_params = quantization_params self.nodes_to_quantize = nodes_to_quantize # specific nodes to quantize if not self.mode in quantization_modes: @@ -242,31 +246,35 @@ class ONNXQuantizer: # QuantizeRange tensor name and zero tensor name for scale and zero point calculation. # Used when static is False - self.fixed_qrange_non_scaled_name = "fixed_quantization_range_non_scaled" - self.fixed_qrange_scaled_name = "fixed_quantization_range_scaled" - # In non scaled mode, to compute zero point, we subtract rmin from 0 (represented by fixed_zero_name tensor) + self.fixed_qrange_uint8_name = "fixed_quantization_range_uint8" + self.fixed_qrange_int8_name = "fixed_quantization_range_int8" + # For uint8 data-type, to compute zero point, we subtract rmin from 0 (represented by fixed_zero_name tensor) self.fixed_zero_name = "fixed_zero" - # In scaled mode, zero point is always zero (respresented by fixed_zero_point_name tensor) + # For int8 data-type, zero point is always zero (respresented by fixed_zero_point_name tensor) self.fixed_zero_zp_name = "fixed_zero_zp" - # List of weights quantized. + # List of quantized weights self._quantized_weights = [] + # Map of all original value names to quantized value names + self.quantized_value_map = {} def quantize_model(self): # Create a new topologically sorted list for quantizing a model new_list = [] for node in self.model.graph.node: if self.nodes_to_quantize is not None and node.name not in self.nodes_to_quantize: - new_list.append(node) + new_list +=self._handle_other_ops(node, new_list) else: - if node.op_type == 'Conv' and len(node.input) == 2: + if node.op_type == 'Conv': new_list += self._quantize_convolution(node, new_list) elif node.op_type == 'MatMul': new_list += self._quantize_matmul(node, new_list) elif node.op_type == 'Gather': new_list += self._quantize_gather_ops(node, new_list) + elif node.op_type == 'Relu' or node.op_type == 'Clip': + new_list +=self._handle_activation_ops(node, new_list) else: - new_list.append(node) + new_list +=self._handle_other_ops(node, new_list) # extend is used to append to the list for a protobuf fields # https://developers.google.com/protocol-buffers/docs/reference/python-generated?csw=1#fields @@ -322,9 +330,11 @@ class ONNXQuantizer: - remove old weight input, update with new inputs for quantized weight, zero point, and scale This function does NOT update the nodes in the graph, just initializers and inputs ''' - packed_weight_name = weight.name + '_quantized' - scale_name = weight.name + '_scale' - zero_point_name = weight.name + '_zero_point' + quantized_value = self.quantized_value_map[weight.name] + assert(quantized_value is not None) + packed_weight_name = quantized_value.q_name + scale_name = quantized_value.scale_name + zero_point_name = quantized_value.zp_name # Update packed weight, zero point, and scale initializers packed_weight_np_data = np.asarray(weight.quantized_data, @@ -360,9 +370,15 @@ class ONNXQuantizer: ''' 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), mode=DataQuantizationMode.mode_for_data_type(qType)) - weight = Weight(initializer.name, initializer, [rmin], [rmax], [zero_point], [scale], + _get_qrange_for_qType(qType), qType) + weight = QuantizedInitializer(initializer.name, initializer, [rmin], [rmax], [zero_point], [scale], weights_data, quantized_weights_data, axis=None, qType=qType) + + # Log entry for this quantized weight + assert(weight.name not in self.quantized_value_map) + quantized_value = QuantizedValue(weight.name, weight.name + "_quantized", weight.name + "_scale", weight.name + "_zero_point", QuantizedValueType.Initializer, None, qType) + self.quantized_value_map[weight.name] = quantized_value + return weight def _get_quantized_weight_convolution(self, initializer, qType): @@ -388,7 +404,7 @@ class ONNXQuantizer: # for each channel, compute quantization data. Assuming (M x C/group x kH x kW) per_channel_data = np_data[i,:,:,:].flatten() rmin, rmax, zero_point, scale, quantized_per_channel_data = quantize_data(per_channel_data.flatten().tolist(), - _get_qrange_for_qType(qType), mode=DataQuantizationMode.mode_for_data_type(qType)) + _get_qrange_for_qType(qType), qType) rmin_list.append(rmin) rmax_list.append(rmax) zero_point_list.append(zero_point) @@ -403,8 +419,14 @@ class ONNXQuantizer: channel_weights = np.asarray(quantized_per_channel_data_list[i]).reshape(reshape_dims) quantized_weights = np.concatenate((quantized_weights, channel_weights), axis=0) - weight = Weight(initializer.name, initializer, rmin_list, rmax_list, zero_point_list, + weight = QuantizedInitializer(initializer.name, initializer, rmin_list, rmax_list, zero_point_list, scale_list, weights, quantized_weights.flatten().tolist(), channel_index, qType) + + # Make entry for this quantized weight + assert(weight.name not in self.quantized_value_map) + quantized_value = QuantizedValue(weight.name, weight.name + "_quantized", weight.name + "_scale", weight.name + "_zero_point", QuantizedValueType.Initializer, None, qType) + self.quantized_value_map[weight.name] = quantized_value + return weight def _get_dynamic_input_quantization_params(self, input_name, nodes_list, qType): @@ -415,16 +437,14 @@ class ONNXQuantizer: parameter qType: type to quantize to. return: scale_name, zero_point_name, scale_shape, zero_point_shape. ''' - mode = DataQuantizationMode.mode_for_data_type(qType) - if mode == DataQuantizationMode.Linear_Scaled: - return self._get_dynamic_input_quantization_params_scaled(input_name, nodes_list) + if qType == onnx_proto.TensorProto.INT8: + return self._get_dynamic_input_quantization_params_int8(input_name, nodes_list) - return self._get_dynamic_input_quantization_params_non_scaled(input_name, nodes_list) + return self._get_dynamic_input_quantization_params_uint8(input_name, nodes_list) - def _get_dynamic_input_quantization_params_scaled(self, input_name, nodes_list): + def _get_dynamic_input_quantization_params_int8(self, input_name, nodes_list): ''' - Create nodes for dynamic quantization of input and add them to nodes_list - in DataQuantizationMode.Linear_Scaled + Create nodes for dynamic quantization of input to nit8 and add them to nodes_list parameter input_name: Name of the input. parameter nodes_list: new nodes are appended to this list. return: scale_name, zero_point_name, scale_shape, zero_point_shape. @@ -461,10 +481,10 @@ class ONNXQuantizer: [abs_max_name + ":0"], abs_max_name) nodes_list.append(abs_max_node) # and divide by (quantize_range/2.0) which will be equal to max(...)*2.0/quantize_range - _add_initializer_if_not_present(self.model.graph, self.fixed_qrange_scaled_name, + _add_initializer_if_not_present(self.model.graph, self.fixed_qrange_int8_name, [_get_qrange_for_qType(qType)/2.0], [], onnx_proto.TensorProto.FLOAT) scale_div_name = input_name + "scale_Div" - scale_div_node = onnx.helper.make_node("Div", [abs_max_node.output[0], self.fixed_qrange_scaled_name], + scale_div_node = onnx.helper.make_node("Div", [abs_max_node.output[0], self.fixed_qrange_int8_name], [input_scale_name], scale_div_name) nodes_list.append(scale_div_node) @@ -474,10 +494,9 @@ class ONNXQuantizer: return input_scale_name, self.fixed_zero_zp_name, [], [] - def _get_dynamic_input_quantization_params_non_scaled(self, input_name, nodes_list): + def _get_dynamic_input_quantization_params_uint8(self, input_name, nodes_list): ''' - Create nodes for dynamic quantization of input and add them to nodes_list - in DataQuantizationMode.Linear_NonScaled + Create nodes for dynamic quantization of input to uint8 and add them to nodes_list parameter input_name: Name of the input. parameter nodes_list: new nodes are appended to this list. return: scale_name, zero_point_name, scale_shape, zero_point_shape. @@ -498,7 +517,7 @@ class ONNXQuantizer: nodes_list.append(reduce_max_node) # Add tensors for quantize range and zero value. - _add_initializer_if_not_present(self.model.graph, self.fixed_qrange_non_scaled_name, + _add_initializer_if_not_present(self.model.graph, self.fixed_qrange_uint8_name, [_get_qrange_for_qType(qType)], [], onnx_proto.TensorProto.FLOAT) _add_initializer_if_not_present(self.model.graph, self.fixed_zero_name, [0.0], [], onnx_proto.TensorProto.FLOAT) @@ -511,7 +530,7 @@ class ONNXQuantizer: nodes_list.append(scale_sub_node) # and divide by quantize range scale_div_name = input_name + "_scale_Div" - scale_div_node = onnx.helper.make_node("Div", [scale_sub_node.output[0], self.fixed_qrange_non_scaled_name], + scale_div_node = onnx.helper.make_node("Div", [scale_sub_node.output[0], self.fixed_qrange_uint8_name], [input_scale_name], scale_div_name) nodes_list.append(scale_div_node) @@ -539,67 +558,17 @@ class ONNXQuantizer: return input_scale_name, input_zp_name, [], [] - def _get_static_input_quantization_params(self, input_name, qType): - ''' - Create initializers and inputs in the graph for static quantization of input. - - Zero point and scale values are obtained from self.input_quantization_params if specified. - ValueError is thrown otherwise. - - parameter input_name: Name of the input. - parameter qType: type to quantize to. - return: scale_name, zero_point_name, scale_shape, zero_point_shape. - ''' - if self.input_quantization_params is None or input_name not in self.input_quantization_params: - raise ValueError("Quantization parameters are not specified for input {}.".format(input_name)) - params = self.input_quantization_params[input_name] - if params is None or len(params) != 2: - raise ValueError("Quantization parameters should contain zero point and scale. " - "Specified values for input {}: {}".format(input_name, params)) - - if not np.isscalar(params[0]): - raise ValueError("Zero point for input {} should be a scalar value. Value specified: {}".format( - input_name, params[0])) - if not np.isscalar(params[1]): - raise ValueError("Scale for input {} should be a scalar value. Value specified: {}".format( - input_name, params[1])) - - zero_point_values = [params[0].item()] - zero_point_shape = [] - zero_point_name = input_name + "_zero_point" - - zero_point_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[params[0].dtype] - if zero_point_type != qType: - raise ValueError("Zero point and input data types should be the same. " - "Zero point for input {} is specified as {}, but input is being quantized to {}." - .format(input_name, params[0].dtype, onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[qType])) - - scale_values = [params[1].item()] - scale_shape = [] - scale_name = input_name + "_scale" - - # Add initializers - _add_initializer_if_not_present(self.model.graph, zero_point_name, zero_point_values, - zero_point_shape, qType) - _add_initializer_if_not_present(self.model.graph, scale_name, scale_values, - scale_shape, onnx_proto.TensorProto.FLOAT) - - return scale_name, zero_point_name, scale_shape, zero_point_shape - - def _get_output_quantization_params(self, output_name): + def _get_quantization_params(self, param_name): ''' Create initializers and inputs in the graph for zero point and scale of output. - Used when mode is QuantizationMode.QLinearOps. - - Zero point and scale values are obtained from self.output_quantization_params if specified. - ValueError is thrown otherwise. + Zero point and scale values are obtained from self.quantization_params if specified. parameter output_name: Name of the output. return: scale_name, zero_point_name, scale_shape, zero_point_shape. - ''' - if self.output_quantization_params is None or output_name not in self.output_quantization_params: - raise ValueError("Quantization parameters are not specified for output {}.".format(output_name)) - params = self.output_quantization_params[output_name] + ''' + if self.quantization_params is None or param_name not in self.quantization_params: + return False, "", "", "", "" + params = self.quantization_params[param_name] if params is None or len(params) != 2: raise ValueError("Quantization parameters should contain zero point and scale. " "Specified values for output {}: {}".format(output_name, params)) @@ -613,12 +582,12 @@ class ONNXQuantizer: zero_point_values = [params[0].item()] zero_point_shape = [] - zero_point_name = output_name + "_zero_point" + zero_point_name = param_name + "_zero_point" zero_point_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[params[0].dtype] scale_values = [params[1].item()] scale_shape = [] - scale_name = output_name + "_scale" + scale_name = param_name + "_scale" # Add initializers _add_initializer_if_not_present(self.model.graph, zero_point_name, zero_point_values, zero_point_shape, @@ -626,7 +595,7 @@ class ONNXQuantizer: _add_initializer_if_not_present(self.model.graph, scale_name, scale_values, scale_shape, onnx_proto.TensorProto.FLOAT) - return scale_name, zero_point_name, scale_shape, zero_point_shape + return True, scale_name, zero_point_name, scale_shape, zero_point_shape def _get_quantize_input_nodes(self, node, input_index, qType): ''' @@ -640,20 +609,72 @@ class ONNXQuantizer: return: List of newly created nodes in NodeProto format. ''' input_name = node.input[input_index] - - nodes = [] - if self.static: - scale_name, zp_name, scale_shape, zp_shape = \ - self._get_static_input_quantization_params(input_name, qType) - else: - scale_name, zp_name, scale_shape, zp_shape = \ - self._get_dynamic_input_quantization_params(input_name, nodes, qType) - - # Add QuantizeLinear Node output_name = input_name + "_quantized" - qlinear_node = onnx.helper.make_node("QuantizeLinear", [input_name, scale_name, zp_name], - [output_name], input_name + "_QuantizeLinear") - return nodes + [qlinear_node] + + data_found, scale_name, zp_name, scale_shape, zp_shape = \ + self._get_quantization_params(input_name) + + if self.static: + if data_found == False: + raise ValueError("Quantization parameters are not specified for param {}." + "In static mode quantization params for inputs and outputs of odes to be quantized are required.".format(input_name)) + + qlinear_node = onnx.helper.make_node("QuantizeLinear", [input_name, scale_name, zp_name], + [output_name], input_name + "_QuantizeLinear") + + return [qlinear_node] + + else: + if data_found == True: + qlinear_node = onnx.helper.make_node("QuantizeLinear", [input_name, scale_name, zp_name], + [output_name], input_name + "_QuantizeLinear") + else: + # Scale and Zero Points not available for this input. Add nodes to dynamically compute it + if self.fuse_dynamic_quant and qType == onnx_proto.TensorProto.UINT8: + scale_name = input_name + "_scale" + zeropoint_name = input_name + "_zero_point" + qlinear_node = onnx.helper.make_node("DynamicQuantizeLinear", [input_name], + [output_name, scale_name, zeropoint_name], input_name + "_QuantizeLinear") + return [qlinear_node] + + else: + nodes = [] + scale_name, zp_name, scale_shape, zp_shape = \ + self._get_dynamic_input_quantization_params(input_name, nodes, qType) + qlinear_node = onnx.helper.make_node("QuantizeLinear", [input_name, scale_name, zp_name], + [output_name], input_name + "_QuantizeLinear") + + return nodes + [qlinear_node] + + 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 + parameter node: current node (Conv) + parameter last_output: output of previous node (input to bias add) + return: the name of output + ''' + # Add an Add operation for bias + # Add reshape for correct broadcase + reshape_input = [quantized_bias_name] + + # Add tensors for the shape to be reshaped to + _add_initializer_if_not_present(self.model.graph, "reshape_shape", + [1,-1,1,1], [4], onnx_proto.TensorProto.INT64) + reshape_input.append('reshape_shape') + reshape_op_output = node.output[0] + "_reshape" + reshape_node = onnx.helper.make_node("Reshape", reshape_input, [reshape_op_output], + quantized_bias_name+"reshape") + nodes.append(reshape_node) + + bias_add_input = [last_output] + bias_add_input.append(reshape_op_output) + add_node_output = node.output[0] + "_bias_add" + add_node = onnx.helper.make_node("Add", bias_add_input, [add_node_output], + quantized_bias_name + "bias_add") + nodes.append(add_node) + return add_node_output def _update_unsupported_nodes_using_weight(self, weight, new_nodes_list): '''Find all nodes using a weight that do not support quantization and @@ -686,35 +707,91 @@ class ONNXQuantizer: return nodes_list - def _is_quantized(self, weight): + def _dynamic_quantize_bias(self, input_name, weight_scale_name, bias_name, quantized_bias_name, new_node_list): ''' - Check if this weight is already quantized to the expected type and quantization axis. - If it is already quantized to (type, axis) different from expected values, - this function will throw an exception and stop the quantization. - - parameter weight: Weight object. - return: Boolean indicating if quantized weight is already added to graph. + Adds series of nodes required to quantize the bias dynamically. + parameter input_name: Input name + parameter weight_scale_name: Weight scale. + parameter bias_scale_name: Bias to quantize. + parameter quantied_bias_name: Output name to use for quantized bias. ''' - quantized_initializer_name = weight.name + "_quantized" - quantized_initializer = _find_by_name(quantized_initializer_name, self.model.graph.initializer) - zero_point = _find_by_name(weight.name + "_zero_point", self.model.graph.initializer) - if quantized_initializer is None: - return False + qType = onnx_proto.TensorProto.INT32 + + input_scale_name = input_name + "_scale" + bias_scale_node = onnx.helper.make_node("Mul", [input_scale_name, weight_scale_name], [bias_name + "_scale"], bias_name + "_scale_node") + new_node_list.append(bias_scale_node) - # Compare type - if quantized_initializer.data_type != weight.qType: - raise ValueError("{} is being used by multiple nodes which are being quantized to different types. " - "Please use different initializers for these nodes.", weight.name) + quantize_bias_node = onnx.helper.make_node("Div", [bias_name, bias_scale_node.output[0]], + [bias_name + "_tmp_quant:0"], bias_name + "_tmp_qaunt") + new_node_list.append(quantize_bias_node) - expected_dims = [] if weight.axis is None else [len(weight.zero_points)] - # Compare quantization axis - if zero_point.dims != expected_dims: - raise ValueError("{} is being used by multiple nodes which are being quantized to different shapes. " - "Please use different initializers for these nodes.", weight.name) + bias_rounded_node = onnx.helper.make_node("Floor", quantize_bias_node.output, + [bias_name + "_quant_rounded:0"], bias_name + "_quant_rounded") + new_node_list.append(bias_rounded_node) + + bias_cast_node = onnx.helper.make_node("Cast", bias_rounded_node.output, + [quantized_bias_name], quantized_bias_name + "_node", to=qType) + new_node_list.append(bias_cast_node) + + return - return True - def _quantize_inputs(self, node, indices, weight_index, new_nodes_list): + def _quantize_bias(self, node, new_node_list): + ''' + Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale + ''' + + # 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.graph.initializer) + weight_scale = self.find_weight_data(weight_initializer) + + # get bias + bias_name = node.input[2] + bias_initializer = _find_by_name(bias_name, self.model.graph.initializer) + bias_data = self.find_weight_data(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 + # so resort to dynamic quantization for bias + if self.quantization_params is None or node.input[0] not in self.quantization_params and node.input[0] not in self.quantized_value_map: + self._dynamic_quantize_bias(node.input[0], weight_scale_name, bias_name, quantized_bias_name, new_node_list) + else: + # get scale for input + input_scale_name = self.quantized_value_map[node.input[0]].scale_name + inputscale_initializer = _find_by_name(input_scale_name, self.model.graph.initializer) + input_scale = self.find_weight_data(inputscale_initializer) + + # calcuate scale for bias + bias_scale_name = node.input[2] + "_scale" + bias_scale = input_scale * weight_scale + print(bias_scale) + + # quantize bias + quantized_data = (np.asarray(bias_data) / bias_scale).round().astype(np.int32) + print(quantized_data) + + #update bias initializer + bias_np_data = np.asarray(quantized_data, dtype=np.int32).reshape(bias_initializer.dims) + packed_bias_initializer = onnx.numpy_helper.from_array(bias_np_data, quantized_bias_name) + self.model.graph.initializer.extend([packed_bias_initializer]) + + bias_value_info = onnx.helper.make_tensor_value_info(quantized_bias_name, onnx_proto.TensorProto.INT32, bias_initializer.dims) + self.model.graph.input.extend([bias_value_info]) + + # log entries for this quantized bias value + quantized_bias_entry = QuantizedInitializer(bias_name, bias_initializer, [0], [0], [0], [bias_scale], + bias_data, quantized_data, qType=onnx_proto.TensorProto.INT32) + self._quantized_weights.append(quantized_bias_entry) + + assert(bias_name not in self.quantized_value_map) + quantized_value = QuantizedValue(bias_name, quantized_bias_name, "", "", QuantizedValueType.Initializer, None, onnx_proto.TensorProto.INT32) + self.quantized_value_map[bias_name] = quantized_value + + return quantized_bias_name + + + def _quantize_inputs(self, node, indices, new_nodes_list): ''' Given a node, this function quantizes the inputs as follows: - If input is a initializer, quantize the initializer data, replace old initializer @@ -723,8 +800,6 @@ class ONNXQuantizer: parameter node: node being quantized in NodeProto format. parameter indices: input indices to quantize. - parameter weight_index: index of weight input. - In Asymmetric mode, this input is quantized into signed integer. 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, @@ -740,56 +815,129 @@ class ONNXQuantizer: nodes = [] for input_index in indices: - qType = self.weight_qType if input_index == weight_index else self.input_qType node_input = node.input[input_index] + + # 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) + continue + + # Quantize the input initializer = _find_by_name(node_input, self.model.graph.initializer) if initializer is not None: - # Quantize the data - if node.op_type == "Conv" and input_index == weight_index: - weight = self._get_quantized_weight_convolution(initializer, qType) + if node.op_type == "Conv": + weight = self._get_quantized_weight_convolution(initializer, self.weight_qType) else: - weight = self._get_quantized_weight(initializer, qType) + weight = self._get_quantized_weight(initializer, self.weight_qType) - if not self._is_quantized(weight): - nodes.extend(self._update_unsupported_nodes_using_weight(weight, new_nodes_list)) - self._update_graph(weight) + # Update graph + nodes.extend(self._update_unsupported_nodes_using_weight(weight, new_nodes_list)) + self._update_graph(weight) quantized_input_names.append(weight.name + "_quantized") zero_point_names.append(weight.name + "_zero_point") scale_names.append(weight.name + "_scale") else: - # Not an initializer input. Add QuantizeLinear node. - # Find if there is already a quantizeLinear node for this input + # Add QuantizeLinear node. qlinear_node = _find_node_by_name(node_input + "_QuantizeLinear", self.model.graph, new_nodes_list) if qlinear_node is None: - quantize_input_nodes = self._get_quantize_input_nodes(node, input_index, qType) + quantize_input_nodes = self._get_quantize_input_nodes(node, input_index, self.input_qType) nodes.extend(quantize_input_nodes) qlinear_node = quantize_input_nodes[-1] - quantized_input_names.extend(qlinear_node.output) - scale_names.append(qlinear_node.input[1]) - zero_point_names.append(qlinear_node.input[2]) + if qlinear_node.op_type == "QuantizeLinear": + quantized_input_names.extend(qlinear_node.output) + scale_names.append(qlinear_node.input[1]) + zero_point_names.append(qlinear_node.input[2]) + else: + quantized_input_names.append(qlinear_node.output[0]) + scale_names.append(qlinear_node.output[1]) + zero_point_names.append(qlinear_node.output[2]) + return (quantized_input_names, zero_point_names, scale_names, nodes) + + def _handle_other_ops(self, node, new_nodes_list): + ''' + Given a node which does not support quantization(Conv, Matmul, Gather), this method + checks whether the input to this node is quantized and adds a DequantizeLinear node + to dequantize this input back to FP32 + + parameter node: Current node + parameter new_nodes_list: List of new nodes created before processing current node + return: List of new nodes created + ''' + nodes = [] + for index, node_input in enumerate(node.input): + if node_input in self.quantized_value_map: + node_input_altered = True + input_name = node.input[index] + quantized_value = self.quantized_value_map[input_name] + # Add DequantizeLinear Node for this input + dqlinear_name = input_name + "_DequantizeLinear" + dqlinear_node = _find_node_by_name(dqlinear_name, self.model.graph, new_nodes_list) + if dqlinear_node is None: + dqlinear_inputs = [quantized_value.q_name, quantized_value.scale_name, quantized_value.zp_name] + dequantize_node = onnx.helper.make_node("DequantizeLinear", dqlinear_inputs, [input_name], dqlinear_name) + nodes.append(dequantize_node) + else: + # DQ op is already present, assert it's output matches the input of current node + assert(input_name == dqlinear_node.output[0]) + + # Append the original node + nodes.append(node) + return nodes + + def _handle_activation_ops(self, node, new_node_list): + ''' + Checks whether the give activation op can be removed from the graph. When mode is QLinearOps, + the output quatization params are calculated based on outputs from activation nodes, + therefore these nodes can be removed from the graph if they follow a quantized op. + + parameter node: Current node + parameter new_nodes_list: List of new nodes created before processing current node + return: List of nodes + ''' + assert(node.op_type == "Relu" or node.op_type == 'Clip') + if self.mode is not QuantizationMode.QLinearOps: + return [node] + # When mode is QLinearOps, the output quatization params are calculated based on outputs from + # activation nodes, therefore these nodes can be removed from the graph if they follow a quantized op. + # If input to this node is not quantized then keep this node + if node.input[0] not in self.quantized_value_map: + return [node] + + # Prepare to remove this node + quantized_value = self.quantized_value_map[node.input[0]] + self.quantized_value_map[node.output[0]] = quantized_value + + return [] def _quantize_gather_ops(self, node, new_nodes_list): assert (node.op_type == "Gather") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self._quantize_inputs(node, [0], 0, new_nodes_list) + self._quantize_inputs(node, [0], new_nodes_list) gather_new_output = node.output[0] + "_quantized" + + # Create an entry for this quantized value + q_output = QuantizedValue(node.output[0], gather_new_output, scale_names[0], zero_point_names[0], QuantizedValueType.Input) + self.quantized_value_map[node.output[0]] = q_output + gather_original_output = node.output[0] node.output[0] = gather_new_output node.input[0] = quantized_input_names[0] nodes.append(node) - # Add DequantizeLinear node. - dqlinear_name = node.output[0] + "_DequantizeLinear" - dqlinear_inputs = [gather_new_output, scale_names[0], zero_point_names[0]] - dqlinear_node = onnx.helper.make_node("DequantizeLinear", dqlinear_inputs, [gather_original_output], dqlinear_name) - print(dqlinear_node.name) - nodes.append(dqlinear_node) - return nodes + return nodes def _quantize_convolution_integer_ops(self, node, new_nodes_list): ''' @@ -801,7 +949,14 @@ class ONNXQuantizer: assert (node.op_type == "Conv") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self._quantize_inputs(node, [0, 1], 1, new_nodes_list) + self._quantize_inputs(node, [0, 1], new_nodes_list) + + # quantize bias if exist + quantized_bias_name = "" + bias_present = False + if len(node.input) == 3: + quantized_bias_name = self._quantize_bias(node, nodes) + bias_present = True conv_integer_output = node.output[0] + "_quantized" conv_integer_name = "" @@ -814,6 +969,10 @@ class ONNXQuantizer: [conv_integer_output], conv_integer_name, **kwargs) nodes.append(conv_integer_node) + # Add bias add nodes + if bias_present: + conv_integer_output = self._get_bias_add_nodes(nodes, node, conv_integer_output, quantized_bias_name) + # Add cast operation to cast convInteger output to float. cast_op_output = conv_integer_output + "_cast_output" cast_node = onnx.helper.make_node("Cast", [conv_integer_output], [cast_op_output], @@ -840,6 +999,7 @@ class ONNXQuantizer: if conv_integer_name != "": output_scale_mul_op = conv_integer_name + "_output_scale_mul" nodes.append(_get_mul_node([cast_op_output, scales_mul_op_output], node.output[0], output_scale_mul_op)) + return nodes def _quantize_matmul_integer_ops(self, node, new_nodes_list): @@ -852,7 +1012,7 @@ class ONNXQuantizer: assert (node.op_type == "MatMul") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self._quantize_inputs(node, [0, 1], 1, new_nodes_list) + self._quantize_inputs(node, [0, 1], new_nodes_list) matmul_integer_output = node.output[0] + "_quantized" matmul_integer_name = "" @@ -901,10 +1061,17 @@ class ONNXQuantizer: assert (node.op_type == "Conv") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self._quantize_inputs(node, [0, 1], 1, new_nodes_list) + self._quantize_inputs(node, [0, 1], new_nodes_list) + + quantized_bias_name = "" + bias_present = False + if len(node.input) == 3: + quantized_bias_name = self._quantize_bias(node, nodes) + bias_present = True + data_found, output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ + self._get_quantization_params(node.output[0]) - output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ - self._get_output_quantization_params(node.output[0]) + assert(data_found) qlinear_conv_output = node.output[0] + "_quantized" qlinear_conv_name = "" @@ -922,19 +1089,22 @@ class ONNXQuantizer: qlinear_conv_inputs.append(quantized_input_names[1]) qlinear_conv_inputs.append(scale_names[1]) qlinear_conv_inputs.append(zero_point_names[1]) + # Output qlinear_conv_inputs.append(output_scale_name) qlinear_conv_inputs.append(output_zp_name) + if bias_present: + qlinear_conv_inputs.append(quantized_bias_name) + qlinear_conv_node = onnx.helper.make_node("QLinearConv", qlinear_conv_inputs, [qlinear_conv_output], qlinear_conv_name, **kwargs) nodes.append(qlinear_conv_node) - # Add DequantizeLinear node. - dqlinear_name = node.output[0] + "_DequantizeLinear" - dqlinear_inputs = [qlinear_conv_output, output_scale_name, output_zp_name] - dqlinear_node = onnx.helper.make_node("DequantizeLinear", dqlinear_inputs, [node.output[0]], dqlinear_name) - nodes.append(dqlinear_node) + # Create an entry for this quantized value + q_output = QuantizedValue(node.output[0], qlinear_conv_output, output_scale_name, output_zp_name, QuantizedValueType.Input) + self.quantized_value_map[node.output[0]] = q_output + return nodes def _quantize_matmul_qlinear_ops(self, node, new_nodes_list): @@ -947,10 +1117,12 @@ class ONNXQuantizer: assert (node.op_type == "MatMul") (quantized_input_names, zero_point_names, scale_names, nodes) = \ - self._quantize_inputs(node, [0, 1], 1, new_nodes_list) + self._quantize_inputs(node, [0, 1], new_nodes_list) - output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ - self._get_output_quantization_params(node.output[0]) + data_found, output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ + self._get_quantization_params(node.output[0]) + + assert(data_found) qlinear_matmul_output = node.output[0] + "_quantized" qlinear_matmul_name = "" @@ -974,11 +1146,10 @@ class ONNXQuantizer: [qlinear_matmul_output], qlinear_matmul_name) nodes.append(qlinear_matmul_node) - # Add DequantizeLinear node. - dqlinear_name = node.output[0] + "_DequantizeLinear" - dqlinear_inputs = [qlinear_matmul_output, output_scale_name, output_zp_name] - dqlinear_node = onnx.helper.make_node("DequantizeLinear", dqlinear_inputs, [node.output[0]], dqlinear_name) - nodes.append(dqlinear_node) + # Create an entry for this quantized value + q_output = QuantizedValue(node.output[0], qlinear_matmul_output, output_scale_name, output_zp_name, QuantizedValueType.Input) + self.quantized_value_map[node.output[0]] = q_output + return nodes def _quantize_convolution(self, node, new_nodes_list): @@ -1015,9 +1186,38 @@ class ONNXQuantizer: return [node] +def check_opset_version(org_model, force_fusions): + ''' + Check opset version of original model and set opset version and fuse_dynamic_quant accordingly. + If opset version < 10, set quantized model opset version to 10. + If opset version == 10, do quantization without using dynamicQuantizeLinear operator. + If opset version == 11, do quantization using dynamicQuantizeLinear operator. + + :return: fuse_dynamic_quant boolean value. + ''' + global onnx_op_set_version + opset_version = org_model.opset_import[0].version + fuse_dynamic_quant = False + + if opset_version < 11 and force_fusions == True: + print("Warning: The original model opset version is {}, which does not support node fusions.\n\ + Forcing fusions can break other nodes in the model.".format(opset_version)) + fuse_dynamic_quant = True + + if opset_version < 10: + print("Warning: 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)) + onnx_op_set_version = 10 + elif opset_version == 10: + onnx_op_set_version = 10 + else: + onnx_op_set_version > 10 + fuse_dynamic_quant = True + return fuse_dynamic_quant def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMode.IntegerOps, - static=False, asymmetric_input_types=False, input_quantization_params=None, output_quantization_params=None, nodes_to_quantize=None): + static=False, force_fusions=False, asymmetric_input_types=False, + quantization_params=None, nodes_to_quantize=None): ''' Given an onnx model, create a quantized onnx model and save it into a file @@ -1031,16 +1231,20 @@ def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMo the function will use QLinear ops. Only QLinearConv and QLinearMatMul ops are supported now. :param static: True: The inputs/activations are quantized using static scale and zero point values - specified through input_quantization_params. + specified through quantization_params. False: The inputs/activations are quantized using dynamic scale and zero point values computed while running the model. + :param force_fusions: + True: Fuses nodes added for dynamic quantization + False: No fusion is applied for nodes which are added for dynamic quantization. + Should be only used in cases where backends want to apply special fusion routines :param asymmetric_input_types: True: Weights are quantized into signed integers and inputs/activations into unsigned integers. False: Weights and inputs/activations are quantized into unsigned integers. - :param input_quantization_params: + :param quantization_params: Dictionary to specify the zero point and scale values for inputs to conv and matmul nodes. Should be specified when static is set to True. - The input_quantization_params should be specified in the following format: + The quantization_params should be specified in the following format: { "input_name": [zero_point, scale] }. @@ -1049,20 +1253,7 @@ def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMo { 'resnet_model/Relu_1:0': [np.uint8(0), np.float32(0.019539741799235344)], 'resnet_model/Relu_2:0': [np.uint8(0), np.float32(0.011359662748873234)] - } - :param output_quantization_params: - Dictionary to specify the zero point and scale values for outputs of conv and matmul nodes. - Should be specified in QuantizationMode.QLinearOps mode. - The output_quantization_params should be specified in the following format: - { - "output_name": [zero_point, scale] - } - zero_point can be of type np.uint8/np.int8 and scale should be of type np.float32. - example: - { - 'resnet_model/Relu_3:0': [np.int8(0), np.float32(0.011359662748873234)], - 'resnet_model/Relu_4:0': [np.uint8(0), np.float32(0.011359662748873234)] - } + } :return: ModelProto with quantization :param nodes_to quantize: List of nodes names to quantize. When this list is not None only the nodes in this list @@ -1079,8 +1270,9 @@ def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMo mode = quantization_mode copy_model = onnx_proto.ModelProto() copy_model.CopyFrom(model) - quantizer = ONNXQuantizer(copy_model, per_channel, mode, static, weight_qType, input_qType, - input_quantization_params, output_quantization_params, nodes_to_quantize) + fuse_dynamic_quant = check_opset_version(copy_model, force_fusions) + quantizer = ONNXQuantizer(copy_model, per_channel, mode, static, fuse_dynamic_quant, weight_qType, input_qType, + quantization_params, nodes_to_quantize) quantizer.quantize_model() quantizer.model.producer_name = __producer__ quantizer.model.producer_version = __version__