pytorch模型部署到onnx的详细过程

更新时间：2023年08月28日 10:25:53 作者：Liang-ml

这篇文章主要介绍了如何简单的将pytorch模型部署到onnx,本文结合示例代码给大家介绍的非常详细,对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下

1. 创建一个pytorch模型

这里我用的U2Net，直接加载好训练出的权重

model = U2Net(class_nums=4)
model.load_state_dict(torch.load(checkpoint_path))

2. 将pytorch模型转成onnx格式

x = torcg.randn(1,3,512,512)
with torch.no_grad():
torch.onnx.export(
model,
x,
"onnx_file_name.onnx",
opset_version=11,
input_names=['input'],
output_names=['output']
)

torch.onnx.export()函数用于将pytorch模型转成onnx格式的函数，其中的参数如下

model：需要转换的模型，加载好权重
args：模型的任意一组输入，注意维度
f：导出的onnx模型的文件名
opset_version：表示ONNX算子集的版本，随着深度学习的发展，新算子会不断诞生，算子集就需要扩充，截至20230824，官网已经发布20个版本
input_names：输入tensor的名称
output_names：输出tensor的名称

转换成功会在指定的路径f下生成onnx模型文件

可用onnx中的函数验证模型文件是否正确

import onnx
onnx_model = onnx.load("onnx_file_name.onnx")
try:
onnx.checker.check_model(onnx_model)
except Exception as ex:
print(f'ERROR: {ex}')
else:
print(f'model correct')

可以使用netron查看转换的onnx模型的详细内容链接 https://netron.app/

点击某个算子节点，可以查看算子的具体信息，每个算子记录了算子的属性、图结构和权重三类信息

算子属性信息(attributes)，对于卷积来说，算子属性包含了卷积核的大小、步长等属性
图结构信息(node properties)，算子节点在计算图中的名称、邻边的信息，如上图节点名为conv2，输入数据叫683，权重叫1324
权重信息指算子存储的网络权重信息，点击w后的+号可以查看权重信息的具体内容

3. 推理引擎ONNXRuntime

使用onnxruntime可以运行输出的onnx模型，代码如下

input_img 是一个与onnx模型输入维度一致的归一化的np.array
model = onnxruntime.InferenceSession(onnx_file)
inputs = {'input':input_img}
output = model.run(['output'],inputs)[0]

onnxruntime.InferenceSession用于获取一个ONNX Runtime的推理器，其参数为用于推理的ONNX模型文件

run方法用于模型推理，第一个参数为输出张量名的列表，第二个参数为输入值的字典，key为导出模型时设置的张量名，value为输入张量

到此这篇关于简单的将pytorch模型部署到onnx的文章就介绍到这了,更多相关pytorch模型部署到onnx内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:

python opencv3实现人脸识别(windows)
这篇文章主要为大家详细介绍了python opencv3实现人脸识别程序，文中示例代码介绍的非常详细，具有一定的参考价值，感兴趣的小伙伴们可以参考一下
2018-05-05
使用Dataframe.info()显示空值与类型信息
这篇文章主要介绍了使用Dataframe.info()显示空值与类型信息,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
2024-02-02
浅析python 字典嵌套
这篇文章主要介绍了python 字典嵌套的相关资料，帮助大家更好的理解和学习python，感兴趣的朋友可以了解下
2020-09-09
python中Flask框架简单入门实例
这篇文章主要介绍了python中Flask框架简单入门实例,以一个hello程序简单分析了Flask框架的使用技巧,需要的朋友可以参考下
2015-03-03
Python使用PyAudio制作录音工具的实现代码
这篇文章主要介绍了Python使用PyAudio制作录音工具,音频录制与视频录制相似，也是以数据帧的方式录制保存，这次使用强大的第三方包PyAudio和内置的wave模块编写，需要的朋友可以参考下
2022-04-04
django2.2安装错误最全的解决方案(小结)
这篇文章主要介绍了django2.2安装错误最全的解决方案(小结)，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2019-09-09
python selenium 对浏览器标签页进行关闭和切换的方法
今天小编就为大家分享一篇python selenium 对浏览器标签页进行关闭和切换的方法，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2018-05-05
python代码实现逻辑回归logistic原理
这篇文章主要介绍了python代码实现逻辑回归logistic原理，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2019-08-08
Python依赖管理及打包工具Poetry使用规范
这篇文章主要为大家介绍了Python依赖管理及打包工具Poetry的依赖规范，有需要的朋友可以借鉴参考下，希望能够有所帮助，祝大家多多进步
2021-09-09
解决Jupyter NoteBook输出的图表太小看不清问题
这篇文章主要介绍了解决Jupyter NoteBook输出的图表太小看不清问题，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2020-04-04