当前位置：首页 > news >正文

openvino python推理demo

news 2024/11/4 17:30:11

import openvino
from openvino.runtime import Core
import numpy as np
import argparse
import hashlib
import os
import ioclass OpenvinoInfer:def __init__(self,device_id=0):self.device_id=device_idself.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 Truedef 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_dataif __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"