文章目录
- 1. 下载Netron
- 2. pth/pt转onnx
- 3. 可视化实现
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软件中,即可看到可视化后的模型图了
如下所示: