onnx模型转化与使用

た入场券 2022-10-27 05:19 299阅读 0赞

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_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][https_blog.csdn.net_luohenyj_article_details_110940031_utm_medium_distribute.pc_relevant.none-task-blog-baidujs_title-3_spm_1001.2101.3001.4242]

[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/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]: https://blog.csdn.net/luohenyj/article/details/110940031?utm_medium=distribute.pc_relevant.none-task-blog-baidujs_title-3&spm=1001.2101.3001.4242