|
|
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() |
|
|
|
