yolov8转openvino并量化

yolov8转openvino并进行INT8量化(python)

使用openvino对yolov8模型进行部署，以下为模型转换和验证测试代码

1.开发环境

win11、python3.10、openvino2023.1.0、ultralytics8.3.55

2.模型转换及INT8量化

import argparse
from openvino.tools import mo
from openvino.runtime import serialize
from openvino.runtime import Core
import os
import nncf #nncf（神经网络压缩框架）用于压缩神经网络,它提供了模型剪枝、量化、稀疏化等压缩方法的工具和技术。

from utils.datasets import create_dataloader #原YOLOv5的函数
# from utils.general import check_dataset
from ultralytics.data.utils import check_det_dataset #函数接口修改

def create_data_source(dataset_yaml,image_size):
    data = check_det_dataset(dataset_yaml)
    val_dataloader = create_dataloader(data['train'], imgsz=image_size, batch_size=1, stride=32, pad=0.5, workers=1)[0]
    return val_dataloader

def transform_fn(data_item):
    # unpack input images tensor
    images = data_item[0]
    # convert input tensor into float format
    images = images.float()
    # scale input
    images = images / 255
    # convert torch tensor to numpy array
    images = images.cpu().detach().numpy()
    return images

def quant_nncf(fp32_path,dataset_yaml,nncf_int8_path,image_size):
    core = Core()
    ov_model = core.read_model(fp32_path)
    subset_size = 120
    preset = nncf.QuantizationPreset.MIXED
    # preset 是量化的预设配置，subset_size 则是用于量化校准的子集大小。
    data_source = create_data_source(dataset_yaml,image_size)
    nncf_calibration_dataset = nncf.Dataset(data_source, transform_func=transform_fn)
    print('wrap data source into nncf.dataset object. ')
    quantized_model = nncf.quantize(
        ov_model, nncf_calibration_dataset,preset=preset,subset_size=subset_size
    )
    serialize(quantized_model, nncf_int8_path)

def onnx2opnvino(onnx_path,fp32_path,fp16_path):
    # fp32 IR model
    print(f"Export ONNX to OpenVINO FP32 to: {fp32_path}")
    model = mo.convert_model(onnx_path)
    serialize(model, fp32_path)

    print(f"Export ONNX to OpenVINO FP16 to: {fp16_path}")
    model = mo.convert_model(onnx_path, compress_to_fp16=True)
    serialize(model, fp16_path)

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model',type=str,default=r'runs\detect\train6\weights\last.onnx',help='model path')
    parser.add_argument('--Int8',type=bool,default="True",help='export int8 model')
    parser.add_argument('--dataset',type=str,default=r'D:\project\YOLO\2024_12_16SKF_conver\ultralytics-main\ultralytics\cfg\datasets\my_detdata_all.yaml',help='dataset yaml file path')
    parser.add_argument('--image_size',type=int,default=640, help='model input image size')

    opt = parser.parse_args()
    return opt

if __name__ == "__main__":
    args = parse_opt()
    model_path=args.model
    sep=os.sep
    model=os.path.split(model_path)[-1]
    model_name=model.split('.')[0]
    model_type=model.split('.')[1]
    model_save_dir=os.path.join(os.path.split(model_path)[0],"opnvino_model")
    fp32_path = f"{model_save_dir}{sep}FP32_openvino_model{sep}{model_name}_fp32.xml"
    fp16_path = f"{model_save_dir}{sep}FP16_openvino_model{sep}{model_name}_fp16.xml"
    onnx2opnvino(args.model,fp32_path,fp16_path)

    #使用nccf进行Int8量化
    if(args.Int8):
        nncf_int8_path = f"{model_save_dir}{sep}NNCF_INT8_openvino_model{sep}{model_name}_int8.xml"
        print("start nccf qua")
        quant_nncf(fp32_path,args.dataset,nncf_int8_path,args.image_size)
        print("nccf qua done")

注意：进行INT8量化后的量化数据的加载，使用了原YOLOv5项目中的datasets.py,该python文件在utils文佳夹中

3.转换后的模型测试

from openvino.runtime import Co