Open Neural Network Exchange(ONNX,开放神经网络交换)格式,是一个用于表示深度学习模型的标准,可使模型在不同框架之间进行转移。
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')
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)
"加载load"
model=onnx.load('net.onnx')
# 检查模型格式是否完整及正确
onnx.checker.check_model(model)
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))
input=model.graph.input
output = model.graph.output
"""输入输出层"""
print(input,output)
"""推断"""
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)
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