onnx模型转换及推理

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

1.模型转换

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

import torch
 
device = 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 cv2
import numpy as np
import onnxruntime as rt
 
def image_process(image_path):
    mean = np.array([[[0.485, 0.456, 0.406]]])      # 训练的时候用来mean和std
    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.0
    image = (image - mean)/ std
 
    image = image.transpose((2, 0, 1))              # (3, 96, 96)
    image = image[np.newaxis,:,:,:]                 # (1, 3, 96, 96)
 
    image = np.array(image, dtype=np.float32)
    
    return image
 
def onnx_runtime():
    imgdata = image_process('test.jpg')
    
    sess = rt.InferenceSession('test.onnx')
    input_name = sess.get_inputs()[0].name  
    output_name = sess.get_outputs()[0].name
 
    pred_onnx = sess.run([output_name], {input_name: imgdata})
 
    print("outputs:")
    print(np.array(pred_onnx))
 
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模型预测的概率几乎完全一样，可见一开始我遇到问题不是pytorch->onnx 模型转换造成的，可能是在移动端图片预处理没弄好造成的。

 

参考：onnx 模型转换及推理

一、pytorch模型保存/加载

有两种方式可用于保存/加载pytorch模型 1）文件中保存模型结构和权重参数 2）文件只保留模型权重.

1、文件中保存模型结构和权重参数

1）pytorch模型保存

import torch
torch.save(selfmodel,"save.pt")

2）pytorch模型加载

import torch
torch.load("save.pt")

2、文件只保留模型权重

1）pytorch模型保存

import torch
torch.save(selfmodel.state_dict(),"save.pt")

2）pytorch模型加载

selfmodel.load_state_dict(torch.load("save.pt"))

二、pytorch模型转ONNX模型

1、文件中保存模型结构和权重参数

import torch
torch_model = torch.load("save.pt") # pytorch模型加载
batch_size = 1  #批处理大小
input_shape = (3,244,244)   #输入数据

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

x = torch.randn(batch_size,*input_shape)		# 生成张量
export_onnx_file = "test.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"}})

注：dynamic_axes字段用于批处理.若不想支持批处理或固定批处理大小,移除dynamic_axes字段即可.

2、文件中只保留模型权重

import torch
torch_model = selfmodel()  					# 由研究员提供python.py文件
batch_size = 1 								# 批处理大小
input_shape = (3, 244, 244) 				# 输入数据

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

x = torch.randn(batch_size,*input_shape) 	# 生成张量
export_onnx_file = "test.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"}})

链接:
1）pytorch官方文档.

 

参考：pytorch模型转ONNX模型(支持批处理)