pytorch转onnx模型并进行推理

前言： 项目需求，还是之前那个缺陷检测的项目，因为要落地到实际生产线上，所以pytorch模型要进行部署，因为libtorch在使用的时候不知道啥原因不能用opencv，所以试试转onnx模型再用dnn模块进行推理，记录一下。

首先是pytorch转onnx模型

参考

import torch
import torchvision
# torch_model = torchvision.models.resnet50(pretrained=True)  # pytorch模型加载
torch_model = torch.load("epoch_40.pth")['model']  # load自己的模型

batch_size = 1  #批处理大小
input_shape = (3,224,224)   #输入数据

# set the model to inference mode
torch_model.eval().cuda()

x = torch.randn(batch_size,*input_shape).cuda()		# 生成张量
export_onnx_file = "resnet50----.onnx"					# 目的ONNX文件名
torch.onnx.export(torch_model,
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],		# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={
     "input":{
     0:"batch_size"},		# 批处理变量
                                    "output":{
     0:"batch_size"}})

**

推理代码：

这个代码也是网上参考的，找不到链接了,注意代码里面的图像前处理操作要跟自己pytorch里面的一致。
**

# -*-coding: utf-8 -*-

import os, sys

sys.path.append(os.getcwd())
import onnxruntime
import onnx
import cv2
import torch
import numpy as np
import torchvision.transforms as transforms
import time
import math



class ONNXModel():
    def __init__(self, onnx_path):
        """
        :param onnx_path:
        """
        self.onnx_session = onnxruntime.InferenceSession(onnx_path)
        self.input_name = self.get_input_name(self.onnx_session)
        self.output_name = self.get_output_name(self.onnx_session)
        print("input_name:{}".format(self.input_name))
        print("output_name:{}".format(self.output_name))

    def get_output_name(self, onnx_session):
        """
        output_name = onnx_session.get_outputs()[0].name
        :param onnx_session:
        :return:
        """
        output_name = []
        for node in onnx_session.get_outputs():
            output_name.append(node.name)
        return output_name

    def get_input_name(self, onnx_session):
        """
        input_name = onnx_session.get_inputs()[0].name
        :param onnx_session:
        :return:
        """
        input_name = []
        for node in onnx_session.get_inputs():
            input_name.append(node.name)
        return input_name

    def get_input_feed(self, input_name, image_numpy):
        """
        input_feed={self.input_name: image_numpy}
        :param input_name:
        :param image_numpy:
        :return:
        """
        input_feed = {
     }
        for name in input_name:
            input_feed[name] = image_numpy
        return input_feed

    def forward(self, image_numpy):
        '''
        # image_numpy = image.transpose(2, 0, 1)
        # image_numpy = image_numpy[np.newaxis, :]
        # onnx_session.run([output_name], {input_name: x})
        # :param image_numpy:
        # :return:
        '''
        # 输入数据的类型必须与模型一致,以下三种写法都是可以的
        # scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})
        # scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})
        input_feed = self.get_input_feed(self.input_name, image_numpy)
        scores = self.onnx_session.run(self.output_name, input_feed=input_feed)
        return scores


def to_numpy(tensor):
    print(tensor.device)
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


labels = ['crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale', 'scratches']
r_model_path = "resnet50----.onnx"

time_start1 = time.time()
rnet1 = ONNXModel(r_model_path)
time_end2 = time.time()
print('load model cost', time_end2 - time_start1)

# 测时间
for i in range(5):
    time_start = time.time()
    img_ori = cv2.imread(r"C:\Users\59593\Desktop\pytorch_classification-master\samples\rolled-in_scale_155.jpg")
    img = cv2.resize(img_ori, (224, 224), interpolation = cv2.INTER_CUBIC)
    img_input = img[..., ::-1] # BGR to RGB
    img_input = (np.float32(img)/255.0-[0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
    img_input = img_input.transpose((2, 0, 1))
    img_input = torch.from_numpy(img_input).unsqueeze(0)
    img_input = img_input.type(torch.FloatTensor)

    out = rnet1.forward(to_numpy(img_input))

    print(out)
    print(labels[np.argmax(out[0][0])])

    time_end=time.time()
    print('infer cost',time_end-time_start)

    # cv2.putText(img_ori, labels[np.argmax(out[0][0])], (50,50), cv2.FONT_HERSHEY_COMPLEX, 0.5, (100, 200, 200), 1)
    # cv2.imshow("1", img_ori)
    # cv2.waitKey(0)

实验：

按上面下来的话，两个模型的输出是一致的，这一点没问题。
在自己的笔记本上实验，图像大小224*224，缺陷六分类，1660Ti，可以看出推理速度确实快了。

1. pytorch下加载模型要1.4s，onnx加载模型大概2s
2. pytorch一张图要0.02s-0.03s左右，而转换成onnx模型后一张图要0.01s左右

• 然后接下来尝试C++ 编译的GPU版opencv dnn模块来调用onnx模型，并且配上QT界面实现。