Pytorch与Onnx的转换与推理

Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移。

一、pytorch模型保存/加载

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

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

模型保存与调用方式一(只保存权重):

保存:

torch.save(model.state_dict(), mymodel.pth)#只保存模型权重参数,不保存模型结构

调用:

model = My_model(*args, **kwargs)  #这里需要重新创建模型,My_model
model.load_state_dict(torch.load(mymodel.pth))#这里根据模型结构,导入存储的模型参数
model.eval()

模型保存与调用方式二(保存完整模型):

保存:

bash 复制代码
torch.save(model, mymodel.pth)#保存整个model的状态

调用:

bash 复制代码
model=torch.load(mymodel.pth)#这里已经不需要重构模型结构了,直接load就可以
model.eval()

.pt表示pytorch的模型,.onnx表示onnx的模型,后缀名为.pt, .pth, .pkl的pytorch模型文件之间其实没有任何区别

二、pytorch模型转ONNX模型

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

bash 复制代码
import torch
torch_model = torch.load("/home/pytorch/save.pth") # pytorch模型加载

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

x = torch.randn(1,3,320,640)        # 生成张量(模型输入格式)
export_onnx_file = "/home/pytorch/test.onnx"   # 目的ONNX文件名

// 导出export:pt->onnx
torch.onnx.export(torch_model,                    # pytorch模型
                    x,                            # 生成张量(模型输入格式)
                    export_onnx_file,            # 目的ONNX文件名
                    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、文件中只保留模型权重

bash 复制代码
import torch
torch_model = selfmodel()                      # 由研究员提供python.py文件

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

x = torch.randn(1,3,320,640)        # 生成张量(模型输入格式)
export_onnx_file = "/home/pytorch/test.onnx"   # 目的ONNX文件名

// 导出export:pt->onnx
torch.onnx.export(torch_model,                    # pytorch模型
                    x,                            # 生成张量(模型输入格式)
                    export_onnx_file,            # 目的ONNX文件名
                    do_constant_folding=True,    # 是否执行常量折叠优化
                    input_names=["input"],        # 输入名(可略)
                    output_names=["output"],    # 输出名(可略)
                    dynamic_axes={"input":{0:"batch_size"},        # 批处理变量(可略)
                                    "output":{0:"batch_size"}}) 

3、onnx文件操作

3.1 安装onnx,onnxruntime:

bash 复制代码
pip install onnx
pip install onnxruntime(只能用cpu)
pip install onnxruntime-gpu(gpu和cpu都能用)

首先要强调的是,有两个版本的onnxruntime,一个叫onnxruntime,只能使用cpu推理,另一个叫onnxruntime-gpu,既可以使用gpu,也可以使用cpu。

如果自己安装的是onnxruntime,需要卸载后安装gpu版本。

确认一下是否可以使用gpu

注意:

```python

print(onnxruntime.get_device())

```

上面的代码给出的输出是'GPU'时,并不代表就成功了。

而要用下面的代码来验证:

```python

ort_session = onnxruntime.InferenceSession("path/model/model_name.onnx",

providers=['CUDAExecutionProvider'])

print(ort_session.get_providers())

```

当输出是:['CUDAExecutionProvider', 'CPUExecutionProvider']才表示成功了。

版本查询:NVIDIA - CUDA | onnxruntime

安装固定版本的onnxruntime:

pip install onnxruntime-gpu==1.9.0

卸载pip uninstall

3.2 加载onnx文件

bash 复制代码
# "加载load"
model=onnx.load('net.onnx')

检查模型格式是否完整及正确

bash 复制代码
onnx.checker.check_model(model)

3.3 打印onnx模型文件信息

bash 复制代码
session=onnxruntime.InferenceSession('net.onnx')
inp=session.get_inputs()[0]


#conv1=session.get_inputs()['conv1']
#out1=session.get_outputs()[1]
out=session.get_provider_options()
#print(inp,conv1,out1)
print(inp)
#print(out)
"打印图信息:字符串信息"
graph=onnx.helper.printable_graph(model.graph)
print(type(graph))

3.4 获取onnx模型输入输出层

bash 复制代码
input=model.graph.input
output = model.graph.output
"""输入输出层"""
print(input,output)

3.5 推理过程

bash 复制代码
import onnx
import onnxruntime
import torch

inputs=torch.randn(1,3,640,320)
#上述inputs仅用于测试使用,用于图片推理,应该换成自己的图片,如:
#img_path='1.jpg'#图片尺寸与onnx模型的处理尺寸保持一致
#img=cv2.imread(img_path)
#inputs=preprocess_imgae(img)#标准化等预处理操作,与源项目代码保持一致即可
#print('inputs.size():',inputs.size())

model=onnx.load('/home/pytorch_DL/test_320_640.onnx')
onnx.checker.check_model(model)
session =onnxruntime.InferenceSession('/home/pytorch_DL/test_320_640.onnx',
providers['CUDAExecutionProvider','CPUExecutionProvider'])
print('session.get_providers():',session.get_providers())
input_name = session.get_inputs()
output_name=session.get_outputs()[0].name
res=session.run([output_name],{input_name[0].name:inputs.numpy()})

参考:Pytorch与Onnx模型的保存、转换与操作_onnx转pytorch_Yuezero_的博客-CSDN博客

pytorch 模型的保存与加载方法以及使用onnx模型部署推理 | 码农家园

onnxruntime使用gpu推理 - 知乎

相关推荐
说私域2 分钟前
基于开源 AI 智能名片、S2B2C 商城小程序的用户获取成本优化分析
人工智能·小程序
东胜物联22 分钟前
探寻5G工业网关市场,5G工业网关品牌解析
人工智能·嵌入式硬件·5g
南宫理的日知录31 分钟前
99、Python并发编程:多线程的问题、临界资源以及同步机制
开发语言·python·学习·编程学习
皓74133 分钟前
服饰电商行业知识管理的创新实践与知识中台的重要性
大数据·人工智能·科技·数据分析·零售
coberup40 分钟前
django Forbidden (403)错误解决方法
python·django·403错误
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读3】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer
龙哥说跨境1 小时前
如何利用指纹浏览器爬虫绕过Cloudflare的防护?
服务器·网络·python·网络爬虫
AltmanChan1 小时前
大语言模型安全威胁
人工智能·安全·语言模型
985小水博一枚呀1 小时前
【深度学习滑坡制图|论文解读2】基于融合CNN-Transformer网络和深度迁移学习的遥感影像滑坡制图方法
人工智能·深度学习·神经网络·cnn·transformer·迁移学习
数据与后端架构提升之路1 小时前
从神经元到神经网络:深度学习的进化之旅
人工智能·神经网络·学习