pytorch yolov5 onnx推理

包的版本

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模型

1. 我首先用的是yolov5 v3.0的export.py，但是显然这个代码有问题，采用netron 看导出的模型,最后一层的输出维度是 1 * 3 * 20 * 20 * 9，不能用 ；
2. 然后采用 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