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 缺失值处理的方法（Imputation），文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2019-07-07
这篇文章主要介绍了python web应用程序之Django数据库,本文通过实例代码给大家介绍的非常详细,需要的朋友可以参考下2024-06-06
这篇文章主要介绍了基于python实现把json数据转换成Excel表格,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下2020-05-05
CAPL能做超级多的功能,本文主要介绍了CAPL与Python交互的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧2024-05-05
本文主要介绍了Python中typing模块的具体使用，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2022-05-05
这篇文章给大家介绍了Python开发最牛逼的IDE——pycharm，主要是介绍python IDE pycharm的安装与使用教程，非常不错，具有一定的参考借鉴价值，需要的朋友参考下吧2018-08-08
详解Django-channels 实现WebSocket实例
这篇文章主要介绍了详解Django-channels实现WebSocket实例,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2019-08-08
在Python中,__repr__是一个特殊的内置类属性,用于定义类的字符串表示形式,本文将深入探讨__repr__的作用、用法以及一些实际应用场景,希望对大家有所帮助2023-12-12
这篇文章主要介绍了python判断字符串的前两个字母是否是”id",使用 Python 的字符串切片来判断一个字符串的前两个字母是否是 "id"，本文结合示例代码给大家介绍的非常详细，需要的朋友可以参考下2023-04-04
双星号(**)和单星号(*)在参数传递中扮演着关键角色,本文主要介绍了Python参数传递中双星号(**)和单星号(*)是作用,具有一定的参考价值,感兴趣的可以了解一下2024-05-05
最新评论