| import os |
| import sys |
| import numpy as np |
| import cv2 as cv |
|
|
| import onnx |
| from neural_compressor.experimental import Quantization, common |
| from neural_compressor.experimental.metric import BaseMetric |
|
|
|
|
| class Accuracy(BaseMetric): |
| def __init__(self, *args): |
| self.pred_list = [] |
| self.label_list = [] |
| self.samples = 0 |
|
|
| def update(self, predict, label): |
| predict = np.array(predict) |
| label = np.array(label) |
| self.pred_list.append(np.argmax(predict[0])) |
| self.label_list.append(label[0][0]) |
| self.samples += 1 |
|
|
| def reset(self): |
| self.pred_list = [] |
| self.label_list = [] |
| self.samples = 0 |
|
|
| def result(self): |
| correct_num = np.sum(np.array(self.pred_list) == np.array(self.label_list)) |
| return correct_num / self.samples |
|
|
|
|
| class Quantize: |
| def __init__(self, model_path, config_path, custom_dataset=None, eval_dataset=None, metric=None): |
| self.model_path = model_path |
| self.config_path = config_path |
| self.custom_dataset = custom_dataset |
| self.eval_dataset = eval_dataset |
| self.metric = metric |
|
|
| def run(self): |
| print('Quantizing (int8) with Intel\'s Neural Compressor:') |
| print('\tModel: {}'.format(self.model_path)) |
| print('\tConfig: {}'.format(self.config_path)) |
|
|
| output_name = '{}-int8-quantized.onnx'.format(self.model_path[:-5]) |
|
|
| model = onnx.load(self.model_path) |
| quantizer = Quantization(self.config_path) |
| quantizer.model = common.Model(model) |
| if self.custom_dataset is not None: |
| quantizer.calib_dataloader = common.DataLoader(self.custom_dataset) |
| if self.eval_dataset is not None: |
| quantizer.eval_dataloader = common.DataLoader(self.eval_dataset) |
| if self.metric is not None: |
| quantizer.metric = common.Metric(metric_cls=self.metric, name='metric') |
| q_model = quantizer() |
| q_model.save(output_name) |
|
|
|
|
| class Dataset: |
| def __init__(self, root, size=None, dim='chw', scale=1.0, mean=0.0, std=1.0, swapRB=False, toFP32=False): |
| self.root = root |
| self.size = size |
| self.dim = dim |
| self.scale = scale |
| self.mean = mean |
| self.std = std |
| self.swapRB = swapRB |
| self.toFP32 = toFP32 |
|
|
| self.image_list, self.label_list = self.load_image_list(self.root) |
|
|
| def load_image_list(self, path): |
| image_list = [] |
| label_list = [] |
| for f in os.listdir(path): |
| if not f.endswith('.jpg'): |
| continue |
| image_list.append(os.path.join(path, f)) |
| label_list.append(1) |
| return image_list, label_list |
|
|
| def __getitem__(self, idx): |
| img = cv.imread(self.image_list[idx]) |
|
|
| if self.swapRB: |
| img = cv.cvtColor(img, cv.COLOR_BGR2RGB) |
|
|
| if self.size: |
| img = cv.resize(img, dsize=self.size) |
|
|
| if self.toFP32: |
| img = img.astype(np.float32) |
|
|
| img = img * self.scale |
| img = img - self.mean |
| img = img / self.std |
|
|
| if self.dim == 'chw': |
| img = img.transpose(2, 0, 1) |
|
|
| return img, self.label_list[idx] |
|
|
| def __len__(self): |
| return len(self.image_list) |
|
|
|
|
| class FerDataset(Dataset): |
| def __init__(self, root, size=None, dim='chw', scale=1.0, mean=0.0, std=1.0, swapRB=False, toFP32=False): |
| super(FerDataset, self).__init__(root, size, dim, scale, mean, std, swapRB, toFP32) |
|
|
| def load_image_list(self, path): |
| image_list = [] |
| label_list = [] |
| for f in os.listdir(path): |
| if not f.endswith('.jpg'): |
| continue |
| image_list.append(os.path.join(path, f)) |
| label_list.append(int(f.split("_")[2])) |
| return image_list, label_list |
|
|
|
|
| models = dict( |
| mobilenetv1=Quantize(model_path='../../models/image_classification_mobilenet/image_classification_mobilenetv1_2022apr.onnx', |
| config_path='./inc_configs/mobilenet.yaml'), |
| mobilenetv2=Quantize(model_path='../../models/image_classification_mobilenet/image_classification_mobilenetv2_2022apr.onnx', |
| config_path='./inc_configs/mobilenet.yaml'), |
| mp_handpose=Quantize(model_path='../../models/handpose_estimation_mediapipe/handpose_estimation_mediapipe_2022may.onnx', |
| config_path='./inc_configs/mp_handpose.yaml', |
| custom_dataset=Dataset(root='../../benchmark/data/palm_detection', dim='hwc', swapRB=True, mean=127.5, std=127.5, toFP32=True)), |
| fer=Quantize(model_path='../../models/facial_expression_recognition/facial_expression_recognition_mobilefacenet_2022july.onnx', |
| config_path='./inc_configs/fer.yaml', |
| custom_dataset=FerDataset(root='../../benchmark/data/facial_expression_recognition/fer_calibration', size=(112, 112), toFP32=True, swapRB=True, scale=1./255, mean=0.5, std=0.5), |
| eval_dataset=FerDataset(root='../../benchmark/data/facial_expression_recognition/fer_evaluation', size=(112, 112), toFP32=True, swapRB=True, scale=1./255, mean=0.5, std=0.5), |
| metric=Accuracy), |
| ) |
|
|
| if __name__ == '__main__': |
| selected_models = [] |
| for i in range(1, len(sys.argv)): |
| selected_models.append(sys.argv[i]) |
| if not selected_models: |
| selected_models = list(models.keys()) |
| print('Models to be quantized: {}'.format(str(selected_models))) |
|
|
| for selected_model_name in selected_models: |
| q = models[selected_model_name] |
| q.run() |
|
|
