onnx模型转化与使用

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

1、生成pth文件

import torch.nn as nn
import torch
#
#
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layers1=nn.Sequential(
            nn.Conv2d(3,64,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64,128,3,1,1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128,256,3,1,1),
            nn.ReLU()

        )
        self.layers2=nn.Sequential(
            nn.Linear(256 * 7 * 7, 4096),
            nn.ReLU(),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Linear(4096, 10)
        )
    def forward(self,x):
        x=self.layers1(x)
        x=x.view(x.size()[0],-1)
        x=self.layers2(x)
        return x
#
inputs=torch.randn(1,3,28,28)
model=Net()
pred=model(inputs)
print(pred)
torch.save(model,'./net.pt')

2、pth文件转onnx文件

import onnx
import onnxruntime
# import torch

model=torch.load('net.pt')
inputs=torch.randn(1,3,28,28)
model.eval()
#
"""
导出export:pt->onnx
"""

inputname=['input','conv1','relu1','pool1','conv2','relu2','pool2','conv3','relu3','linear1','linear2','linear3']
# outname=['output%d'%i for i in range(1,11) ]
outname=['output']
torch.onnx.export(model,inputs,'./net.onnx',input_names=inputname,output_names=outname)

3、onnx文件操作

3.1 加载onnx文件

"加载load"
model=onnx.load('net.onnx')
# 检查模型格式是否完整及正确
onnx.checker.check_model(model)

3.2 打印onnx模型文件信息

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.3 获取onnx模型输入输出层

input=model.graph.input
output = model.graph.output
"""输入输出层"""
print(input,output)

3.4 推断结果

"""推断"""
session=onnxruntime.InferenceSession('net.onnx')
input_name = session.get_inputs()
print(input_name)
output_name=session.get_outputs()[0].name
res=session.run([output_name],{input_name[0].name:inputs.numpy()})
print(res)

3.5 修改

 

 

 

 

 

 

 

 

 

 

 

https://blog.csdn.net/qq_40263477/article/details/112802084?utm_medium=distribute.pc_relevant.none-task-blog-baidujs_utm_term-11&spm=1001.2101.3001.4242

https://blog.csdn.net/luohenyj/article/details/110940031?utm_medium=distribute.pc_relevant.none-task-blog-baidujs_title-3&spm=1001.2101.3001.4242