diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 432f085845..f077dead20 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -21,6 +21,7 @@ import re import subprocess import json + def augment_graph(model): ''' Adds ReduceMin and ReduceMax nodes to all Conv and MatMul nodes in @@ -38,22 +39,33 @@ def augment_graph(model): input_name = node.output[0] # Adding ReduceMin nodes reduce_min_name = node.name + '_ReduceMin' - reduce_min_node = onnx.helper.make_node('ReduceMin', [input_name], - [input_name + '_ReduceMin'], reduce_min_name, keepdims=0) + reduce_min_node = onnx.helper.make_node( + 'ReduceMin', [input_name], [input_name + '_ReduceMin'], + reduce_min_name, + keepdims=0) added_nodes.append(reduce_min_node) - added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, ())) + added_outputs.append( + helper.make_tensor_value_info(reduce_min_node.output[0], + TensorProto.FLOAT, ())) # Adding ReduceMax nodes reduce_max_name = node.name + '_ReduceMax' - reduce_max_node = onnx.helper.make_node('ReduceMax', [input_name], - [input_name + '_ReduceMax'], reduce_max_name, keepdims=0) + reduce_max_node = onnx.helper.make_node( + 'ReduceMax', [input_name], [input_name + '_ReduceMax'], + reduce_max_name, + keepdims=0) added_nodes.append(reduce_max_node) - added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, ())) + added_outputs.append( + helper.make_tensor_value_info(reduce_max_node.output[0], + TensorProto.FLOAT, ())) model.graph.node.extend(added_nodes) model.graph.output.extend(added_outputs) return model + # Using augmented outputs to generate inputs to quantize.py + + def get_intermediate_outputs(model_path, session, inputs, calib_mode='naive'): ''' Gather intermediate model outputs after running inference @@ -69,33 +81,52 @@ def get_intermediate_outputs(model_path, session, inputs, calib_mode='naive'): return: dictionary mapping added node names to (ReduceMin, ReduceMax) pairs ''' model = onnx.load(model_path) - num_model_outputs = len(model.graph.output) # number of outputs in original model + # number of outputs in original model + num_model_outputs = len(model.graph.output) num_inputs = len(inputs) input_name = session.get_inputs()[0].name - intermediate_outputs = [session.run([], {input_name: inputs[i]}) for i in range(num_inputs)] + intermediate_outputs = [ + session.run([], {input_name: inputs[i]}) for i in range(num_inputs) + ] # Creating dictionary with output results from multiple test inputs - node_output_names = [session.get_outputs()[i].name for i in range(len(intermediate_outputs[0]))] - output_dicts = [dict(zip(node_output_names, intermediate_outputs[i])) for i in range(num_inputs)] + node_output_names = [ + session.get_outputs()[i].name + for i in range(len(intermediate_outputs[0])) + ] + output_dicts = [ + dict(zip(node_output_names, intermediate_outputs[i])) + for i in range(num_inputs) + ] merged_dict = {} for d in output_dicts: for k, v in d.items(): merged_dict.setdefault(k, []).append(v) added_node_output_names = node_output_names[num_model_outputs:] - node_names = [added_node_output_names[i].rpartition('_')[0] for i in range(0, len(added_node_output_names), 2)] # output names + node_names = [ + added_node_output_names[i].rpartition('_')[0] + for i in range(0, len(added_node_output_names), 2) + ] # output names # Characterizing distribution of a node's values across test data sets - clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict if i != list(merged_dict.keys())[0]) + clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict + if i != list(merged_dict.keys())[0]) if calib_mode == 'naive': - pairs = [tuple([float(min(clean_merged_dict[added_node_output_names[i]])), - float(max(clean_merged_dict[added_node_output_names[i+1]]))]) - for i in range(0, len(added_node_output_names), 2)] + pairs = [ + tuple([ + float(min(clean_merged_dict[added_node_output_names[i]])), + float(max(clean_merged_dict[added_node_output_names[i + 1]])) + ]) for i in range(0, len(added_node_output_names), 2) + ] else: - raise ValueError('Unknown value for calib_mode. Currently only naive mode is supported.') + raise ValueError( + 'Unknown value for calib_mode. Currently only naive mode is supported.' + ) final_dict = dict(zip(node_names, pairs)) return final_dict + def calculate_scale_zeropoint(node, next_node, rmin, rmax): zp_and_scale = [] # adjust rmin and rmax such that 0 is included in the range. This is required @@ -117,7 +148,7 @@ def calculate_scale_zeropoint(node, next_node, rmin, rmax): if rmin < 0: rmin = 0 - scale = np.float32((rmax - rmin)/255 if rmin != rmax else 1) + scale = np.float32((rmax - rmin) / 255 if rmin != rmax else 1) initial_zero_point = (0 - rmin) / scale zero_point = np.uint8(round(max(0, min(255, initial_zero_point)))) @@ -125,6 +156,7 @@ def calculate_scale_zeropoint(node, next_node, rmin, rmax): zp_and_scale.append(scale) return zp_and_scale + def calculate_quantization_params(model, quantization_thresholds): ''' Given a model and quantization thresholds, calculates the quantization params. @@ -145,16 +177,20 @@ def calculate_quantization_params(model, quantization_thresholds): { "param_name": [zero_point, scale] } - ''' + ''' if quantization_thresholds is None: - raise ValueError('quantization thresholds is required to calculate quantization params (zero point and scale)') + raise ValueError( + 'quantization thresholds is required to calculate quantization params (zero point and scale)' + ) quantization_params = {} for index, node in enumerate(model.graph.node): node_output_name = node.output[0] if node_output_name in quantization_thresholds: node_thresholds = quantization_thresholds[node_output_name] - node_params = calculate_scale_zeropoint(node, model.graph.node[index+1], node_thresholds[0], node_thresholds[1]) + node_params = calculate_scale_zeropoint( + node, model.graph.node[index + 1], node_thresholds[0], + node_thresholds[1]) quantization_params[node_output_name] = node_params return quantization_params @@ -185,21 +221,38 @@ def load_pb_file(data_file_name, size_limit, samples, channels, height, width): inputs = inputs[:size_limit] dataset_size = size_limit - inputs = inputs.reshape(dataset_size, samples, channels, height, width) + inputs = inputs.reshape(dataset_size, samples, channels, height, + width) except: - sys.exit("Input .pb file contains incorrect input size. \nThe required size is: (%s). The real size is: (%s)" - %((dataset_size, samples, channels, height, width), shape)) + sys.exit( + "Input .pb file contains incorrect input size. \nThe required size is: (%s). The real size is: (%s)" + % ((dataset_size, samples, channels, height, width), shape)) return inputs + def main(): # Parsing command-line arguments - parser = argparse.ArgumentParser(description='parsing model and test data set paths') + parser = argparse.ArgumentParser( + description='parsing model and test data set paths') parser.add_argument('--model_path', required=True) parser.add_argument('--dataset_path', required=True) - parser.add_argument('--output_model_path', type=str, default='calibrated_quantized_model.onnx') - parser.add_argument('--dataset_size', type=int, default=0, help="Number of images or tensors to load. Default is 0 which means all samples") - parser.add_argument('--data_preprocess', type=str, required=True, choices=['preprocess_method1', 'preprocess_method2', 'None'], help="Refer to Readme.md for guidance on choosing this option.") + parser.add_argument('--output_model_path', + type=str, + default='calibrated_quantized_model.onnx') + parser.add_argument( + '--dataset_size', + type=int, + default=0, + help= + "Number of images or tensors to load. Default is 0 which means all samples" + ) + parser.add_argument( + '--data_preprocess', + type=str, + required=True, + choices=['preprocess_method1', 'preprocess_method2', 'None'], + help="Refer to Readme.md for guidance on choosing this option.") args = parser.parse_args() model_path = args.model_path output_model_path = args.output_model_path @@ -219,16 +272,24 @@ def main(): # Generating inputs for quantization if args.data_preprocess == "None": - inputs = load_pb_file(images_folder, args.dataset_size, samples, channels, height, width) + inputs = load_pb_file(images_folder, args.dataset_size, samples, + channels, height, width) else: - inputs = load_batch(images_folder, height, width, size_limit, args.data_preprocess) + inputs = load_batch(images_folder, height, width, size_limit, + args.data_preprocess) print(inputs.shape) - dict_for_quantization = get_intermediate_outputs(model_path, session, inputs, calib_mode) - quantization_params_dict = calculate_quantization_params(model, quantization_thresholds=dict_for_quantization) - calibrated_quantized_model = quantize(onnx.load(model_path), quantization_mode=QuantizationMode.QLinearOps, quantization_params=quantization_params_dict) + dict_for_quantization = get_intermediate_outputs(model_path, session, + inputs, calib_mode) + quantization_params_dict = calculate_quantization_params( + model, quantization_thresholds=dict_for_quantization) + calibrated_quantized_model = quantize( + onnx.load(model_path), + quantization_mode=QuantizationMode.QLinearOps, + quantization_params=quantization_params_dict) onnx.save(calibrated_quantized_model, output_model_path) print("Calibrated, quantized model saved.") + if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/onnxruntime/python/tools/quantization/data_preprocess.py b/onnxruntime/python/tools/quantization/data_preprocess.py index 2bd0187d4a..2803dd9ca6 100644 --- a/onnxruntime/python/tools/quantization/data_preprocess.py +++ b/onnxruntime/python/tools/quantization/data_preprocess.py @@ -10,16 +10,20 @@ import os import sys import numpy as np + def set_preprocess(preprocess_func_name): ''' Set up the data preprocess function name and function dict. parameter preprocess_func_name: name of the preprocess function return: function pointer ''' - funcdict = {'preprocess_method1': preprocess_method1, - 'preprocess_method2': preprocess_method2} + funcdict = { + 'preprocess_method1': preprocess_method1, + 'preprocess_method2': preprocess_method2 + } return funcdict[preprocess_func_name] + def preprocess_method1(image_filepath, height, width): ''' Resizes image to NCHW format. Image is scaled to range [-1, 1]. @@ -30,12 +34,13 @@ def preprocess_method1(image_filepath, height, width): return: matrix characterizing image ''' pillow_img = Image.open(image_filepath).resize((width, height)) - input_data = np.float32(pillow_img)/127.5 - 1.0 # normalization - input_data -= np.mean(input_data) # normalization + input_data = np.float32(pillow_img) / 127.5 - 1.0 # normalization + input_data -= np.mean(input_data) # normalization nhwc_data = np.expand_dims(input_data, axis=0) - nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard + nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard return nchw_data + def preprocess_method2(image_filepath, height, width): ''' Resizes and normalizes image to NCHW format. @@ -46,13 +51,18 @@ def preprocess_method2(image_filepath, height, width): return: matrix characterizing image ''' pillow_img = Image.open(image_filepath).resize((width, height)) - input_data = np.float32(pillow_img) - np.array([123.68, 116.78, 103.94], dtype=np.float32) + input_data = np.float32(pillow_img) - \ + np.array([123.68, 116.78, 103.94], dtype=np.float32) nhwc_data = np.expand_dims(input_data, axis=0) - nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard + nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard return nchw_data -def load_batch(images_folder, height, width, preprocess_func_name, size_limit=0): +def load_batch(images_folder, + height, + width, + preprocess_func_name, + size_limit=0): ''' Loads a batch of images parameter images_folder: path to folder storing images @@ -68,11 +78,13 @@ def load_batch(images_folder, height, width, preprocess_func_name, size_limit=0) else: batch_filenames = image_names unconcatenated_batch_data = [] - + preprocess_func = set_preprocess(preprocess_func_name) for image_name in batch_filenames: image_filepath = images_folder + '/' + image_name nchw_data = preprocess_func(image_filepath, height, width) unconcatenated_batch_data.append(nchw_data) - batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0) - return batch_data \ No newline at end of file + batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, + axis=0), + axis=0) + return batch_data diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 598c21bb97..947545b433 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -37,40 +37,65 @@ type_to_name = { # Quantization mode # IntegerOps: Use IntegerOps in quantized model. Only ConvInteger and MatMulInteger ops are supported now. # QLinearOps: Use QLinearOps in quantized model. Only QLinearConv and QLinearMatMul ops are supported now. + + class QuantizationMode(): IntegerOps = 0 QLinearOps = 1 -quantization_modes = [getattr(QuantizationMode, attr) for attr in dir(QuantizationMode)\ - if not callable(getattr(QuantizationMode, attr)) and not attr.startswith("__")] + +quantization_modes = [ + getattr(QuantizationMode, attr) for attr in dir(QuantizationMode) if + not callable(getattr(QuantizationMode, attr)) and not attr.startswith("__") +] + class QuantizedInitializer: ''' Represents a linearly quantized weight input from ONNX operators ''' - def __init__(self, name, initializer, rmins, rmaxs, zero_points, scales, data=[], quantized_data=[], axis=None, + def __init__(self, + name, + initializer, + rmins, + rmaxs, + zero_points, + scales, + data=[], + quantized_data=[], + axis=None, qType=onnx_proto.TensorProto.UINT8): self.name = name self.initializer = initializer # TensorProto initializer in ONNX graph self.rmins = rmins # List of minimum range for each axis self.rmaxs = rmaxs # List of maximum range for each axis - self.zero_points = zero_points # 1D tensor of zero points computed for each axis. scalar if axis is empty + # 1D tensor of zero points computed for each axis. scalar if axis is empty + self.zero_points = zero_points self.scales = scales # 1D tensor of scales computed for each axis. scalar if axis is empty self.data = data # original data from initializer TensorProto self.quantized_data = quantized_data # weight-packed data from data - self.axis = axis # Scalar to specify which dimension in the initializer to weight pack. - # If empty, single zero point and scales computed from a single rmin and rmax - self.qType = qType # type of quantized data. + # Scalar to specify which dimension in the initializer to weight pack. + self.axis = axis + # 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 + 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, + 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 @@ -80,6 +105,7 @@ class QuantizedValue: self.axis = axis self.qType = qType + def quantize_data(data, quantize_range, qType): ''' :parameter quantize_range: list of data to weight pack. @@ -103,15 +129,19 @@ 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 zero_point = 0 - quantized_data = (np.asarray(data) / scale).round().astype('b') #signed byte type + # signed byte type + quantized_data = (np.asarray(data) / scale).round().astype('b') 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 + 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.") + raise ValueError( + "Unexpected data type {} requested. Only INT8 and UINT8 are supported." + ) return rmin, rmax, zero_point, scale, quantized_data @@ -123,7 +153,8 @@ def _attribute_to_kwarg(attribute): :return: attribute in {key: value} format. ''' if (attribute.type == 0): - raise ValueError('attribute {} does not have type specified.'.format(attribute.name)) + raise ValueError('attribute {} does not have type specified.'.format( + attribute.name)) # Based on attribute type definitions from AttributeProto # definition in https://github.com/onnx/onnx/blob/master/onnx/onnx.proto @@ -148,10 +179,12 @@ def _attribute_to_kwarg(attribute): elif (attribute.type == 10): value = attribute.graphs else: - raise ValueError('attribute {} has unsupported type {}.'.format(attribute.name, attribute.type)) + raise ValueError('attribute {} has unsupported type {}.'.format( + attribute.name, attribute.type)) return {attribute.name: value} + def _find_by_name(item_name, item_list): ''' Helper function to find item by name in a list. @@ -162,6 +195,7 @@ def _find_by_name(item_name, item_list): items = [item for item in item_list if item.name == item_name] return items[0] if len(items) > 0 else None + def _get_mul_node(inputs, output, name): ''' Helper function to create a Mul node. @@ -172,6 +206,7 @@ def _get_mul_node(inputs, output, name): ''' return onnx.helper.make_node("Mul", inputs, [output], name) + def _find_node_by_name(node_name, graph, new_nodes_list): ''' Helper function to check if a node exists in a graph or @@ -181,11 +216,12 @@ def _find_node_by_name(node_name, graph, new_nodes_list): parameter new_nodes_list: list of nodes added during quantization. return: NodeProto if found. None otherwise. ''' - graph_nodes_list = list(graph.node) # deep copy + graph_nodes_list = list(graph.node) # deep copy graph_nodes_list.extend(new_nodes_list) node = _find_by_name(node_name, graph_nodes_list) return node + def _add_initializer_if_not_present(graph, name, value, shape, type): ''' Helper function to add an initializer if it is not present in the graph. @@ -199,6 +235,7 @@ def _add_initializer_if_not_present(graph, name, value, shape, type): initializer = onnx.helper.make_tensor(name, type, shape, value) graph.initializer.extend([initializer]) + def _get_qrange_for_qType(qType): ''' Helper function to get the quantization range for a type. @@ -212,6 +249,7 @@ def _get_qrange_for_qType(qType): else: raise ValueError('unsupported quantization data type') + def _find_nodes_using_initializer(graph, initializer): ''' Helper function to find all nodes with an initializer as a input. @@ -226,21 +264,24 @@ def _find_nodes_using_initializer(graph, initializer): result.append(node) return result + class ONNXQuantizer: - def __init__(self, model, per_channel, mode, static, fuse_dynamic_quant, weight_qType, input_qType, - 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.mode = mode # QuantizationMode.Value - self.static = static # use static quantization for inputs. + 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_qType = input_qType # quantize input type self.weight_qType = weight_qType # quantize data type self.quantization_params = quantization_params - self.nodes_to_quantize = nodes_to_quantize # specific nodes to quantize + self.nodes_to_quantize = nodes_to_quantize # specific nodes to quantize if not self.mode in quantization_modes: - raise ValueError('unsupported quantization mode {}'.format(self.mode)) + raise ValueError('unsupported quantization mode {}'.format( + self.mode)) # QuantizeRange tensor name and zero tensor name for scale and zero point calculation. # Used when static is False @@ -255,17 +296,17 @@ class ONNXQuantizer: 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 a list of ops to be quantized is provided then only quantize those ops if self.nodes_to_quantize is not None and node.name not in self.nodes_to_quantize: - new_list +=self._handle_other_ops(node, new_list) + new_list += self._handle_other_ops(node, new_list) # only onnx domain ops can be quantized today elif node.domain != "ai.onnx" and node.domain != '': - new_list +=self._handle_other_ops(node, new_list) + new_list += self._handle_other_ops(node, new_list) else: if node.op_type == 'Conv': new_list += self._quantize_convolution(node, new_list) @@ -274,9 +315,9 @@ class ONNXQuantizer: 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) + new_list += self._handle_activation_ops(node, new_list) else: - new_list +=self._handle_other_ops(node, new_list) + 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 @@ -286,11 +327,14 @@ class ONNXQuantizer: # Remove weights which are already quantized from graph. self._remove_quantized_weights() - # update opset. - opset_info = next((opset for opset in self.model.opset_import if opset.domain == '' or opset.domain == onnx_domain), None) + # update opset. + opset_info = next( + (opset for opset in self.model.opset_import + if opset.domain == '' or opset.domain == onnx_domain), None) if opset_info is not None: self.model.opset_import.remove(opset_info) - self.model.opset_import.extend([onnx.helper.make_opsetid(onnx_domain, onnx_op_set_version)]) + self.model.opset_import.extend( + [onnx.helper.make_opsetid(onnx_domain, onnx_op_set_version)]) return self.model @@ -302,8 +346,9 @@ class ONNXQuantizer: if initializer.data_type == onnx_proto.TensorProto.FLOAT: weights = onnx.numpy_helper.to_array(initializer) else: - raise ValueError('Model contains conv operator weights in {}. Only float type quantization is supported.'.format( - type_to_name[initializer.data_type])) + raise ValueError( + 'Model contains conv operator weights in {}. Only float type quantization is supported.' + .format(type_to_name[initializer.data_type])) return weights def _remove_quantized_weights(self): @@ -318,12 +363,14 @@ class ONNXQuantizer: # Removing input weight to a convolution try: - weight_input = next(val for val in self.model.graph.input if val.name == weight.name) + weight_input = next(val for val in self.model.graph.input + if val.name == weight.name) self.model.graph.input.remove(weight_input) except StopIteration: if self.model.ir_version < 4: - raise ValueError('invalid weight name {} found in the graph (not a graph input) '.format(weight.name)) - + raise ValueError( + 'invalid weight name {} found in the graph (not a graph input) ' + .format(weight.name)) def _update_graph(self, weight): ''' @@ -333,25 +380,34 @@ class ONNXQuantizer: This function does NOT update the nodes in the graph, just initializers and inputs ''' quantized_value = self.quantized_value_map[weight.name] - assert(quantized_value is not None) + 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, - dtype=onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[weight.qType]).reshape(weight.initializer.dims) - packed_weight_initializer = onnx.numpy_helper.from_array(packed_weight_np_data, packed_weight_name) + packed_weight_np_data = np.asarray( + weight.quantized_data, + dtype=onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[weight.qType]).reshape( + weight.initializer.dims) + packed_weight_initializer = onnx.numpy_helper.from_array( + packed_weight_np_data, packed_weight_name) if weight.axis is not None: zero_scale_shape = [weight.initializer.dims[weight.axis]] - else: # scale and zero point must be scalar + else: # scale and zero point must be scalar zero_scale_shape = [] zero_point_type = weight.qType - scale_initializer = onnx.helper.make_tensor(scale_name, onnx_proto.TensorProto.FLOAT, zero_scale_shape, weight.scales) - zero_initializer = onnx.helper.make_tensor(zero_point_name, zero_point_type, zero_scale_shape, weight.zero_points) + scale_initializer = onnx.helper.make_tensor( + scale_name, onnx_proto.TensorProto.FLOAT, zero_scale_shape, + weight.scales) + zero_initializer = onnx.helper.make_tensor(zero_point_name, + zero_point_type, + zero_scale_shape, + weight.zero_points) - self.model.graph.initializer.extend([packed_weight_initializer, scale_initializer, zero_initializer]) + self.model.graph.initializer.extend( + [packed_weight_initializer, scale_initializer, zero_initializer]) self._quantized_weights.append(weight) @@ -362,14 +418,25 @@ class ONNXQuantizer: :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) - weight = QuantizedInitializer(initializer.name, initializer, [rmin], [rmax], [zero_point], [scale], - weights_data, quantized_weights_data, axis=None, qType=qType) + rmin, rmax, zero_point, scale, quantized_weights_data = quantize_data( + weights_data.flatten().tolist(), _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) + 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 @@ -395,34 +462,47 @@ class ONNXQuantizer: 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() - rmin, rmax, zero_point, scale, quantized_per_channel_data = quantize_data(per_channel_data.flatten().tolist(), + 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), 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) + 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 - quantized_weights = np.asarray(quantized_per_channel_data_list[0]).reshape(reshape_dims) + 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) + channel_weights = np.asarray( + quantized_per_channel_data_list[i]).reshape(reshape_dims) + quantized_weights = np.concatenate( + (quantized_weights, channel_weights), axis=0) + + weight = QuantizedInitializer(initializer.name, initializer, rmin_list, + rmax_list, zero_point_list, scale_list, + weights, + quantized_weights.flatten().tolist(), + channel_index, qType) - 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) + 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): + def _get_dynamic_input_quantization_params(self, input_name, nodes_list, + qType): ''' Create nodes for dynamic quantization of input and add them to nodes_list. parameter input_name: Name of the input. @@ -431,11 +511,14 @@ class ONNXQuantizer: return: scale_name, zero_point_name, scale_shape, zero_point_shape. ''' 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_int8( + input_name, nodes_list) - return self._get_dynamic_input_quantization_params_uint8(input_name, nodes_list) + return self._get_dynamic_input_quantization_params_uint8( + input_name, nodes_list) - def _get_dynamic_input_quantization_params_int8(self, input_name, nodes_list): + def _get_dynamic_input_quantization_params_int8(self, input_name, + nodes_list): ''' Create nodes for dynamic quantization of input to nit8 and add them to nodes_list parameter input_name: Name of the input. @@ -449,45 +532,58 @@ class ONNXQuantizer: reduce_min_name = input_name + "_ReduceMin" reduce_min_node = onnx.helper.make_node("ReduceMin", [input_name], - [reduce_min_name + ":0"], reduce_min_name, keepdims=0) + [reduce_min_name + ":0"], + reduce_min_name, + keepdims=0) nodes_list.append(reduce_min_node) reduce_max_name = input_name + "_ReduceMax" reduce_max_node = onnx.helper.make_node("ReduceMax", [input_name], - [reduce_max_name + ":0"], reduce_max_name, keepdims=0) + [reduce_max_name + ":0"], + reduce_max_name, + keepdims=0) nodes_list.append(reduce_max_node) # Compute scale # Find abs(rmin) reduce_min_abs_name = reduce_min_name + "_Abs" - reduce_min_abs_node = onnx.helper.make_node("Abs", [reduce_min_node.output[0]], - [reduce_min_abs_name + ":0"], reduce_min_abs_name) + reduce_min_abs_node = onnx.helper.make_node( + "Abs", [reduce_min_node.output[0]], [reduce_min_abs_name + ":0"], + reduce_min_abs_name) nodes_list.append(reduce_min_abs_node) # Find abs(rmax) reduce_max_abs_name = reduce_max_name + "_Abs" - reduce_max_abs_node = onnx.helper.make_node("Abs", [reduce_max_node.output[0]], - [reduce_max_abs_name + ":0"], reduce_max_abs_name) + reduce_max_abs_node = onnx.helper.make_node( + "Abs", [reduce_max_node.output[0]], [reduce_max_abs_name + ":0"], + reduce_max_abs_name) nodes_list.append(reduce_max_abs_node) # Compute max of abs(rmin) and abs(rmax) abs_max_name = input_name + "_Abs_Max" - abs_max_node = onnx.helper.make_node("Max", [reduce_min_abs_node.output[0], reduce_max_abs_node.output[0]], + abs_max_node = onnx.helper.make_node( + "Max", + [reduce_min_abs_node.output[0], reduce_max_abs_node.output[0]], [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_int8_name, - [_get_qrange_for_qType(qType)/2.0], [], onnx_proto.TensorProto.FLOAT) + _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_int8_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) # Zero point - _add_initializer_if_not_present(self.model.graph, self.fixed_zero_zp_name, - [0], [], qType) + _add_initializer_if_not_present(self.model.graph, + self.fixed_zero_zp_name, [0], [], + qType) return input_scale_name, self.fixed_zero_zp_name, [], [] - def _get_dynamic_input_quantization_params_uint8(self, input_name, nodes_list): + def _get_dynamic_input_quantization_params_uint8(self, input_name, + nodes_list): ''' Create nodes for dynamic quantization of input to uint8 and add them to nodes_list parameter input_name: Name of the input. @@ -501,52 +597,67 @@ class ONNXQuantizer: reduce_min_name = input_name + "_ReduceMin" reduce_min_node = onnx.helper.make_node("ReduceMin", [input_name], - [reduce_min_name + ":0"], reduce_min_name, keepdims=0) + [reduce_min_name + ":0"], + reduce_min_name, + keepdims=0) nodes_list.append(reduce_min_node) reduce_max_name = input_name + "_ReduceMax" reduce_max_node = onnx.helper.make_node("ReduceMax", [input_name], - [reduce_max_name + ":0"], reduce_max_name, keepdims=0) + [reduce_max_name + ":0"], + reduce_max_name, + keepdims=0) 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_uint8_name, - [_get_qrange_for_qType(qType)], [], onnx_proto.TensorProto.FLOAT) + _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) + [0.0], [], + onnx_proto.TensorProto.FLOAT) # Compute Scale # Subtract rmax and rmin scale_sub_name = input_name + "_scale_Sub" - scale_sub_node = onnx.helper.make_node("Sub", [reduce_max_node.output[0], reduce_min_node.output[0]], + scale_sub_node = onnx.helper.make_node( + "Sub", [reduce_max_node.output[0], reduce_min_node.output[0]], [scale_sub_name + ":0"], scale_sub_name) 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_uint8_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) # Compute zero point # Subtract zero and rmin zp_sub_name = input_name + "_zero_point_Sub" - zp_sub_node = onnx.helper.make_node("Sub", [self.fixed_zero_name, reduce_min_node.output[0]], + zp_sub_node = onnx.helper.make_node( + "Sub", [self.fixed_zero_name, reduce_min_node.output[0]], [zp_sub_name + ":0"], zp_sub_name) nodes_list.append(zp_sub_node) # Divide by scale zp_div_name = input_name + "_zero_point_Div" - zp_div_node = onnx.helper.make_node("Div", [zp_sub_node.output[0], input_scale_name], + zp_div_node = onnx.helper.make_node( + "Div", [zp_sub_node.output[0], input_scale_name], [zp_div_name + ":0"], zp_div_name) nodes_list.append(zp_div_node) # Compute floor zp_floor_name = input_name + "_zero_point_Floor" zp_floor_node = onnx.helper.make_node("Floor", zp_div_node.output, - [zp_floor_name + ":0"], zp_floor_name) + [zp_floor_name + ":0"], + zp_floor_name) nodes_list.append(zp_floor_node) # Cast to integer zp_cast_name = input_name + "_zero_point_Cast" - zp_cast_node = onnx.helper.make_node("Cast", zp_floor_node.output, - [input_zp_name], zp_cast_name, to=qType) + zp_cast_node = onnx.helper.make_node("Cast", + zp_floor_node.output, + [input_zp_name], + zp_cast_name, + to=qType) nodes_list.append(zp_cast_node) return input_scale_name, input_zp_name, [], [] @@ -556,22 +667,26 @@ class ONNXQuantizer: Create initializers and inputs in the graph for zero point and scale of output. Zero point and scale values are obtained from self.quantization_params if specified. - parameter output_name: Name of the output. + parameter param_name: Name of the quantization parameter. return: 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, "", "", "", "" 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)) + raise ValueError( + "Quantization parameters should contain zero point and scale. " + "Specified values for output {}: {}".format( + output_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])) + raise ValueError( + "Zero point for output {} should be a scalar value. Value specified: {}" + .format(output_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])) + raise ValueError( + "Scale for output {} should be a scalar value. Value specified: {}" + .format(output_name, params[1])) zero_point_values = [params[0].item()] zero_point_shape = [] @@ -583,18 +698,20 @@ class ONNXQuantizer: scale_name = param_name + "_scale" # Add initializers - _add_initializer_if_not_present(self.model.graph, zero_point_name, zero_point_values, zero_point_shape, - zero_point_type) - _add_initializer_if_not_present(self.model.graph, scale_name, scale_values, scale_shape, - onnx_proto.TensorProto.FLOAT) + _add_initializer_if_not_present(self.model.graph, zero_point_name, + zero_point_values, zero_point_shape, + zero_point_type) + _add_initializer_if_not_present(self.model.graph, scale_name, + scale_values, scale_shape, + onnx_proto.TensorProto.FLOAT) return True, scale_name, zero_point_name, scale_shape, zero_point_shape def _get_quantize_input_nodes(self, node, input_index, qType): ''' Given a input for a node (which is not a initializer), this function - - add elements to graph to compute zero point and scale for this input. - - add new QuantizeLinear nodes to quantize the input. + - add nodes to compute zero point and scale for this input if they don't exist. + - add new QuantizeLinear node to quantize the input. parameter node: node being quantized in NodeProto format. parameter input_index: index of input in node.input. @@ -605,41 +722,50 @@ class ONNXQuantizer: output_name = input_name + "_quantized" data_found, scale_name, zp_name, scale_shape, zp_shape = \ - self._get_quantization_params(input_name) + 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)) + 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], + 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], + qlinear_node = onnx.helper.make_node( + "QuantizeLinear", [input_name, scale_name, zp_name], [output_name], input_name + "_QuantizeLinear") + return [qlinear_node] 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") + 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], + 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): + 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 @@ -654,22 +780,25 @@ class ONNXQuantizer: # 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) + [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") + 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") + 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): + def _update_unsupported_nodes_using_weight(self, weight, new_nodes_list): '''Find all nodes using a weight that do not support quantization and add a DequantizeLinear node before those nodes. This includes all nodes except Conv, MatMul. @@ -677,8 +806,12 @@ class ONNXQuantizer: parameter new_nodes_list: List of new nodes created before processing current node. return: List of new nodes created. ''' - nodes_using_weight = _find_nodes_using_initializer(self.model.graph, weight.initializer) - unsupported_nodes = [node for node in nodes_using_weight if node.op_type not in ["Conv", "MatMul", "Gather"]] + nodes_using_weight = _find_nodes_using_initializer( + self.model.graph, weight.initializer) + unsupported_nodes = [ + node for node in nodes_using_weight + if node.op_type not in ["Conv", "MatMul", "Gather"] + ] nodes_list = [] dequantize_linear_name = weight.name + "_DequantizeLinear" @@ -686,10 +819,13 @@ class ONNXQuantizer: # Check if DequantizeLinear node needs to be added to graph. if len(unsupported_nodes) != 0 and \ - _find_node_by_name(dequantize_linear_name, self.model.graph, new_nodes_list) is None: - inputs = [weight.name + "_quantized", weight.name + "_scale", weight.name + "_zero_point"] - node = onnx.helper.make_node("DequantizeLinear", inputs, [output_name], - dequantize_linear_name) + _find_node_by_name(dequantize_linear_name, self.model.graph, new_nodes_list) is None: + inputs = [ + weight.name + "_quantized", weight.name + "_scale", + weight.name + "_zero_point" + ] + node = onnx.helper.make_node("DequantizeLinear", inputs, + [output_name], dequantize_linear_name) nodes_list.append(node) # Update unsupported nodes to take dequantized weight as input. @@ -700,7 +836,8 @@ class ONNXQuantizer: return nodes_list - def _dynamic_quantize_bias(self, input_name, weight_scale_name, bias_name, quantized_bias_name, new_node_list): + def _dynamic_quantize_bias(self, input_name, weight_scale_name, bias_name, + quantized_bias_name, new_node_list): ''' Adds series of nodes required to quantize the bias dynamically. parameter input_name: Input name @@ -709,78 +846,102 @@ class ONNXQuantizer: parameter quantied_bias_name: Output name to use for quantized bias. ''' 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") + 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) - quantize_bias_node = onnx.helper.make_node("Div", [bias_name, bias_scale_node.output[0]], + 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) - bias_rounded_node = onnx.helper.make_node("Floor", quantize_bias_node.output, + 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 + 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 def _quantize_bias(self, node, new_node_list): ''' Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale ''' - # get scale for weight + # 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) + 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_initializer = _find_by_name(bias_name, + self.model.graph.initializer) bias_data = self.find_weight_data(bias_initializer) - quantized_bias_name = bias_name + "_quantized" + 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) + 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) + 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) + 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) + # 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]) # 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) + 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) + + 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: @@ -797,7 +958,8 @@ class ONNXQuantizer: List of scale names used for input quantization, List of new QuantizeLinear nodes created) ''' - assert (node.op_type == "Conv" or node.op_type == "MatMul" or node.op_type == "Gather") + assert (node.op_type == "Conv" or node.op_type == "MatMul" + or node.op_type == "Gather") quantized_input_names = [] zero_point_names = [] @@ -811,9 +973,10 @@ class ONNXQuantizer: 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) + 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) @@ -821,15 +984,20 @@ class ONNXQuantizer: continue # Quantize the input - initializer = _find_by_name(node_input, self.model.graph.initializer) + initializer = _find_by_name(node_input, + self.model.graph.initializer) if initializer is not None: if node.op_type == "Conv": - weight = self._get_quantized_weight_convolution(initializer, self.weight_qType) + 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 - nodes.extend(self._update_unsupported_nodes_using_weight(weight, new_nodes_list)) + nodes.extend( + self._update_unsupported_nodes_using_weight( + weight, new_nodes_list)) self._update_graph(weight) quantized_input_names.append(weight.name + "_quantized") @@ -837,9 +1005,12 @@ class ONNXQuantizer: scale_names.append(weight.name + "_scale") else: # Add QuantizeLinear node. - qlinear_node = _find_node_by_name(node_input + "_QuantizeLinear", self.model.graph, new_nodes_list) + 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, self.input_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] @@ -852,9 +1023,8 @@ class ONNXQuantizer: 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 @@ -873,14 +1043,21 @@ class ONNXQuantizer: 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) + 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) + 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]) + assert (input_name == dqlinear_node.output[0]) # Append the original node nodes.append(node) @@ -891,12 +1068,12 @@ class ONNXQuantizer: 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') + 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 @@ -915,11 +1092,13 @@ class ONNXQuantizer: assert (node.op_type == "Gather") (quantized_input_names, zero_point_names, scale_names, nodes) = \ 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) + 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] @@ -955,18 +1134,22 @@ class ONNXQuantizer: kwargs = {} for attribute in node.attribute: kwargs.update(_attribute_to_kwarg(attribute)) - conv_integer_node = onnx.helper.make_node("ConvInteger", quantized_input_names + zero_point_names, + 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) # Add bias add nodes if bias_present: - conv_integer_output = self._get_bias_add_nodes(nodes, node, conv_integer_output, quantized_bias_name) + 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], - conv_integer_output + "_cast", to=onnx_proto.TensorProto.FLOAT) + cast_node = onnx.helper.make_node("Cast", [conv_integer_output], + [cast_op_output], + conv_integer_output + "_cast", + to=onnx_proto.TensorProto.FLOAT) nodes.append(cast_node) # Add mul operation to multiply scales of two inputs. @@ -976,9 +1159,11 @@ class ONNXQuantizer: else: scales_mul_op = scale_names[0] + "_" + scale_names[1] + "_mul" - scales_mul_node = _find_node_by_name(scales_mul_op, self.model.graph, new_nodes_list) + scales_mul_node = _find_node_by_name(scales_mul_op, self.model.graph, + new_nodes_list) 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] @@ -988,7 +1173,9 @@ class ONNXQuantizer: output_scale_mul_op = "" 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)) + nodes.append( + _get_mul_node([cast_op_output, scales_mul_op_output], + node.output[0], output_scale_mul_op)) return nodes @@ -1008,14 +1195,17 @@ class ONNXQuantizer: matmul_integer_name = "" if node.name != "": matmul_integer_name = node.name + "_quant" - matmul_integer_node = onnx.helper.make_node("MatMulInteger", quantized_input_names + zero_point_names, + matmul_integer_node = onnx.helper.make_node( + "MatMulInteger", quantized_input_names + zero_point_names, [matmul_integer_output], matmul_integer_name) nodes.append(matmul_integer_node) # Add cast operation to cast matmulInteger output to float. cast_op_output = matmul_integer_output + "_cast_output" - cast_node = onnx.helper.make_node("Cast", [matmul_integer_output], [cast_op_output], - matmul_integer_output + "_cast", to=onnx_proto.TensorProto.FLOAT) + cast_node = onnx.helper.make_node("Cast", [matmul_integer_output], + [cast_op_output], + matmul_integer_output + "_cast", + to=onnx_proto.TensorProto.FLOAT) nodes.append(cast_node) # Add mul operation to multiply scales of two inputs. @@ -1025,9 +1215,11 @@ class ONNXQuantizer: else: scales_mul_op = scale_names[0] + "_" + scale_names[1] + "_mul" - scales_mul_node = _find_node_by_name(scales_mul_op, self.model.graph, new_nodes_list) + scales_mul_node = _find_node_by_name(scales_mul_op, self.model.graph, + new_nodes_list) 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] @@ -1037,8 +1229,9 @@ class ONNXQuantizer: 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)) return nodes def _quantize_convolution_qlinear_ops(self, node, new_nodes_list): @@ -1052,16 +1245,16 @@ class ONNXQuantizer: (quantized_input_names, zero_point_names, scale_names, nodes) = \ 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 + bias_present = True data_found, output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ self._get_quantization_params(node.output[0]) - assert(data_found) + assert (data_found) qlinear_conv_output = node.output[0] + "_quantized" qlinear_conv_name = "" @@ -1087,14 +1280,18 @@ class ONNXQuantizer: 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) + qlinear_conv_node = onnx.helper.make_node("QLinearConv", + qlinear_conv_inputs, + [qlinear_conv_output], + qlinear_conv_name, **kwargs) nodes.append(qlinear_conv_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) + 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): @@ -1111,8 +1308,8 @@ class ONNXQuantizer: data_found, output_scale_name, output_zp_name, output_scale_shape, output_zp_shape = \ self._get_quantization_params(node.output[0]) - - assert(data_found) + + assert (data_found) qlinear_matmul_output = node.output[0] + "_quantized" qlinear_matmul_name = "" @@ -1132,14 +1329,18 @@ class ONNXQuantizer: qlinear_matmul_inputs.append(output_scale_name) qlinear_matmul_inputs.append(output_zp_name) - qlinear_matmul_node = onnx.helper.make_node("QLinearMatMul", qlinear_matmul_inputs, - [qlinear_matmul_output], qlinear_matmul_name) + qlinear_matmul_node = onnx.helper.make_node("QLinearMatMul", + qlinear_matmul_inputs, + [qlinear_matmul_output], + qlinear_matmul_name) nodes.append(qlinear_matmul_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) + 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): @@ -1166,7 +1367,7 @@ class ONNXQuantizer: :param new_nodes_list: List of new nodes created before processing this node. :return: a list of nodes in topological order that represents quantized MatMul node ''' - assert(node.op_type == 'MatMul') + assert (node.op_type == 'MatMul') if self.mode == QuantizationMode.IntegerOps: return self._quantize_matmul_integer_ops(node, new_nodes_list) @@ -1176,6 +1377,7 @@ 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. @@ -1190,15 +1392,19 @@ def check_opset_version(org_model, force_fusions): 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)) + 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)) onnx_op_set_version = 11 fuse_dynamic_quant = True return fuse_dynamic_quant 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)) + 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 @@ -1206,9 +1412,17 @@ def check_opset_version(org_model, force_fusions): fuse_dynamic_quant = True return fuse_dynamic_quant -def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMode.IntegerOps, - static=False, force_fusions=False, symmetric_activation=False, symmetric_weight=False, - quantization_params=None, nodes_to_quantize=None): + +def quantize(model, + per_channel=False, + nbits=8, + quantization_mode=QuantizationMode.IntegerOps, + static=False, + force_fusions=False, + symmetric_activation=False, + symmetric_weight=False, + quantization_params=None, + nodes_to_quantize=None): ''' Given an onnx model, create a quantized onnx model and save it into a file @@ -1265,11 +1479,15 @@ def quantize(model, per_channel=False, nbits=8, quantization_mode=QuantizationMo copy_model = onnx_proto.ModelProto() copy_model.CopyFrom(model) 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 = 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__ return quantizer.model else: - raise ValueError('Unknown value for nbits. only 8 bit quantization is currently supported') \ No newline at end of file + raise ValueError( + 'Unknown value for nbits. only 8 bit quantization is currently supported' + ) diff --git a/onnxruntime/python/tools/quantization/test_calibrate.py b/onnxruntime/python/tools/quantization/test_calibrate.py index 704642e2ec..84a573b2be 100644 --- a/onnxruntime/python/tools/quantization/test_calibrate.py +++ b/onnxruntime/python/tools/quantization/test_calibrate.py @@ -16,8 +16,8 @@ import calibrate import numpy as np from onnx import numpy_helper -class TestCalibrate(unittest.TestCase): +class TestCalibrate(unittest.TestCase): def test_augment_graph(self): # Creating graph A = helper.make_tensor_value_info('A', TensorProto.FLOAT, [1, 1, 5, 5]) @@ -26,10 +26,15 @@ class TestCalibrate(unittest.TestCase): D = helper.make_tensor_value_info('D', TensorProto.FLOAT, [1, 1, 5, 5]) E = helper.make_tensor_value_info('E', TensorProto.FLOAT, [1, 1, 5, 1]) F = helper.make_tensor_value_info('F', TensorProto.FLOAT, [1, 1, 5, 1]) - conv_node = onnx.helper.make_node('Conv', ['A', 'B'], ['C'], name='Conv', kernel_shape=[3, 3], pads=[1, 1, 1, 1]) + conv_node = onnx.helper.make_node('Conv', ['A', 'B'], ['C'], + name='Conv', + kernel_shape=[3, 3], + pads=[1, 1, 1, 1]) clip_node = onnx.helper.make_node('Clip', ['C'], ['D'], name='Clip') - matmul_node = onnx.helper.make_node('MatMul', ['D', 'E'], ['F'], name='MatMul') - graph = helper.make_graph([conv_node, clip_node, matmul_node], 'test_graph', [A, B, E], [F]) + matmul_node = onnx.helper.make_node('MatMul', ['D', 'E'], ['F'], + name='MatMul') + graph = helper.make_graph([conv_node, clip_node, matmul_node], + 'test_graph', [A, B, E], [F]) model = helper.make_model(graph) onnx.save(model, 'test_model.onnx') @@ -38,12 +43,23 @@ class TestCalibrate(unittest.TestCase): onnx.save(augmented_model, 'augmented_test_model.onnx') # Checking if each added ReduceMin and ReduceMax node and its output exists - augmented_model_node_names = [node.name for node in augmented_model.graph.node] - augmented_model_outputs = [output.name for output in augmented_model.graph.output] - added_node_names = ['Conv_ReduceMin', 'Conv_ReduceMax', 'MatMul_ReduceMin', 'MatMul_ReduceMax'] - added_outputs = ['C_ReduceMin', 'C_ReduceMax', 'F_ReduceMin', 'F_ReduceMax'] - self.assertEqual(len(augmented_model_node_names), 7) # original 3 nodes with 4 added ones - self.assertEqual(len(augmented_model_outputs), 5) # original single graph output with 4 added ones + augmented_model_node_names = [ + node.name for node in augmented_model.graph.node + ] + augmented_model_outputs = [ + output.name for output in augmented_model.graph.output + ] + added_node_names = [ + 'Conv_ReduceMin', 'Conv_ReduceMax', 'MatMul_ReduceMin', + 'MatMul_ReduceMax' + ] + added_outputs = [ + 'C_ReduceMin', 'C_ReduceMax', 'F_ReduceMin', 'F_ReduceMax' + ] + # original 3 nodes with 4 added ones + self.assertEqual(len(augmented_model_node_names), 7) + # original single graph output with 4 added ones + self.assertEqual(len(augmented_model_outputs), 5) for name in added_node_names: self.assertTrue(name in augmented_model_node_names) for output in added_outputs: @@ -55,9 +71,12 @@ class TestCalibrate(unittest.TestCase): I = helper.make_tensor_value_info('I', TensorProto.FLOAT, [1, 1, 5, 1]) J = helper.make_tensor_value_info('J', TensorProto.FLOAT, [1, 1, 5, 1]) K = helper.make_tensor_value_info('K', TensorProto.FLOAT, [1, 1, 5, 1]) - matmul_node_1 = onnx.helper.make_node('MatMul', ['G', 'H'], ['I'], name='MatMul_1') - matmul_node_2 = onnx.helper.make_node('MatMul', ['I', 'J'], ['K'], name='MatMul_2') - graph = helper.make_graph([matmul_node_1, matmul_node_2], 'test_graph_order', [G, H, J], [K]) + matmul_node_1 = onnx.helper.make_node('MatMul', ['G', 'H'], ['I'], + name='MatMul_1') + matmul_node_2 = onnx.helper.make_node('MatMul', ['I', 'J'], ['K'], + name='MatMul_2') + graph = helper.make_graph([matmul_node_1, matmul_node_2], + 'test_graph_order', [G, H, J], [K]) model = helper.make_model(graph) onnx.save(model, 'test_model_matmul.onnx') @@ -65,7 +84,9 @@ class TestCalibrate(unittest.TestCase): augmented_model = calibrate.augment_graph(model) onnx.save(augmented_model, 'augmented_test_model_matmul.onnx') - augmented_model_outputs = [output.name for output in augmented_model.graph.output] + augmented_model_outputs = [ + output.name for output in augmented_model.graph.output + ] self.assertEqual(augmented_model_outputs[0], 'K') self.assertEqual(augmented_model_outputs[1], 'I_ReduceMin') self.assertEqual(augmented_model_outputs[2], 'I_ReduceMax') @@ -76,20 +97,28 @@ class TestCalibrate(unittest.TestCase): images_folder = 'test_images' session = onnxruntime.InferenceSession('augmented_test_model.onnx') (samples, channels, height, width) = session.get_inputs()[0].shape - batch_data = calibrate.load_batch(images_folder, height, width, preprocess_func_name="preprocess_method1") - self.assertEqual(len(batch_data.shape), 5) # for 2D images like the ones in test_images + batch_data = calibrate.load_batch( + images_folder, + height, + width, + preprocess_func_name="preprocess_method1") + # for 2D images like the ones in test_images + self.assertEqual(len(batch_data.shape), 5) self.assertEqual(batch_data.shape[0], len(os.listdir(images_folder))) - self.assertEqual(batch_data.shape[2], 3) # checking for 3 channels for colored image - self.assertEqual(batch_data.shape[3], height) # checking if resized height is correct - self.assertEqual(batch_data.shape[4], width) # checking if resized width is correct - + # checking for 3 channels for colored image + self.assertEqual(batch_data.shape[2], 3) + # checking if resized height is correct + self.assertEqual(batch_data.shape[3], height) + # checking if resized width is correct + self.assertEqual(batch_data.shape[4], width) + def test_load_pb(self): numpy_array = np.random.randn(3, 1, 3, 5, 5).astype(np.float32) tensor = numpy_helper.from_array(numpy_array) test_file_name = 'test_tensor.pb' with open(test_file_name, 'wb') as f: f.write(tensor.SerializeToString()) - + # test size_limit < than number of samples in data set # expecting to load size_limit number of samples batch_data = calibrate.load_pb_file('test_tensor.pb', 2, 1, 3, 5, 5) @@ -97,9 +126,9 @@ class TestCalibrate(unittest.TestCase): self.assertEqual(batch_data.shape[0], 2) self.assertEqual(batch_data.shape[2], 3) self.assertEqual(batch_data.shape[3], 5) - self.assertEqual(batch_data.shape[4], 5) + self.assertEqual(batch_data.shape[4], 5) - # test size_limit == 0 + # test size_limit == 0 # expecting to load all samples batch_data = calibrate.load_pb_file('test_tensor.pb', 0, 1, 3, 5, 5) self.assertEqual(len(batch_data.shape), 5) @@ -120,9 +149,9 @@ class TestCalibrate(unittest.TestCase): try: os.remove('test_tensor.pb') except: - print("Warning: Trying to remove test file {} failed.".format(test_file_name)) - + print("Warning: Trying to remove test file {} failed.".format( + test_file_name)) if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main()