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

最新推荐文章于 2024-05-22 13:47:11 发布

张○不慌○浩洋

最新推荐文章于 2024-05-22 13:47:11 发布

阅读量4k

点赞数 9

分类专栏：笔记文章标签： pytorch python 深度学习

本文链接：https://blog.csdn.net/weixin_44391370/article/details/120804541

版权

本文记录了一个使用PyTorch的vgg11模型进行二分类任务，然后将模型转换为ONNX格式的过程。在转换后，通过ONNX Runtime进行推理，避免了依赖PyTorch环境。尽管尝试使用opencv.dnn调用ONNX模型遇到了问题，但最终通过其他方法实现了成功转换和推理。文中引用了多位作者的文章作为参考。

摘要由CSDN通过智能技术生成

我做的是一个简单的二分类任务，用了个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