Netron pth文件转onnx并显示出shape(size)的网络结构图

---

前言

提示：Netron pth文件转onnx并显示出shape(size)的网络结构图。

---

一、pth转onnx

提示：pth文件是pytorch训练后保存的网络模型

开放式神经网络交换（Open Neural Network Exchange）ONNX格式
所谓开放就是ONNX定义了一组和环境，平台均无关的标准格式，可使模型在不同框架之间进行转移。

换句话说，无论你使用何种训练框架训练模型（比如TensorFlow/Pytorch/OneFlow/Paddle），在训练完毕后你都可以将这些框架的模型统一转换为ONNX这种统一的格式进行存储。注意ONNX文件不仅仅存储了神经网络模型的权重，同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其它的辅助信息。

其实pth文件也可以直接用netron展示，但是pth文件展示模型网络结构有较多缺陷（因为模型通常只保存了网络参数，故展示时会变成下面这样没有网络结构）
eg

import torch
from torchvision.models import resnet50

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

modelpath=r"D:\Model\MA598_20220107.pth"
#pth模型路径
model=resnet50(pretrained=False)
#加载网络结构
w=torch.load(modelpath)
#torch加载模型
model.load_state_dict(w,strict=False)
#网络加载模型参数
model.to(device)
#模型放到device上
model.eval()
#将模型置于评估模式 不启用 BatchNormalization 和 Dropout,不改变权值
inputs = torch.ones((1, 1, 768,768)).to(device)
#网络输入size 同样放到device上
onnxpath=r"D:\Model\MA598_20220107.onnx"
#onnx格式的模型保存路径
torch.onnx.export(model,inputs,onnxpath,export_params=True,verbose=True,input_names=['input'],output_names=['output'],opset_version=12)
#模型格式转换

torch.onnx.export的各参数解析链接
onnx格式模型图片

有了网络结构，但是网络各层的shape没有

二、onnx格式给网络结构增加shape功能

代码如下（示例）：

import onnx
from onnx import shape_inference

model = r'D:\Model\MA598_20220107.onnx'
#上一步保存好的onnx格式的模型路径
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)
#增加节点的shape信息

增加节点shape信息的onnx格式模型在netron中的结构显示

Netron中显示

import netron
modelpath=r"D:\Model\MA598_20220107.onnx"
netron.start(modelpath)

总结

提示：pytorch模型在netron中显示网络结构和节点shape的路径为：
.pth—>.onnx–>增加节点shape信息–>netron