这里写自定义目录标题
包的版本
cuda 10.1
torch 1.6.0+cu101
torchvision 0.7.0+cu101
onnx 1.7.0
onnx-simplifier 0.3.6
onnxoptimizer 0.2.6
onnxruntime 1.8.0
onnxruntime-gpu 1.3.0
或者
cuda 10.2
onnx 1.8.0
onnx-simplifier 0.3.4
onnxoptimizer 0.2.5
onnxruntime 1.7.0
onnxruntime-gpu 1.5.1
torch 1.8.1
torch-tb-profiler 0.3.1
torchsummary 1.5.1
torchvision 0.9.1
能不能推理成功,跟各个包的版本很大关系,多试onnxruntime的版本吧!
pt转onnx
python export.py --weights exp8/weights/best.pt --img 640 --batch 1
python -m onnxsim best.onnx best-sm.onnx
生成best-sm.onnx模型
- 我首先用的是yolov5 v3.0的export.py,但是显然这个代码有问题,采用netron 看导出的模型,最后一层的输出维度是 1 * 3 * 20 * 20 * 9,
不能用; - 然后采用 yolov5 v6.0的export.py,采用netron 看导出的模型,最后一层的输出维度是 1 * 25200 * 9,能用;
onnx 推理
编写推理脚本,参考这里
# coding=utf-8
import cv2.cv2 as cv2
import numpy as np
import onnxruntime
import torch
import torchvision
import time
import random
from utils.general import non_max_suppression
class YOLOV5_ONNX(object):
def __init__(self,onnx_path):
'''初始化onnx'''
self.onnx_session=onnxruntime.InferenceSession(onnx_path)
print(onnxruntime.get_device())
self.input_name=self.get_input_name()
self.output_name=self.get_output_name()
self.classes=['person', 'car', 'special_person', 'truck']
def get_input_name(self):
'''获取输入节点名称'''
input_name=[]
for node in self.onnx_session.get_inputs():
input_name.append(node.name)
return input_name
def get_output_name(self):
'''获取输出节点名称'''
output_name=[]
for node in self.onnx_session.get_outputs():
output_name.append(node.name)
return output_name
def get_input_feed(self,image_tensor):
'''获取输入tensor'''
input_feed={
}
for name in self.input_name:
input_feed[name]=image_tensor
return input_feed
def letterbox(self,img, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True,
stride=32):
'''图片归一化'''
# Resize and pad image while meeting stride-multiple constraints
shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup: # only scale down, do not scale up (for better test mAP)
r = min(r, 1.0)
# Compute padding
ratio = r, r # width, height ratios
new_unpad = int(round