【模型推理】使用onnx进行推理

举个栗子:

模型从pytorch转换成onnx后如何使用onnx文件进行推理呢?

今天分享一个脚本用于使用onnx文件进行推理,onnx文件需要利用onnxruntime工具包进行推理。

import numpy as np
import onnxruntime as ort
import torch
import cv2
def preict_one_img(img_path):
    img = cv2.imread(img_path) #读取图片
    img = cv2.resize(img, (299, 299))#调整图片尺寸
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 把图片BGR变成RGB

    img = np.transpose(img,(2,0,1))#调整维度将HWC - CHW
    img = np.expand_dims(img, 0) #添加一个维度 就是batch维度
    img = img.astype(np.float32)#格式转成float32
    img /= 255
   #调用onnxruntime run函数进行模型推理
    outputs = ort_session.run(
        None,
        {"images": img},
    )
    #outputs的输出类型为list类型,所以要先将list转换成numpy再转换成torch
    outputs1 = torch.from_numpy(np.array(outputs))
    #通过softmax进行最后分数的计算
    outputs_softmax = torch.softmax(outputs1[0], dim=1).numpy()[:, 0].tolist()[0]
   
if __name__ == '__main__':
    #cpu or gpu
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    #onnx路径
    model_path = "1.onnx"
    #加载onnx模型
    ort_session = ort.InferenceSession(model_path, providers=device)
    #图片路径
    i='.jpg'
    preict_one_img(i)

这里我只是拿出一个简单的分类模型做一个演示,比如应用到目标检测最后结果的解码逻辑要符合回归和分类这两部分。 但是推理onnx是不变的都是先加载onnx然后处理图片将处理好的图片送入onnxruntime.run进行推理再进行解码。

猜你喜欢

转载自blog.csdn.net/qq_44992785/article/details/129180382