【Pytorch模型可视化】Netron可视化pytorch框架下的模型

---

1. 下载Netron

github 下载链接：https://github.com/lutzroeder/Netron

打开上面网页，显示界面中的 Windows：Download 处点击直接下载 .exe 可执行文件。

如果觉得下载麻烦，也可以使用网页版的Netron：https://netron.app

2. pth/pt转onnx

由于Netron对pytorch模型的权重格式(pt/pth)不友好，因此要转成onnx格式
代码如下：

import torch
from model import Model   # 这里将自己的模型导入该处
 
pytorch_net_path = 'path1/to/model.pth'   # 原来模型保存的权重路径
onnx_net_path = 'path2/to/net.onnx'       # 设置onnx模型保存的权重路径
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )
 
# 权重导入模型
net = Model().to(device)
net.load_state_dict(torch.load(pytorch_net_path, map_location=device))
net.eval()
 
input = torch.randn(1, 3, 512, 512).to(device)   # (B,C,H,W)  其中Batch必须为1，因为test时一般为1，尺寸 H,W 必须和训练时的尺寸一致
torch.onnx.export(net, input, onnx_net_path, verbose=False)

运行结束后，会在 设置的onnx模型保存的权重路径文件夹 中看到net.onnx文件。

3. 可视化实现

将上述步骤的net.onnx文件拖入到Netron软件中，即可看到可视化后的模型图了
如下所示：