From 4bf72dec1f71feea0d518d8963807ebdc166c3f3 Mon Sep 17 00:00:00 2001 From: Chen Fu <1316708+chenfucn@users.noreply.github.com> Date: Thu, 18 Aug 2022 13:34:44 -0700 Subject: [PATCH] bug fix --- .../cpu/resnet50_data_reader.py | 61 ++++++++++ quantization/image_classification/cpu/run.py | 107 +++++------------- 2 files changed, 90 insertions(+), 78 deletions(-) create mode 100644 quantization/image_classification/cpu/resnet50_data_reader.py diff --git a/quantization/image_classification/cpu/resnet50_data_reader.py b/quantization/image_classification/cpu/resnet50_data_reader.py new file mode 100644 index 000000000..bf929bce6 --- /dev/null +++ b/quantization/image_classification/cpu/resnet50_data_reader.py @@ -0,0 +1,61 @@ +import numpy +import onnxruntime +import os +from onnxruntime.quantization import CalibrationDataReader +from PIL import Image + + +def _preprocess_images(images_folder: str, height: int, width: int, 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 + """ + image_names = os.listdir(images_folder) + if size_limit > 0 and len(image_names) >= size_limit: + batch_filenames = [image_names[i] for i in range(size_limit)] + else: + batch_filenames = image_names + unconcatenated_batch_data = [] + + for image_name in batch_filenames: + image_filepath = images_folder + "/" + image_name + pillow_img = Image.new("RGB", (width, height)) + pillow_img.paste(Image.open(image_filepath).resize((width, height))) + input_data = numpy.float32(pillow_img) - numpy.array( + [123.68, 116.78, 103.94], dtype=numpy.float32 + ) + nhwc_data = numpy.expand_dims(input_data, axis=0) + nchw_data = nhwc_data.transpose(0, 3, 1, 2) # ONNX Runtime standard + unconcatenated_batch_data.append(nchw_data) + batch_data = numpy.concatenate( + numpy.expand_dims(unconcatenated_batch_data, axis=0), axis=0 + ) + return batch_data + + +class ResNet50DataReader(CalibrationDataReader): + def __init__(self, calibration_image_folder: str, model_path: str): + self.image_folder = calibration_image_folder + self.model_path = model_path + self.preprocess_flag = True + self.enum_data_dicts = [] + self.datasize = 0 + + def get_next(self): + if self.preprocess_flag: + self.preprocess_flag = False + session = onnxruntime.InferenceSession(self.model_path, None) + (_, _, height, width) = session.get_inputs()[0].shape + nhwc_data_list = _preprocess_images( + self.image_folder, height, width, size_limit=0 + ) + input_name = session.get_inputs()[0].name + self.datasize = len(nhwc_data_list) + self.enum_data_dicts = iter( + [{input_name: nhwc_data} for nhwc_data in nhwc_data_list] + ) + return next(self.enum_data_dicts, None) diff --git a/quantization/image_classification/cpu/run.py b/quantization/image_classification/cpu/run.py index 3f48057f9..82088b9da 100644 --- a/quantization/image_classification/cpu/run.py +++ b/quantization/image_classification/cpu/run.py @@ -1,67 +1,10 @@ -import os -import sys -import numpy as np -import re -import abc -import subprocess -import json import argparse -import time -from PIL import Image - -import onnx +import numpy as np import onnxruntime -from onnx import helper, TensorProto, numpy_helper -from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantFormat, QuantType - - -class ResNet50DataReader(CalibrationDataReader): - def __init__(self, calibration_image_folder, augmented_model_path='augmented_model.onnx'): - self.image_folder = calibration_image_folder - self.augmented_model_path = augmented_model_path - self.preprocess_flag = True - self.enum_data_dicts = [] - self.datasize = 0 - - def get_next(self): - if self.preprocess_flag: - self.preprocess_flag = False - session = onnxruntime.InferenceSession(self.augmented_model_path, None) - (_, _, height, width) = session.get_inputs()[0].shape - nhwc_data_list = preprocess_func(self.image_folder, height, width, size_limit=0) - input_name = session.get_inputs()[0].name - self.datasize = len(nhwc_data_list) - self.enum_data_dicts = iter([{input_name: nhwc_data} for nhwc_data in nhwc_data_list]) - return next(self.enum_data_dicts, None) - - -def preprocess_func(images_folder, height, width, 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 - ''' - image_names = os.listdir(images_folder) - if size_limit > 0 and len(image_names) >= size_limit: - batch_filenames = [image_names[i] for i in range(size_limit)] - else: - batch_filenames = image_names - unconcatenated_batch_data = [] +import time +from onnxruntime.quantization import QuantFormat, QuantType, quantize_static - for image_name in batch_filenames: - image_filepath = images_folder + '/' + image_name - pillow_img = Image.new("RGB", (width, height)) - pillow_img.paste(Image.open(image_filepath).resize((width, height))) - 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 - unconcatenated_batch_data.append(nchw_data) - batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0) - return batch_data +import resnet50_data_reader def benchmark(model_path): @@ -87,11 +30,15 @@ def get_args(): parser = argparse.ArgumentParser() parser.add_argument("--input_model", required=True, help="input model") parser.add_argument("--output_model", required=True, help="output model") - parser.add_argument("--calibrate_dataset", default="./test_images", help="calibration data set") - parser.add_argument("--quant_format", - default=QuantFormat.QDQ, - type=QuantFormat.from_string, - choices=list(QuantFormat)) + parser.add_argument( + "--calibrate_dataset", default="./test_images", help="calibration data set" + ) + parser.add_argument( + "--quant_format", + default=QuantFormat.QDQ, + type=QuantFormat.from_string, + choices=list(QuantFormat), + ) parser.add_argument("--per_channel", default=False, type=bool) args = parser.parse_args() return args @@ -102,21 +49,25 @@ def main(): input_model_path = args.input_model output_model_path = args.output_model calibration_dataset_path = args.calibrate_dataset - dr = ResNet50DataReader(calibration_dataset_path) - quantize_static(input_model_path, - output_model_path, - dr, - quant_format=args.quant_format, - per_channel=args.per_channel, - weight_type=QuantType.QInt8) - print('Calibrated and quantized model saved.') - - print('benchmarking fp32 model...') + dr = resnet50_data_reader.ResNet50DataReader( + calibration_dataset_path, input_model_path + ) + quantize_static( + input_model_path, + output_model_path, + dr, + quant_format=args.quant_format, + per_channel=args.per_channel, + weight_type=QuantType.QInt8, + ) + print("Calibrated and quantized model saved.") + + print("benchmarking fp32 model...") benchmark(input_model_path) - print('benchmarking int8 model...') + print("benchmarking int8 model...") benchmark(output_model_path) -if __name__ == '__main__': +if __name__ == "__main__": main()