Trying to quantize the object detection model of EdgeNet2 but fail, using the pretrained model from here(ms coco 300x300).
qconfig = torch.quantization.get_default_qconfig('qnnpack')
print(qconfig)
model.eval()
model.qconfig = qconfig
torch.quantization.prepare(model, inplace=True)
predictor = BoxPredictor(cfg=cfg, device=device)
#works fine if I do not run the main_images
main_images(predictor=predictor, model=model, object_names=object_names,
in_dir=args.im_dir, out_dir=args.save_dir, device=device)
# Convert to quantized model
torch.quantization.convert(model, inplace=True)
torch.save(model.state_dict(), "espnet2v_300.pt")
print("save as quantize model")
def main_images(predictor, model, object_names, in_dir, out_dir, device='cuda'):
png_file_names = glob.glob(in_dir + os.sep + '*.png')
jpg_file_names = glob.glob(in_dir + os.sep + '*.jpg')
file_names = png_file_names + jpg_file_names
if len(file_names) == 0:
print_error_message('No image files in the folder')
# model in eval mode
model.eval()
with torch.no_grad():
for img_name in file_names:
image = cv2.imread(img_name)
predictor.predict(model, image, is_scaling=False)
If I do not run the main_images function, I can quantize the model, If I run the main_images function, I received error message.
The expanded size of the tensor (243076) must match the existing size (262144) at non-singleton dimension 0. Target sizes: [243076]. Tensor sizes: [262144]
The main_images function works fine if I do not quantize it, any idea how should I fix this?Or do anyone know an object detection project of pytorch suit for quantization? Thanks
ps : Tried with the model of tflite, but accuracy is not that good.