Pytorch模型转onnx，模型可视化

1、保存模型

torch.save(model,"hr1.pth")
# 注意是model，并不是model.state_dict()

2、转onnx

import torch.nn


model = torch.load('hr1.pth')
model.eval()

input_names = ['input']
output_names = ['output']
# 自己起名字

x = torch.randn(1,3,512,512,requires_grad=True)
# 这里要把握住网络的输入大小，如果模型是在gpu上进行训练，则将x变为
# x = torch.randn(1,3,128,128,requires_grad=True,device="cuda")

torch.onnx.export(model, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')

3、网络可视化

import netron
modelPath = "best.onnx"
netron.start(modelPath)
# 先安装netron模块（pip就可以）
# 这里加载的模型可以是torch，也可以是onnx

4、
Pytorch模型可视化结果

onnx模型可视化（只是一部分）

从上图中可以看出onnx模型更加详细，所以一般要转化成onnx