当我们使用下列代码导出onnx模型,使用Netron查看的结果是这样的:

import torch
import torch.onnx
import torchvision.models as models
from collections import OrderedDict

device = torch.device("cpu")

def convert():
    model = models.inception_v3(pretrained = False)
    pthfile = './inception_v3_google-1a9a5a14.pth'
    state_dict =torch.load(str(pthfile))
    new_state_dict = OrderedDict()
    
    for k, v in state_dict.items():
        if(k[0:7] == "module."):
          name = k[7:]
        else:
          name = k[0:]        
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)
    
    
    print(model)

    input_names = ["actual_input_1"]
    output_names = ["output1"]
    dummy_input = torch.randn(16, 3, 224, 224)
    dynamic_axes = {'actual_input_1': {0: '-1'}, 'output1': {0: '-1'}}

    torch.onnx.export(model, dummy_input, "inception_v3.onnx", input_names = input_names, output_names = output_names, dynamic_axes = dynamic_axes, opset_version=11)
    
if __name__ == "__main__":
    convert()

onnx图
但是有些时候我们想要查看算子输出的shape结果,显然我们没有办法从上面的图中查看。那么这时候我们就需要onnx库中的shape_inference来进行推导,代码如下:

import onnx
from onnx import shape_inference
path = "..." #the path of your onnx model
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(path)), path)

使用Netron查看保存的模型结果:
在这里插入图片描述
顺利达到我们想要的效果!!!

GitHub 加速计划 / ne / netron
27.07 K
2.72 K
下载
lutzroeder/netron: 是一个用于查看和可视化神经网络模型的桌面应用程序,支持多种深度学习框架和常用的神经网络格式。适合用于可视化神经网络模型,尤其是对于需要进行神经网络模型调试和可视化的开发人员和研究人员。
最近提交(Master分支:1 个月前 )
9cee7758 - 3 个月前
63a272e8 - 3 个月前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐