openvino python推理demo

openvino python推理demo

import openvino
from openvino.runtime import Core
import numpy as np
import argparse
import hashlib
import os
import io

class OpenvinoInfer:
    def __init__(self,device_id=0):
        self.device_id=device_id
        self.ie = Core()
        self.available_devices = []
        # 根据可用设备选择
        for name in self.ie.available_devices:
            if name.find("GPU")>=0:
                self.available_devices.append(name)
        print(self.available_devices)
        self.device_name=self.available_devices[self.device_id]
        print("Using device:",self.device_name)
    def build(self,onnx_path,onnx_data=None):  
        if onnx_data is None:
            with open(onnx_path, 'rb') as model:
                onnx_data=model.read()

        md5_hash = hashlib.md5(onnx_data).hexdigest()
        self.cache_path = f"{md5_hash}.engine"
        if not os.path.exists(self.cache_path):
            print("Building engine")
            model = self.ie.read_model(model=onnx_path)
            self.compiled_model = self.ie.compile_model(model=model, device_name=self.device_name)
            user_stream = io.BytesIO()
            self.compiled_model.export_model(user_stream)
            with open(self.cache_path, 'wb') as f:
                f.write(user_stream.getvalue())
        else:
            print("Load engine from cache")
            file = open(self.cache_path, 'br') 
            user_stream = io.BytesIO(file.read())
            self.compiled_model = self.ie.import_model(model_stream=user_stream, device_name=self.device_name)
        
        dtype_map={openvino.Type.f32:np.float32,
                    openvino.Type.i32:np.int32,
                    openvino.Type.i64:np.int64}
        
        self.inputs=[]
        self.outputs=[]
        for input_layer in self.compiled_model.inputs:
            print(f"输入层名称：{input_layer.any_name}, 形状：{input_layer.shape} dtype:{input_layer.element_type}")
            self.inputs.append({"name": input_layer.any_name, "shape": input_layer.shape, "dtype": dtype_map[input_layer.element_type]})
        for output_layer in self.compiled_model.outputs:
            print(f"输入层名称：{output_layer.any_name}, 形状：{output_layer.shape} dtype:{output_layer.element_type}")
            self.outputs.append({"name": output_layer.any_name, "shape": output_layer.shape, "dtype": dtype_map[output_layer.element_type]})
        return True

    def inference(self,inputs):
        args=[]
        for idx,ipt in enumerate(inputs):
            args.append(ipt.reshape(self.inputs[idx]['shape']))
        result = self.compiled_model(args)
        output_data = []
        for item in self.outputs: 
            output_data.append(result[item['name']])
        return output_data

if __name__ == '__main__':
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='ai_model_trt_infer')
    parser.add_argument('--model', type=str, help='model')
    parser.add_argument('--input_paths', type=str, help='inputs')
    parser.add_argument('--input_dtypes', type=str, help='inputs')
    parser.add_argument('--output_paths', type=str, help='outputs')
    parser.add_argument('--output_dtypes', type=str, help='outputs')
    args = parser.parse_args()

    dtype_map={"int64":np.int64,"float32":np.float32,"float16":np.float16,"uint8":np.uint8}
    input_paths=args.input_paths.split(',')
    input_dtypes=args.input_dtypes.split(',')

    output_paths=args.output_paths.split(',')
    output_dtypes=args.output_dtypes.split(',')

    infer = OpenvinoInfer()
    infer.build(args.model)

    inputs=[]
    for idx,file_path in enumerate(input_paths):
        with open(file_path, 'rb') as f:
            input_data = np.frombuffer(f.read(), dtype=infer.inputs[idx]['dtype'])
        inputs.append(input_data)

    outputs_gt=[]
    outputs_pred=[]
    for idx,file_path in enumerate(output_paths):
        with open(file_path, 'rb') as f:
            output_data = np.frombuffer(f.read(), dtype=infer.outputs[idx]['dtype'])
        outputs_gt.append(output_data)        
        outputs_pred.append(np.empty(output_data.shape, dtype=output_data.dtype))


    outputs_pred=infer.inference(inputs)

    for idx,output_data in enumerate(outputs_pred):
        mse = np.mean((output_data.reshape(-1) - outputs_gt[idx].reshape(-1)) ** 2)
        print("均方误差 (MSE):", mse)

用法

python ai_model_openvino_infer.py \
    --model=resnet50.onnx \
    --input_paths=resnet50-input-input.bin \
    --output_paths=resnet50-output-output.bin \
    --input_dtypes="float32" \
    --output_dtypes="float32"

python ai_model_openvino_infer.py \
    --model=yolov5m.onnx \
    --input_paths=yolov5m-images-input.bin \
    --output_paths=yolov5m-output0-output.bin \
    --input_dtypes="float32" \
    --output_dtypes="float32"    

python ai_model_openvino_infer.py \
    --model=bert-base.onnx \
    --input_paths=bert-base-input_ids-input.bin,bert-base-attention_mask-input.bin,bert-base-token_type_ids-input.bin \
    --output_paths=bert-base-uncased-output.bin \
    --input_dtypes="int64,int64" \
    --output_dtypes="float32"