From 67c478ede453fde70fcbd0b3f880aab419afd682 Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Thu, 18 Feb 2021 05:50:59 -0800 Subject: [PATCH] Entropy method for calibration-based quantization (#6619) * Add entropy method * Update pre/post-preprocessing of yolov3 * Code refactor * Code refactor --- .../trt/yolov3/data_reader.py | 17 +- .../trt/yolov3/e2e_user_yolov3_example.py | 5 +- .../object_detection/trt/yolov3/evaluate.py | 262 ++++++++++++++-- .../trt/yolov3/postprocessing.py | 52 ++-- .../trt/yolov3/preprocessing.py | 143 ++++++++- .../python/tools/quantization/__init__.py | 2 +- .../python/tools/quantization/calibrate.py | 284 ++++++++++++++++-- .../python/tools/quantization/quant_utils.py | 26 ++ .../python/tools/quantization/quantize.py | 11 +- 9 files changed, 708 insertions(+), 94 deletions(-) diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/data_reader.py b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/data_reader.py index bba0144e25..7421b64bcb 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/data_reader.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/data_reader.py @@ -1,5 +1,5 @@ from onnxruntime.quantization import CalibrationDataReader -from preprocessing import yolov3_preprocess_func, yolov3_variant_preprocess_func +from preprocessing import yolov3_preprocess_func, yolov3_preprocess_func_2, yolov3_variant_preprocess_func, yolov3_variant_preprocess_func_2, yolov3_variant_preprocess_func_3 import onnxruntime from argparse import Namespace import os @@ -93,8 +93,10 @@ class YoloV3DataReader(ObejctDetectionDataReader): def load_serial(self): width = self.width height = self.width - nchw_data_list, filename_list, image_size_list = yolov3_preprocess_func(self.image_folder, height, width, + nchw_data_list, filename_list, image_size_list = yolov3_preprocess_func_2(self.image_folder, height, width, self.start_index, self.stride) + # nchw_data_list, filename_list, image_size_list = yolov3_preprocess_func(self.image_folder, height, width, + # self.start_index, self.stride) input_name = self.input_name print("Start from index %s ..." % (str(self.start_index))) @@ -179,18 +181,19 @@ class YoloV3VariantDataReader(YoloV3DataReader): annotations='./annotations/instances_val2017.json'): YoloV3DataReader.__init__(self, calibration_image_folder, width, height, start_index, end_index, stride, batch_size, model_path, is_evaluation, annotations) - self.input_name = '000_net' - # self.input_name = 'images' + # self.input_name = '000_net' + self.input_name = 'images' def load_serial(self): width = self.width height = self.height input_name = self.input_name - nchw_data_list, filename_list, image_size_list = yolov3_variant_preprocess_func( + # nchw_data_list, filename_list, image_size_list = yolov3_variant_preprocess_func_2( + # self.image_folder, height, width, self.start_index, self.stride) + nchw_data_list, filename_list, image_size_list = yolov3_variant_preprocess_func_3( self.image_folder, height, width, self.start_index, self.stride) - # nchw_data_list, filename_list, image_size_list = yolov3_variant_2_preprocess_func( - # self.image_folder, height, width, self.start_index, self.stride) + print("Start from index %s ..." % (str(self.start_index))) data = [] if self.is_evaluation: img_name_to_img_id = self.img_name_to_img_id diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/e2e_user_yolov3_example.py b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/e2e_user_yolov3_example.py index 129577cccd..dc607ea54f 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/e2e_user_yolov3_example.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/e2e_user_yolov3_example.py @@ -1,5 +1,5 @@ import os -from onnxruntime.quantization import create_calibrator, write_calibration_table +from onnxruntime.quantization import create_calibrator, write_calibration_table, CalibrationMethod from data_reader import YoloV3DataReader, YoloV3VariantDataReader from evaluate import YoloV3Evaluator, YoloV3VariantEvaluator @@ -64,7 +64,8 @@ def get_prediction_evaluation(model_path, validation_dataset, providers): def get_calibration_table_yolov3_variant(model_path, augmented_model_path, calibration_dataset): - calibrator = create_calibrator(model_path, None, augmented_model_path=augmented_model_path) + calibrator = create_calibrator(model_path, [], augmented_model_path=augmented_model_path, calibrate_method=CalibrationMethod.Entropy) + calibrator.set_execution_providers(["CUDAExecutionProvider"]) # DataReader can handle dataset with batch or serial processing depends on its implementation # Following examples show two different ways to generate calibration table diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/evaluate.py b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/evaluate.py index 93c1180a1c..d647a05b06 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/evaluate.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/evaluate.py @@ -10,6 +10,8 @@ import onnxruntime from onnxruntime.quantization.calibrate import CalibrationDataReader import numpy as np +import torch +import torchvision class YoloV3Evaluator: def __init__(self, @@ -51,6 +53,8 @@ class YoloV3Evaluator: self.generate_class_to_id(ground_truth_object_class_file) print(self.class_to_id) + self.session = onnxruntime.InferenceSession(model_path, providers=providers) + def generate_class_to_id(self, ground_truth_object_class_file): with open(ground_truth_object_class_file) as f: import json @@ -106,7 +110,7 @@ class YoloV3Evaluator: }) def predict(self): - session = onnxruntime.InferenceSession(self.model_path, providers=self.providers) + session = self.session outputs = [] @@ -184,23 +188,20 @@ class YoloV3Evaluator: cocoEval.accumulate() cocoEval.summarize() +class YoloV3VariantEvaluator(YoloV3Evaluator): + def __init__(self, model_path, + data_reader: CalibrationDataReader, + width=608, + height=384, + providers=["CUDAExecutionProvider"], + ground_truth_object_class_file="./coco-object-categories-2017.json", + onnx_object_class_file="./onnx_coco_classes.txt"): -class YoloV3VariantEvaluator(YoloV3Evaluator): - def __init__(self, - model_path, - data_reader: CalibrationDataReader, - width=608, - height=384, - providers=["CUDAExecutionProvider"], - ground_truth_object_class_file="./coco-object-categories-2017.json", - onnx_object_class_file="./onnx_coco_classes.txt"): - - YoloV3Evaluator.__init__(self, model_path, data_reader, width, height, providers, - ground_truth_object_class_file, onnx_object_class_file) + YoloV3Evaluator.__init__(self, model_path, data_reader,width, height, providers, ground_truth_object_class_file, onnx_object_class_file) def predict(self): - from postprocessing import PostprocessYOLOWrapper - session = onnxruntime.InferenceSession(self.model_path, providers=self.providers) + from postprocessing import PostprocessYOLOWrapper + session = self.session outputs = [] image_id_list = [] @@ -224,24 +225,25 @@ class YoloV3VariantEvaluator(YoloV3Evaluator): image_size_list = [image_size_list] image_id_list = [image_id_list] + image_size_batch.append(image_size_list) image_id_batch.append(image_id_list) outputs.append(session.run(None, inputs)) for i in range(len(outputs)): output = outputs[i] - + for batch_i in range(self.data_reader.get_batch_size()): - if batch_i > len(image_size_batch[i]) - 1 or batch_i > len(image_id_batch[i]) - 1: + if batch_i > len(image_size_batch[i])-1 or batch_i > len(image_id_batch[i])-1: continue image_height = image_size_batch[i][batch_i][0] - image_width = image_size_batch[i][batch_i][1] + image_width= image_size_batch[i][batch_i][1] image_id = image_id_batch[i][batch_i] - boxes, classes, scores = postprocess_yolo.postprocessor.process(output, (image_width, image_height), - 0.01) + boxes, classes, scores = postprocess_yolo.postprocessor.process( + output, (image_width, image_height), 0.01) for j in range(len(boxes)): box = boxes[j] @@ -253,13 +255,7 @@ class YoloV3VariantEvaluator(YoloV3Evaluator): y = float(box[1]) w = float(box[2] - box[0] + 1) h = float(box[3] - box[1] + 1) - self.prediction_result_list.append({ - "image_id": int(image_id), - "category_id": int(id), - "bbox": [x, y, w, h], - "score": scores[j] - }) - + self.prediction_result_list.append({"image_id":int(image_id), "category_id":int(id), "bbox":[x,y,w,h], "score":scores[j]}) class YoloV3Variant2Evaluator(YoloV3Evaluator): def __init__(self, @@ -328,7 +324,7 @@ class YoloV3Variant2Evaluator(YoloV3Evaluator): }) def predict(self): - session = onnxruntime.InferenceSession(self.model_path, providers=self.providers) + session = self.session outputs = [] image_id_list = [] @@ -367,3 +363,213 @@ class YoloV3Variant2Evaluator(YoloV3Evaluator): image_width = image_size_batch[i][batch_i][1] image_id = image_id_batch[i][batch_i] self.set_bbox_prediction(bboxes, scores, image_height, image_width, image_id) + +def xywh2xyxy(x): + # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right + y = np.zeros_like(x) + y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x + y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y + y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x + y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y + return y + + +def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): + # Rescale coords (xyxy) from img1_shape to img0_shape + if ratio_pad is None: # calculate from img0_shape + # gain = max(img1_shape) / max(img0_shape) # gain = old / new + gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) + pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding + else: + gain = ratio_pad[0][0] + pad = ratio_pad[1] + + coords[:, [0, 2]] -= pad[0] # x padding + coords[:, [1, 3]] -= pad[1] # y padding + coords[:, :4] /= gain + return coords + + +def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): + # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 + shape = img.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + # Scale ratio (new / old) + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scaleup: # only scale down, do not scale up (for better test mAP) + r = min(r, 1.0) + + # Compute padding + ratio = r, r # width, height ratios + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding + if auto: # minimum rectangle + dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding + elif scaleFill: # stretch + dw, dh = 0.0, 0.0 + new_unpad = new_shape + ratio = new_shape[0] / shape[1], new_shape[1] / shape[0] # width, height ratios + + dw /= 2 # divide padding into 2 sides + dh /= 2 + + if shape[::-1] != new_unpad: # resize + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border + return img, ratio, (dw, dh) + + +def post_process_without_nms(opts): + final_output = [] + for batch_i in range(opt.batch_size): + batch_idx = opts[0][:, 0] == batch_i + bbox = opts[1][batch_idx, :] + score = opts[2][batch_idx, :] + bbox[:, 0] *= opt.input_w #x + bbox[:, 1] *= opt.input_h #y + bbox[:, 2] *= opt.input_w #w + bbox[:, 3] *= opt.input_h #h + bbox = xywh2xyxy(bbox) + bbox0 = scale_coords(img.shape[2:], bbox, img0.shape[0:2]) + if bbox0.shape[0] == 0: + final_output.append(torch.empty(0, 5).numpy()) + continue + + output = np.concatenate((bbox, score), axis=1) + final_output.append(output) + + return final_output + + +def post_process_with_nms(predictions, image_height, image_width, conf_thres=0.35, nms_thres=0.35): + """Performs NMS and score thresholding + """ + final_output = [] + batch_size = 1 + input_w = 512 + input_h = 288 + for batch_i in range(batch_size): + scores = predictions[0][batch_i, :, 0] + keep_idx = scores >= conf_thres + boxes_ = predictions[1][batch_i, keep_idx, :] + boxes_[:, 0] *= input_w #x + boxes_[:, 1] *= input_h #y + boxes_[:, 2] *= input_w #w + boxes_[:, 3] *= input_h #h + boxes_ = xywh2xyxy(boxes_) + img0_shape = (image_height, image_width) + img1_shape = (input_h, input_w) + # bbox = self.scale_coords(img1_shape, bbox, img0_shape) + boxes_ = scale_coords(img1_shape, boxes_, img0_shape) + # boxes_ = scale_coords(img.shape[2:], boxes_, img0.shape[0:2]) + boxes_ = torch.from_numpy(boxes_) + scores = torch.from_numpy(scores[keep_idx]) + if scores.dim() == 0: + final_output.append(torch.empty(0, 5).numpy()) + continue + keep_idx = torchvision.ops.nms(boxes_, scores, nms_thres) + scores = scores[keep_idx].view(-1, 1) + boxes_ = boxes_[keep_idx].view(-1, 4) + output = torch.cat((boxes_, scores), dim=-1) + final_output.append(output.numpy()) + return final_output + +class YoloV3Variant3Evaluator(YoloV3Evaluator): + def __init__(self, + model_path, + data_reader: CalibrationDataReader, + width=512, + height=288, + providers=["CUDAExecutionProvider"], + ground_truth_object_class_file="./coco-object-categories-2017.json", + onnx_object_class_file="./onnx_coco_classes.txt"): + + YoloV3Evaluator.__init__(self, model_path, data_reader, width, height, providers, + ground_truth_object_class_file, onnx_object_class_file) + + + def set_bbox_prediction(self, bboxes, scores, image_height, image_width, image_id): + + for i in range(bboxes.shape[0]): + bbox = bboxes[i] + bbox[0] *= self.width #x + bbox[1] *= self.height #y + bbox[2] *= self.width #w + bbox[3] *= self.height #h + + img0_shape = (image_height, image_width) + img1_shape = (self.height, self.width) + bbox = self.xywh2xyxy(bbox) + bbox = self.scale_coords(img1_shape, bbox, img0_shape) + + class_name = 'person' + if class_name in self.identical_class_map: + class_name = self.identical_class_map[class_name] + id = self.class_to_id[class_name] + + bbox[2] = bbox[2] - bbox[0] + bbox[3] = bbox[3] - bbox[1] + + self.prediction_result_list.append({ + "image_id": int(image_id), + "category_id": int(id), + "bbox": list(bbox), + "score": scores[i][0] + }) + + def predict(self): + session = onnxruntime.InferenceSession(self.model_path, providers=self.providers) + outputs = [] + + image_id_list = [] + image_id_batch = [] + image_size_list = [] + image_size_batch = [] + + class_name = 'person' + id = self.class_to_id[class_name] + + while True: + inputs = self.data_reader.get_next() + if not inputs: + break + image_size_list = inputs["image_size"] + image_id_list = inputs["image_id"] + del inputs["image_size"] + del inputs["image_id"] + + # in the case of batch size is 1 + if type(image_id_list) == int: + image_size_list = [image_size_list] + image_id_list = [image_id_list] + + image_size_batch.append(image_size_list) + image_id_batch.append(image_id_list) + outputs.append(session.run(None, inputs)) + + for j in range(len(outputs)): + output = outputs[j] + image_id = image_id_batch[j][0] + image_height = image_size_batch[j][0][0] + image_width = image_size_batch[j][0][1] + dets = post_process_with_nms(output, image_height, image_width)[0] + + for i in range(dets.shape[0]): + x1 = dets[i, 0] + y1 = dets[i, 1] + x2 = dets[i, 2] + y2 = dets[i, 3] + score = dets[i, 4] + + bbox = [x1, y1, x2-x1, y2-y1] + self.prediction_result_list.append({ + "image_id": int(image_id), + "category_id": int(id), + "bbox": list(bbox), + "score": score + }) + diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/postprocessing.py b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/postprocessing.py index 52daab0bb5..7b4a84380d 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/postprocessing.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/postprocessing.py @@ -1,9 +1,13 @@ import numpy as np - - class PostprocessYOLO(object): """Class for post-processing the three output tensors from YOLO.""" - def __init__(self, yolo_masks, yolo_anchors, nms_threshold, yolo_input_resolution, category_num=80): + + def __init__(self, + yolo_masks, + yolo_anchors, + nms_threshold, + yolo_input_resolution, + category_num=80): """Initialize with all values that will be kept when processing several frames. Assuming 3 outputs of the network in the case of (large) YOLO, or 2 for the Tiny YOLO. @@ -37,7 +41,8 @@ class PostprocessYOLO(object): for output in outputs: outputs_reshaped.append(self._reshape_output(output)) - boxes_xywh, categories, confidences = self._process_yolo_output(outputs_reshaped, resolution_raw, conf_th) + boxes_xywh, categories, confidences = self._process_yolo_output( + outputs_reshaped, resolution_raw, conf_th) if len(boxes_xywh) > 0: # convert (x, y, width, height) to (x1, y1, x2, y2) @@ -46,9 +51,9 @@ class PostprocessYOLO(object): yy = boxes_xywh[:, 1].reshape(-1, 1) ww = boxes_xywh[:, 2].reshape(-1, 1) hh = boxes_xywh[:, 3].reshape(-1, 1) - boxes = np.concatenate([xx, yy, xx + ww, yy + hh], axis=1) + 0.5 - boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0., float(img_w - 1)) - boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0., float(img_h - 1)) + boxes = np.concatenate([xx, yy, xx+ww, yy+hh], axis=1) + 0.5 + boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0., float(img_w-1)) + boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0., float(img_h-1)) boxes = boxes.astype(np.int) else: boxes = np.zeros((0, 4), dtype=np.int) # empty @@ -118,8 +123,9 @@ class PostprocessYOLO(object): nscores.append(confidence[keep]) if not nms_categories and not nscores: - return (np.empty((0, 4), dtype=np.float32), np.empty((0, 1), - dtype=np.float32), np.empty((0, 1), dtype=np.float32)) + return (np.empty((0, 4), dtype=np.float32), + np.empty((0, 1), dtype=np.float32), + np.empty((0, 1), dtype=np.float32)) boxes = np.concatenate(nms_boxes) categories = np.concatenate(nms_categories) @@ -136,6 +142,7 @@ class PostprocessYOLO(object): output_reshaped -- reshaped YOLO output as NumPy arrays with shape (height,width,3,85) mask -- 2-dimensional tuple with mask specification for this output """ + def sigmoid_v(array): return np.reciprocal(np.exp(-array) + 1.0) @@ -238,24 +245,27 @@ class PostprocessYOLO(object): keep = np.array(keep) return keep - class PostprocessYOLOWrapper(object): """This class encapsulates things needed to run yolo.""" """Reference from here https://github.com/jkjung-avt/tensorrt_demos/blob/3fb15c908b155d5edc1bf098c6b8c31886cd8e8d/utils/yolo.py""" + def _init_yolov3_postprocessor(self): h, w = self.input_shape filters = (self.category_num + 5) * 3 if 'tiny' in self.model: - self.output_shapes = [(1, filters, h // 32, w // 32), (1, filters, h // 16, w // 16)] + self.output_shapes = [(1, filters, h // 32, w // 32), + (1, filters, h // 16, w // 16)] else: - self.output_shapes = [(1, filters, h // 32, w // 32), (1, filters, h // 16, w // 16), - (1, filters, h // 8, w // 8)] + self.output_shapes = [(1, filters, h // 32, w // 32), + (1, filters, h // 16, w // 16), + (1, filters, h // 8, w // 8)] if 'tiny' in self.model: postprocessor_args = { # A list of 2 three-dimensional tuples for the Tiny YOLO masks 'yolo_masks': [(3, 4, 5), (0, 1, 2)], # A list of 6 two-dimensional tuples for the Tiny YOLO anchors - 'yolo_anchors': [(10, 14), (23, 27), (37, 58), (81, 82), (135, 169), (344, 319)], + 'yolo_anchors': [(10, 14), (23, 27), (37, 58), + (81, 82), (135, 169), (344, 319)], # Threshold for non-max suppression algorithm, float # value between 0 and 1 'nms_threshold': 0.5, @@ -267,16 +277,14 @@ class PostprocessYOLOWrapper(object): # A list of 3 three-dimensional tuples for the YOLO masks 'yolo_masks': [(6, 7, 8), (3, 4, 5), (0, 1, 2)], # A list of 9 two-dimensional tuples for the YOLO anchors - 'yolo_anchors': [(10, 13), (16, 30), (33, 23), (30, 61), (62, 45), (59, 119), (116, 90), (156, 198), - (373, 326)], + 'yolo_anchors': [(10, 13), (16, 30), (33, 23), + (30, 61), (62, 45), (59, 119), + (116, 90), (156, 198), (373, 326)], # Threshold for non-max suppression algorithm, float # value between 0 and 1 - 'nms_threshold': - 0.5, - 'yolo_input_resolution': - self.input_shape, - 'category_num': - self.category_num + 'nms_threshold': 0.5, + 'yolo_input_resolution': self.input_shape, + 'category_num': self.category_num } self.postprocessor = PostprocessYOLO(**postprocessor_args) diff --git a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/preprocessing.py b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/preprocessing.py index 4ed9ba709f..dc62470a5c 100644 --- a/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/preprocessing.py +++ b/onnxruntime/python/tools/quantization/E2E_example_model/object_detection/trt/yolov3/preprocessing.py @@ -6,7 +6,6 @@ from PIL import Image import cv2 import pdb - def yolov3_preprocess_func(images_folder, height, width, start_index=0, size_limit=0): ''' Loads a batch of images and preprocess them @@ -16,20 +15,19 @@ def yolov3_preprocess_func(images_folder, height, width, start_index=0, size_lim parameter size_limit: number of images to load. Default is 0 which means all images are picked. return: list of matrices characterizing multiple images ''' - # this function is from yolo3.utils.letterbox_image # https://github.com/qqwweee/keras-yolo3/blob/master/yolo3/utils.py def letterbox_image(image, size): '''resize image with unchanged aspect ratio using padding''' iw, ih = image.size w, h = size - scale = min(w / iw, h / ih) - nw = int(iw * scale) - nh = int(ih * scale) + scale = min(w/iw, h/ih) + nw = int(iw*scale) + nh = int(ih*scale) - image = image.resize((nw, nh), Image.BICUBIC) - new_image = Image.new('RGB', size, (128, 128, 128)) - new_image.paste(image, ((w - nw) // 2, (h - nh) // 2)) + image = image.resize((nw,nh), Image.BICUBIC) + new_image = Image.new('RGB', size, (128,128,128)) + new_image.paste(image, ((w-nw)//2, (h-nh)//2)) return new_image image_names = os.listdir(images_folder) @@ -44,6 +42,66 @@ def yolov3_preprocess_func(images_folder, height, width, start_index=0, size_lim else: batch_filenames = image_names + + unconcatenated_batch_data = [] + image_size_list = [] + + print(batch_filenames) + print("size: %s" % str(len(batch_filenames))) + + for image_name in batch_filenames: + image_filepath = images_folder + '/' + image_name + img = Image.open(image_filepath) + model_image_size = (height, width) + boxed_image = letterbox_image(img, tuple(reversed(model_image_size))) + image_data = np.array(boxed_image, dtype='float32') + image_data /= 255. + image_data = np.transpose(image_data, [2, 0, 1]) + image_data = np.expand_dims(image_data, 0) + unconcatenated_batch_data.append(image_data) + image_size_list.append(np.array([img.size[1], img.size[0]], dtype=np.float32).reshape(1, 2)) + + batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0) + return batch_data, batch_filenames, image_size_list + +def yolov3_preprocess_func_2(images_folder, height, width, start_index=0, size_limit=0): + ''' + Loads a batch of images and preprocess them + parameter images_folder: path to folder storing images + parameter height: image height in pixels + parameter width: image width in pixels + parameter size_limit: number of images to load. Default is 0 which means all images are picked. + return: list of matrices characterizing multiple images + ''' + + # reference from here: + # https://github.com/jkjung-avt/tensorrt_demos/blob/3fb15c908b155d5edc1bf098c6b8c31886cd8e8d/utils/yolo.py#L60 + def _preprocess_yolo(img, input_shape): + """Preprocess an image before TRT YOLO inferencing. + # Args + img: int8 numpy array of shape (img_h, img_w, 3) + input_shape: a tuple of (H, W) + # Returns + preprocessed img: float32 numpy array of shape (3, H, W) + """ + img = cv2.resize(img, (input_shape[1], input_shape[0])) + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + img = img.transpose((2, 0, 1)).astype(np.float32) + img /= 255.0 + return img + + image_names = os.listdir(images_folder) + if start_index >= len(image_names): + return np.asanyarray([]), np.asanyarray([]), np.asanyarray([]) + elif size_limit > 0 and len(image_names) >= size_limit: + end_index = start_index + size_limit + if end_index > len(image_names): + end_index = len(image_names) + + batch_filenames = [image_names[i] for i in range(start_index, end_index)] + else: + batch_filenames = image_names + unconcatenated_batch_data = [] image_size_list = [] @@ -52,7 +110,66 @@ def yolov3_preprocess_func(images_folder, height, width, start_index=0, size_lim for image_name in batch_filenames: image_filepath = images_folder + '/' + image_name - img = Image.open(image_filepath) + model_image_size = (height, width) + + img = cv2.imread(image_filepath) + image_data = _preprocess_yolo(img, tuple(model_image_size)) + image_data = np.ascontiguousarray(image_data) + image_data = np.expand_dims(image_data, 0) + unconcatenated_batch_data.append(image_data) + _height, _width, _ = img.shape + # image_size_list.append(img.shape[0:2]) # img.shape is h, w, c + image_size_list.append(np.array([img.shape[0], img.shape[1]], dtype=np.float32).reshape(1, 2)) + + batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0) + return batch_data, batch_filenames, image_size_list + +def yolov3_variant_preprocess_func(images_folder, height, width, start_index=0, size_limit=0): + ''' + Loads a batch of images and preprocess them + parameter images_folder: path to folder storing images + parameter height: image height in pixels + parameter width: image width in pixels + parameter size_limit: number of images to load. Default is 0 which means all images are picked. + return: list of matrices characterizing multiple images + ''' + # this function is from yolo3.utils.letterbox_image + # https://github.com/qqwweee/keras-yolo3/blob/master/yolo3/utils.py + def letterbox_image(image, size): + '''resize image with unchanged aspect ratio using padding''' + iw, ih = image.size + w, h = size + scale = min(w/iw, h/ih) + nw = int(iw*scale) + nh = int(ih*scale) + + image = image.resize((nw,nh), Image.BICUBIC) + new_image = Image.new('RGB', size, (128,128,128)) + new_image.paste(image, ((w-nw)//2, (h-nh)//2)) + return new_image + + image_names = os.listdir(images_folder) + if start_index >= len(image_names): + return np.asanyarray([]), np.asanyarray([]), np.asanyarray([]) + elif size_limit > 0 and len(image_names) >= size_limit: + end_index = start_index + size_limit + if end_index > len(image_names): + end_index = len(image_names) + + batch_filenames = [image_names[i] for i in range(start_index, end_index)] + else: + batch_filenames = image_names + + + unconcatenated_batch_data = [] + image_size_list = [] + + print(batch_filenames) + print("size: %s" % str(len(batch_filenames))) + + for image_name in batch_filenames: + image_filepath = images_folder + '/' + image_name + img = Image.open(image_filepath) model_image_size = (height, width) boxed_image = letterbox_image(img, tuple(reversed(model_image_size))) image_data = np.array(boxed_image, dtype='float32') @@ -60,13 +177,13 @@ def yolov3_preprocess_func(images_folder, height, width, start_index=0, size_lim image_data = np.transpose(image_data, [2, 0, 1]) image_data = np.expand_dims(image_data, 0) unconcatenated_batch_data.append(image_data) - image_size_list.append(np.array([img.size[1], img.size[0]], dtype=np.float32).reshape(1, 2)) + image_size_list.append((img.size[1], img.size[0])) # img.shape is h, w, c + # image_size_list.append(np.array([img.size[1], img.size[0]], dtype=np.float32).reshape(1, 2)) batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0) return batch_data, batch_filenames, image_size_list - -def yolov3_variant_preprocess_func(images_folder, height, width, start_index=0, size_limit=0): +def yolov3_variant_preprocess_func_2(images_folder, height, width, start_index=0, size_limit=0): ''' Loads a batch of images and preprocess them parameter images_folder: path to folder storing images @@ -127,7 +244,7 @@ def yolov3_variant_preprocess_func(images_folder, height, width, start_index=0, # This is for special tuned yolov3 model -def yolov3_variant_2_preprocess_func(images_folder, height, width, start_index=0, size_limit=0): +def yolov3_variant_preprocess_func_3(images_folder, height, width, start_index=0, size_limit=0): def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True): # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 shape = img.shape[:2] # current shape [height, width] diff --git a/onnxruntime/python/tools/quantization/__init__.py b/onnxruntime/python/tools/quantization/__init__.py index c08331ca81..f16c20ce8b 100644 --- a/onnxruntime/python/tools/quantization/__init__.py +++ b/onnxruntime/python/tools/quantization/__init__.py @@ -1,4 +1,4 @@ from .quantize import quantize, quantize_static, quantize_dynamic, quantize_qat from .quantize import QuantizationMode -from .calibrate import CalibrationDataReader, CalibraterBase, MinMaxCalibrater, create_calibrator +from .calibrate import CalibrationDataReader, CalibraterBase, MinMaxCalibrater, create_calibrator, CalibrationMethod from .quant_utils import QuantType, QuantFormat, write_calibration_table diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index c07e99738c..2b8c9693d5 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -15,7 +15,7 @@ from onnx import onnx_pb as onnx_proto from six import string_types from enum import Enum -from .quant_utils import QuantType +from .quant_utils import QuantType, smooth_distribution from .registry import QLinearOpsRegistry import abc @@ -24,6 +24,7 @@ import itertools class CalibrationMethod(Enum): MinMax = 0 + Entropy = 1 class CalibrationDataReader(metaclass=abc.ABCMeta): @@ -80,6 +81,33 @@ class CalibraterBase: sess_options=sess_options, providers=self.execution_providers) + def select_tensors_to_calibrate(self, model): + ''' + select all quantization_candidates op type nodes' input/output tensors. + returns: + tensors (set): set of tensor name. + value_infos (dict): tensor name to value info. + ''' + value_infos = {vi.name: vi for vi in model.graph.value_info} + value_infos.update({ot.name: ot for ot in model.graph.output}) + value_infos.update({it.name: it for it in model.graph.input}) + initializer = set(init.name for init in model.graph.initializer) + + tensors_to_calibrate = set() + tensor_type_to_calibrate = set([TensorProto.FLOAT, TensorProto.FLOAT16]) + + for node in model.graph.node: + if len(self.op_types_to_calibrate) == 0 or node.op_type in self.op_types_to_calibrate: + for tensor_name in itertools.chain(node.input, node.output): + if tensor_name in value_infos.keys(): + vi = value_infos[tensor_name] + if vi.type.HasField('tensor_type') and ( + vi.type.tensor_type.elem_type in tensor_type_to_calibrate) and ( + tensor_name not in initializer): + tensors_to_calibrate.add(tensor_name) + + return tensors_to_calibrate, value_infos + def get_augment_model(self): ''' return: augmented onnx model @@ -129,27 +157,12 @@ class MinMaxCalibrater(CalibraterBase): model = onnx_proto.ModelProto() model.CopyFrom(self.model) model = onnx.shape_inference.infer_shapes(model) - value_infos = {vi.name: vi for vi in model.graph.value_info} - value_infos.update({ot.name: ot for ot in model.graph.output}) - value_infos.update({it.name: it for it in model.graph.input}) - initializer = set(init.name for init in model.graph.initializer) added_nodes = [] added_outputs = [] - tensors_to_calibrate = set() - tensor_type_to_calibrate = set([TensorProto.FLOAT, TensorProto.FLOAT16]) + tensors, _ = self.select_tensors_to_calibrate(model) - for node in model.graph.node: - if len(self.op_types_to_calibrate) == 0 or node.op_type in self.op_types_to_calibrate: - for tensor_name in itertools.chain(node.input, node.output): - if tensor_name in value_infos.keys(): - vi = value_infos[tensor_name] - if vi.type.HasField('tensor_type') and ( - vi.type.tensor_type.elem_type in tensor_type_to_calibrate) and ( - tensor_name not in initializer): - tensors_to_calibrate.add(tensor_name) - - for tensor in tensors_to_calibrate: + for tensor in tensors: # Adding ReduceMin nodes reduce_min_name = tensor + '_ReduceMin' reduce_min_node = onnx.helper.make_node('ReduceMin', [tensor], [tensor + '_ReduceMin'], @@ -226,6 +239,239 @@ class MinMaxCalibrater(CalibraterBase): return self.calibrate_tensors_range +class EntropyCalibrater(CalibraterBase): + def __init__(self, model, op_types_to_calibrate=[], augmented_model_path='augmented_model.onnx'): + ''' + :param model: ONNX model to calibrate. It can be a ModelProto or a model path + :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. + :param augmented_model_path: save augmented model to this path. + ''' + super(EntropyCalibrater, self).__init__(model, op_types_to_calibrate, augmented_model_path) + self.intermediate_outputs = [] + self.calibrate_tensors_range = None + self.num_model_outputs = len(self.model.graph.output) + self.model_original_outputs = set(output.name for output in self.model.graph.output) + self.collector = None + + def augment_graph(self): + ''' + make all quantization_candidates op type nodes as part of the graph output. + :return: augmented ONNX model + ''' + model = onnx_proto.ModelProto() + model.CopyFrom(self.model) + model = onnx.shape_inference.infer_shapes(model) + + added_nodes = [] + added_outputs = [] + tensors, value_infos = self.select_tensors_to_calibrate(model) + + for tensor in tensors: + added_outputs.append(value_infos[tensor]) + + model.graph.node.extend(added_nodes) + model.graph.output.extend(added_outputs) + onnx.save(model, self.augmented_model_path) + self.augment_model = model + + def clear_collected_data(self): + self.intermediate_outputs = [] + + def collect_data(self, data_reader: CalibrationDataReader): + ''' + Entropy Calibrator collects operators' tensors as well as generates tensor histogram for each operator. + ''' + while True: + inputs = data_reader.get_next() + if not inputs: + break + self.intermediate_outputs.append(self.infer_session.run(None, inputs)) + + + if len(self.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [self.infer_session.get_outputs()[i].name for i in range(len(self.intermediate_outputs[0]))] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) for intermediate_output in self.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict if i not in self.model_original_outputs) + + if not self.collector: + self.collector = HistogramCollector() + self.collector.collect(clean_merged_dict) + + def compute_range(self): + ''' + Compute the min-max range of tensor + :return: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs } + ''' + if not self.collector: + raise ValueError("No collector created and can't generate calibration data.") + + return self.collector.get_optimal_collection_result() + + +class CalibrationDataCollector(metaclass=abc.ABCMeta): + """ + Base class for collecting data for calibration-based quantization. + """ + + @abc.abstractmethod + def collect(self, name_to_arr): + """ + Generate informative data based on given data. + name_to_arr : dict + tensor name to NDArray data + """ + raise NotImplementedError + + @abc.abstractmethod + def get_optimal_collection_result(self): + """ + Get the optimal result among collection data. + """ + raise NotImplementedError + +class HistogramCollector(CalibrationDataCollector): + """ + Implementation of collecting histogram data as dict for each tensor targeting on entropy calibration. + + ref: https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py + """ + def __init__(self, num_quantized_bins=128): + self.histogram_dict = {} + self.num_quantized_bins= num_quantized_bins + + def get_histogram_dict(self): + return self.histogram_dict + + def collect(self, name_to_arr): + for tensor, data_arr in name_to_arr.items(): + data_arr = np.asarray(data_arr) + data_arr = data_arr.flatten() + + if data_arr.size > 0: + min_value = np.min(data_arr) + max_value = np.max(data_arr) + else: + min_value = 0 + max_value = 0 + + threshold = max(abs(min_value), abs(max_value)) + + if tensor in self.histogram_dict: + old_histogram = self.histogram_dict[tensor] + self.histogram_dict[tensor] = self.merge_histogram(old_histogram, data_arr, min_value, max_value, threshold) + else: + # hist, hist_edges = np.histogram(data_arr, self.num_quantized_bins, range=(min_value, max_value)) + hist, hist_edges = np.histogram(data_arr, self.num_quantized_bins, range=(-threshold, threshold)) + self.histogram_dict[tensor] = (hist, hist_edges, min_value, max_value, threshold) + + def merge_histogram(self, old_histogram, data_arr, new_min, new_max, new_threshold): + + (old_hist, old_hist_edges, old_min, old_max, old_threshold) = old_histogram + + if new_threshold <= old_threshold: + new_hist, _ = np.histogram(data_arr, len(old_hist), range=(-old_threshold, old_threshold)) + return (new_hist + old_hist, old_hist_edges, min(old_min, new_min), max(old_max, new_max), old_threshold) + else: + if old_threshold == 0: + hist, hist_edges = np.histogram(data_arr, new_num_bins, range=(-new_threshold, new_threshold)) + hist[len(hist) // 2] += len(old_hist) + else: + old_num_bins = len(old_hist) + old_stride = 2 * old_threshold / old_num_bins + half_increased_bins = int((new_threshold - old_threshold) // old_stride + 1) + new_num_bins = old_num_bins + 2 * half_increased_bins + new_threshold = half_increased_bins * old_stride + old_threshold + hist, hist_edges = np.histogram(data_arr, new_num_bins, range=(-new_threshold, new_threshold)) + hist[half_increased_bins:new_num_bins-half_increased_bins] += old_hist + return (hist, hist_edges, min(old_min, new_min), max(old_max, new_max), new_threshold) + + def get_optimal_collection_result(self): + histogram_dict = self.histogram_dict + num_quantized_bins = self.num_quantized_bins + + thresholds_dict = {} # per tensor thresholds + + for tensor, histogram in histogram_dict.items(): + optimal_threshold = self.get_optimal_threshold(histogram, num_quantized_bins) + thresholds_dict[tensor] = optimal_threshold + + return thresholds_dict + + def get_optimal_threshold(self, histogram, num_quantized_bins): + from scipy.stats import entropy + import copy + + hist, hist_edges, _, _, _ = histogram + num_bins = hist.size + zero_bin_index = num_bins // 2 + num_half_quantized_bin = num_quantized_bins // 2 + + kl_divergence = np.zeros(zero_bin_index - num_half_quantized_bin + 1) + thresholds = [(0, 0) for i in range(kl_divergence.size)] + + for i in range(num_half_quantized_bin, zero_bin_index + 1, 1): + start_index = zero_bin_index - i + end_index = zero_bin_index + i + 1 if (zero_bin_index + i + 1) <= num_bins else num_bins + + thresholds[i - num_half_quantized_bin] = (float(hist_edges[start_index]), float(hist_edges[end_index])) + + sliced_distribution = copy.deepcopy(hist[start_index:end_index]) + + # reference distribution p + p = sliced_distribution.copy() # a copy of np array + left_outliers_count = sum(hist[:start_index]) + right_outliers_count = sum(hist[end_index:]) + p[0] += left_outliers_count + p[-1] += right_outliers_count + + # nonzeros[i] incidates whether p[i] is non-zero + nonzeros = (p != 0).astype(np.int64) + + # quantize p.size bins into quantized bins (default 128 bins) + quantized_bins = np.zeros(num_quantized_bins, dtype=np.int64) + num_merged_bins = sliced_distribution.size // num_quantized_bins + + # merge bins into quantized bins + for index in range(num_quantized_bins): + start = index * num_merged_bins + end = start + num_merged_bins + quantized_bins[index] = sum(sliced_distribution[start:end]) + quantized_bins[-1] += sum(sliced_distribution[num_quantized_bins * num_merged_bins:]) + + # in order to compare p and q, we need to make length of q equals to length of p + # expand quantized bins into p.size bins + q = np.zeros(p.size, dtype=np.int64) + for index in range(num_quantized_bins): + start = index * num_merged_bins + end = start + num_merged_bins + + norm = sum(nonzeros[start:end]) + if norm != 0: + q[start:end] = float(quantized_bins[index]) / float(norm) + + p = smooth_distribution(p) + q = smooth_distribution(q) + + if isinstance(q, np.ndarray): + kl_divergence[i - num_half_quantized_bin] = entropy(p, q) + else: + kl_divergence[i - num_half_quantized_bin] = float('inf') + + min_kl_divergence_idx = np.argmin(kl_divergence) + optimal_threshold = thresholds[min_kl_divergence_idx] + + return optimal_threshold + def create_calibrator(model, op_types_to_calibrate=[], @@ -233,5 +479,7 @@ def create_calibrator(model, calibrate_method=CalibrationMethod.MinMax): if calibrate_method == CalibrationMethod.MinMax: return MinMaxCalibrater(model, op_types_to_calibrate, augmented_model_path) + elif calibrate_method == CalibrationMethod.Entropy: + return EntropyCalibrater(model, op_types_to_calibrate, augmented_model_path) raise ValueError('Unsupported calibration method {}'.format(calibrate_method)) diff --git a/onnxruntime/python/tools/quantization/quant_utils.py b/onnxruntime/python/tools/quantization/quant_utils.py index 5ddfbd2783..d90504bf8f 100644 --- a/onnxruntime/python/tools/quantization/quant_utils.py +++ b/onnxruntime/python/tools/quantization/quant_utils.py @@ -368,3 +368,29 @@ def write_calibration_table(calibration_cache): s = key + ' ' + str(max(abs(value[0]), abs(value[1]))) file.write(s) file.write('\n') + +def smooth_distribution(p, eps=0.0001): + """Given a discrete distribution (may have not been normalized to 1), + smooth it by replacing zeros with eps multiplied by a scaling factor + and taking the corresponding amount off the non-zero values. + Ref: http://web.engr.illinois.edu/~hanj/cs412/bk3/KL-divergence.pdf + https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py + """ + import numpy as np + + is_zeros = (p == 0).astype(np.float32) + is_nonzeros = (p != 0).astype(np.float32) + n_zeros = is_zeros.sum() + n_nonzeros = p.size - n_zeros + + if not n_nonzeros: + # raise ValueError('The discrete probability distribution is malformed. All entries are 0.') + return -1 + eps1 = eps * float(n_zeros) / float(n_nonzeros) + assert eps1 < 1.0, 'n_zeros=%d, n_nonzeros=%d, eps1=%f' % (n_zeros, n_nonzeros, eps1) + + hist = p.astype(np.float32) + hist += eps * is_zeros + (-eps1) * is_nonzeros + assert (hist <= 0).sum() == 0 + + return hist diff --git a/onnxruntime/python/tools/quantization/quantize.py b/onnxruntime/python/tools/quantization/quantize.py index 5120883328..c671c99af9 100644 --- a/onnxruntime/python/tools/quantization/quantize.py +++ b/onnxruntime/python/tools/quantization/quantize.py @@ -24,7 +24,7 @@ from .registry import QLinearOpsRegistry, IntegerOpsRegistry from .onnx_model import ONNXModel from .onnx_quantizer import ONNXQuantizer from .qdq_quantizer import QDQQuantizer -from .calibrate import CalibrationDataReader, create_calibrator +from .calibrate import CalibrationDataReader, create_calibrator, CalibrationMethod def optimize_model(model_path: Path): @@ -145,7 +145,9 @@ def quantize_static(model_input, nodes_to_quantize=[], nodes_to_exclude=[], optimize_model=True, - use_external_data_format=False): + use_external_data_format=False, + calibrate_method=CalibrationMethod.MinMax): + ''' Given an onnx model and calibration data reader, create a quantized onnx model and save it into a file :param model_input: file path of model to quantize @@ -173,6 +175,9 @@ def quantize_static(model_input, when it is not None. :param optimize_model: optimize model before quantization. :parma use_external_data_format: option used for large size (>2GB) model. Set to False by default. + :param calibrate_method: + Current calibration methods supported are MinMax and Entropy. + Please use CalibrationMethod.MinMax or CalibrationMethod.Entropy as options. ''' if activation_type != QuantType.QUInt8: @@ -185,7 +190,7 @@ def quantize_static(model_input, model = load_model(Path(model_input), optimize_model) - calibrator = create_calibrator(model, op_types_to_quantize) + calibrator = create_calibrator(model, op_types_to_quantize, calibrate_method=calibrate_method) calibrator.collect_data(calibration_data_reader) tensors_range = calibrator.compute_range()