pytorch的pth模型转换为onnx模型,并使用netron可视化

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

pytorch的pth模型转换为onnx模型,并使用netron可视化_第1张图片

使用方法:

pip install netron
import netron
netron.start('my.pth')

但是netron对pytorch模型的支持还不成熟。自己试的效果是生成的模型图没有连线。

pytorch的pth模型转换为onnx模型,并使用netron可视化_第2张图片

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

Pytorch模型转onnx

import torch
from model import Model

old_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 netron
netron.start('lenet.onnx')

效果图:

pytorch的pth模型转换为onnx模型,并使用netron可视化_第3张图片


你可能感兴趣的:(pytorch的pth模型转换为onnx模型,并使用netron可视化)