当前位置: 代码迷 >> 综合 >> onnx 模型转换及推理
  详细解决方案

onnx 模型转换及推理

热度:60   发布时间:2024-02-22 03:19:26.0

最近在部署模型在安卓端时发现ncnn模型不及原来的pytorch精度高,误差挺大的,分类误差超过了0.1,怀疑是模型转换,这里就验证一下是不是pytorch -> onnx 模型中产生的误差。

1.模型转换

先来看一下pyotrch 到 onnx 模型转换,这个网上已经很多资料了,这里再贴一下代码:

import torchdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")def model_converter():model = torch.load('test.pth').to(device)  # 这里保存的是完整模型model.eval()dummy_input = torch.randn(1, 3, 96, 96, device=device)input_names = ['data']output_names = ['fc']torch.onnx.export(model, dummy_input, 'test.onnx', export_params=True, verbose=True, input_names=input_names, output_names=output_names)

会输出test.onnx模型。

 

2.onnx 模型 onnx_runtime 推理

import torch
import cv2
import numpy as np
import onnxruntime as rtdef image_process(image_path):mean = np.array([[[0.485, 0.456, 0.406]]])std = np.array([[[0.229, 0.224, 0.225]]])img = cv2.imread(image_path)img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)img = cv2.resize(img, (96, 96))                 # (96, 96, 3)image = img.astype(np.float32)/255.0image = (image - mean)/ stdimage = image.transpose((2, 0, 1))              # (3, 96, 96)image = image[np.newaxis,:,:,:]                 # (1, 3, 96, 96)image = np.array(image, dtype=np.float32)return imagedef onnx_runtime():imgdata = image_process('test.jpg')sess = rt.InferenceSession('test.onnx')input_name = sess.get_inputs()[0].name  label_name = sess.get_outputs()[0].namepred_onnx = sess.run([label_name], {input_name: imgdata})print("outputs:")print(np.array(pred_onnx))if __name__ == '__main__':onnx_runtime()

onnx模型预测结果:

[1.7607212e-04 3.5554171e-05 2.4718046e-04 9.9977958e-01 1.0728532e-01
   3.8951635e-05 4.4435263e-05]

pytorch模型预测结果:

[1.7604830e-04 3.5546487e-05 2.4715674e-04 9.9977964e-01 1.0728550e-01
  3.8929291e-05 4.4419667e-05]

可以看到onnx输出的分类概率和pytorch模型预测的概率几乎完全一样,可见一开始我遇到问题不是模型转换造成的,可能是在移动端图片预处理没弄好造成的。