关于深度学习:pytorch模型转换为onnx并使用netron可视化

netron 是一个十分好用的网络结构可视化工具。

应用办法：

pip install netron

import netronnetron.start('my.pth')

然而netron对pytorch模型的反对还不成熟。本人试的成果是生成的模型图没有连线。

这里就有一个把 .pth 模型转化为 .onnx 模型。

Pytorch模型转onnx

import torchfrom model import Modelold_net_path = 'lenet.pth'new_net_path = 'lenet.onnx'device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )# 导入模型net = Model().to(device)net.load_state_dict(torch.load(old_net_path, map_location=device))net.eval()input = torch.randn(1, 1, 30, 30).to(device)   # BCHW  其中Batch必须为1，因为测试时个别为1，尺寸HW必须和训练时的尺寸统一torch.onnx.export(net, input, new_net_path, verbose=False)

可视化onnx模型

import netronnetron.start('lenet.onnx')

效果图：