Pytorch的pth模型转onnx，再用ONNX Runtime调用推理（附python代码）

Posted 2021-11-04 张○不慌○浩洋

我做的是一个简单的二分类任务，用了个vgg11，因为要部署到应用，所以将 PyTorch 中定义的模型转换为 ONNX 格式，然后在 ONNX Runtime 中运行它，那就不用了在机子上配pytorch环境了。然后也试过转出来的onnx用opencv.dnn来调用，发现识别完全不对，据说是opencv的那个包只能做二维的pooling层，不能做三维的。

然后具体的模型转换以及使用如下代码所示，仅作为学习笔记咯~（亲测可用）

pip install onnx
pip install onnxruntime

首先，将Pytorch模型转成onnx格式，然后验证一波onnx模型有没有什么毛病

# coding=gbk
#_*_ coding=utf-8 _*_
import torch
import torchvision
import torch.nn as nn
from vgg import vgg11_bn
from torchvision import models
import time

out_onnx = 'model.onnx'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dummy = torch.randn(1, 1, 128, 128)   # 模型的输入格式
model = vgg11_bn()					  # 模型
model = nn.DataParallel(model)        #　因为这里我用了多线程训练，所以得加上
model.load_state_dict(torch.load('model.pth', map_location='cuda'))
if isinstance(model,torch.nn.DataParallel):
	model = model.module

model = model.to(device)
dummy = dummy.to(device)
#　定义输入的名字和输出的名字，好像可有可无
input_names = ["input"]
output_names = ["output"]
#　输出pytorch to onnx
torch_out = torch.onnx.export(model, dummy, out_onnx, input_names=input_names, output_names=output_names)
print("finish!") # 搞定

time.sleep(5)
# 验证 Check the model
import onnx
onnx_model = onnx.load(out_onnx)
print('The model is:\\n{}'.format(onnx_model))
try:
    onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
    print('The model is invalid: %s' % e)
else:
    print('The model is valid!')

output = onnx_model.graph.output
print(output)

然后就是用onnx在python不import torch的情况下做推理的测试：

# coding=gbk
#_*_ coding=utf-8 _*_
import numpy as np
from PIL import Image

img_input = "test.jpg" 
imagec = Image.open(img_input).convert("L")                 # 加载图像，我的是灰度图
imagec = imagec.resize((128, 128))								# resize
# 我参考的代码是先ToTensor，然后做ununsqueeze到(1, 1, 128, 128)，
# 再to_numpy，像素值是0-1，而PIL的是0-255，所以这里除了个255，再做ununsqueeze
image = np.array(imagec,dtype=np.float32)/255.
image = np.expand_dims(image, axis=0)
image = np.expand_dims(image, axis=0)
print(image)  # ---> (1, 1, 128, 128)

import onnxruntime
import onnx
##onnx测试
onnx_model_path = "models.onnx"
session = onnxruntime.InferenceSession(onnx_model_path)
#compute ONNX Runtime output prediction
inputs = {session.get_inputs()[0].name: image}
logits = session.run(None, inputs)[0]

print("onnx weights", logits)
print("onnx prediction", logits.argmax(axis=1)[0])

参考了很多大佬的文章

https://blog.csdn.net/ouening/article/details/109245243
https://zhuanlan.zhihu.com/p/363177241
https://www.cnblogs.com/sddai/p/14537381.html
还有的大佬链接找不到了，感谢大佬们