2. YOLOv8模型ONNX导出与OpenVINO部署
任务目标
将训练好的YOLOv8检测模型导出为ONNX格式,使用OpenVINO Runtime完成端到端推理,并与PyTorch原始模型在同一输入下的推理耗时进行对比,验证部署加速效果。
环境依赖
- Python 3.9+
- ultralytics(YOLOv8官方库)
- openvino==2024.1.0(或更高版本)
- onnxruntime(仅用于验证ONNX导出正确性,可选)
- numpy、opencv-python
实现代码
import time
import numpy as np
import cv2
from ultralytics import YOLO
import openvino as ov
# ========== 1. 导出ONNX ==========
model = YOLO("yolov8n.pt")
onnx_path = model.export(format="onnx", opset=12, imgsz=640, simplify=True)
print(f"ONNX已导出: {onnx_path}")
# ========== 2. 准备输入数据 ==========
img = cv2.imread("bus.jpg") # 任意测试图片
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
resized = cv2.resize(img_rgb, (640, 640))
input_tensor = resized.astype(np.float32) / 255.0
input_tensor = np.transpose(input_tensor, (2, 0, 1))[None, ...] # [1,3,640,640]
# ========== 3. PyTorch推理(基准) ==========
model.model.eval()
# 预热
with torch.no_grad():
_ = model.model(torch.from_numpy(input_tensor))
torch.cuda.synchronize() if torch.cuda.is_available() else None
start = time.perf_counter()
with torch.no_grad():
torch_out = model.model(torch.from_numpy(input_tensor))
if torch.cuda.is_available():
torch.cuda.synchronize()
torch_time = (time.perf_counter() - start) * 1000
# ========== 4. OpenVINO推理 ==========
core = ov.Core()
ov_model = core.read_model(onnx_path)
compiled_model = core.compile_model(ov_model, "CPU")
infer_request = compiled_model.create_infer_request()
# 获取输入输出名称
input_name = compiled_model.input(0).get_any_name()
output_names = [out.get_any_name() for out in compiled_model.outputs]
# 预热
infer_request.infer({input_name: input_tensor})
start = time.perf_counter()
infer_request.infer({input_name: input_tensor})
ov_time = (time.perf_counter() - start) * 1000
# 获取输出(YOLOv8输出格式: [1, 84, 8400])
ov_output = infer_request.get_output_tensor(0).data
print(f"OpenVINO输出shape: {ov_output.shape}")
# ========== 5. 结果对比 ==========
print(f"\n=== 推理耗时对比 (CPU) ===")
print(f"PyTorch: {torch_time:.2f} ms")
print(f"OpenVINO: {ov_time:.2f} ms")
print(f"加速比: {torch_time / ov_time:.2f}x")
# 验证输出一致性(取前10个值比较)
torch_np = torch_out[0].detach().cpu().numpy().reshape(84, -1)
ov_np = ov_output.reshape(84, -1)
diff = np.abs(torch_np[:, :10] - ov_np[:, :10]).max()
print(f"前10个输出最大差异: {diff:.6f}")
print("验证通过" if diff < 0.01 else "输出差异过大,请检查导出配置")
运行方式
- 安装依赖:
pip install ultralytics openvino onnxruntime - 准备一张测试图片(如YOLO官方示例
bus.jpg),放在脚本同目录 - 直接运行脚本,首次运行会自动下载yolov8n.pt权重
验证方法与判定标准
- 导出验证:脚本打印
ONNX已导出,且目录下生成yolov8n.onnx文件 - 推理正确性:PyTorch与OpenVINO输出张量shape均为
[1, 84, 8400](COCO 80类+4坐标+1置信度),前10个数值最大差异小于0.01 - 性能对比:OpenVINO在CPU上的推理耗时通常显著低于PyTorch(预期加速1.5~3倍),但具体数值取决于硬件,只需确认OpenVINO耗时低于PyTorch即可判定部署生效
常见问题
- 导出失败:确认ultralytics版本≥8.0.0,且opset≥11
- 输出shape不一致:YOLOv8的ONNX输出为转置后的
[1,84,8400],若得到[1,8400,84]需在导出时添加nms=False参数(默认已包含) - CPU推理仍慢:检查是否误用GPU设备,OpenVINO的CPU加速依赖
openvino包而非onnxruntime