diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 9494a2b56f..68c47a8746 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -108,6 +108,7 @@ class QuantizedValue: def quantize_data(data, quantize_range, qType): ''' + :parameter data: data to quantize :parameter quantize_range: list of data to weight pack. :parameter qType: data type to quantize to. Supported types UINT8 and INT8 :return: minimum, maximum, zero point, scale, and quantized weights @@ -141,7 +142,7 @@ def quantize_data(data, quantize_range, qType): else: raise ValueError( "Unexpected data type {} requested. Only INT8 and UINT8 are supported." - ) + .format(qType)) return rmin, rmax, zero_point, scale, quantized_data @@ -668,7 +669,7 @@ class ONNXQuantizer: Zero point and scale values are obtained from self.quantization_params if specified. parameter param_name: Name of the quantization parameter. - return: scale_name, zero_point_name, scale_shape, zero_point_shape. + return: result, scale_name, zero_point_name, scale_shape, zero_point_shape. ''' if self.quantization_params is None or param_name not in self.quantization_params: return False, "", "", "", "" @@ -677,16 +678,16 @@ class ONNXQuantizer: raise ValueError( "Quantization parameters should contain zero point and scale. " "Specified values for output {}: {}".format( - output_name, params)) + param_name, params)) if not np.isscalar(params[0]): raise ValueError( - "Zero point for output {} should be a scalar value. Value specified: {}" - .format(output_name, params[0])) + "Zero point for param {} should be a scalar value. Value specified: {}" + .format(param_name, params[0])) if not np.isscalar(params[1]): raise ValueError( - "Scale for output {} should be a scalar value. Value specified: {}" - .format(output_name, params[1])) + "Scale for param {} should be a scalar value. Value specified: {}" + .format(param_name, params[1])) zero_point_values = [params[0].item()] zero_point_shape = [] @@ -721,7 +722,7 @@ class ONNXQuantizer: input_name = node.input[input_index] output_name = input_name + "_quantized" - data_found, scale_name, zp_name, scale_shape, zp_shape = \ + data_found, scale_name, zp_name, _, _ = \ self._get_quantization_params(input_name) if self.static: @@ -900,14 +901,22 @@ class ONNXQuantizer: new_node_list) else: # get scale for input - input_scale_name = self.quantized_value_map[ - node.input[0]].scale_name + if node.input[0] in self.quantized_value_map: + input_scale_name = self.quantized_value_map[ + node.input[0]].scale_name + elif node.input[0] in self.quantization_params: + _, input_scale_name, _, _, _ = self._get_quantization_params( + node.input[0]) + else: + 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.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) @@ -1251,10 +1260,13 @@ class ONNXQuantizer: 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 = \ + data_found, output_scale_name, output_zp_name, _, _ = \ self._get_quantization_params(node.output[0]) - assert (data_found) + if not data_found: + raise ValueError( + "Quantization parameters for output:\"{}\" of node:\"{}\" not specified" + .format(node.output[0], node.name)) qlinear_conv_output = node.output[0] + "_quantized" qlinear_conv_name = "" @@ -1306,7 +1318,7 @@ class ONNXQuantizer: (quantized_input_names, zero_point_names, scale_names, nodes) = \ self._quantize_inputs(node, [0, 1], new_nodes_list) - data_found, output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ + data_found, output_scale_name, output_zp_name, _, _ = \ self._get_quantization_params(node.output[0]) assert (data_found) @@ -1488,6 +1500,4 @@ def quantize(model, quantizer.model.producer_version = __version__ return quantizer.model else: - raise ValueError( - 'Unknown value for nbits. only 8 bit quantization is currently supported' - ) + raise ValueError('Only 8 bit quantization is currently supported')