本文介绍 pth 模型转为 onnx 模型、使用 onnx 模型进行可视化以及过程中可能出现的问题。

转换模型

前提

首先，你需要有一个自己的pytorch格式的模型，通常的后缀为.pth。
可以通过

save_path = './Mlp.pth'
torch.save(net.state_dict(), save_path)

进行保存。

可以通过

net.load_state_dict(torch.load('Mlp.pth'))

进行加载参数使用。

当然如果你只想可视化的话，不需要训练得到pth参数，只要有模型就可以。

转换方法

import torch.nn

model = MLP().cuda()	# 声明模型
model.load_state_dict(torch.load('Mlp.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')

如果你的模型需要传入两个参数的话，那就再声明一个 y 变量，和 x 一起以元组的方式 (x, y) 传入.export()方法中。

这样，你就得到了一个onnx模型。
过程中可能出现的问题见最后一章。

模型可视化

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

可视化会进入 http://localhost:8080/ 网址。
效果大致如下图：

可能出现的报错信息

ValueError: torch.nn.DataParallel is not supported by ONNX exporter, please use ‘attribute’ module to unwrap model from torch.nn.DataParallel. Try torch.onnx.export(model.module, …)

这是因为你的模型使用了 DataParallel 包装。
只需要按照报错信息修改为如下

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

即可。

RuntimeError: ONNX export failed on an operator with unrecognized namespace torchvision::roi_align. If you are trying to export a custom operator, make sure you registered it with the right domain and version.

修改为

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

即添加了 opset_version 参数。

类似的报错都可以去查看一下官方文档的版本信息。（查看方法和文档网址见参考资料一章）

参考资料

模型部署入门教程（三）：PyTorch 转 ONNX 详解
onnx/docs/Operators.md
Pytorch模型转onnx，模型可视化
UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9
onnx.export报警告：WARNING: The shape inference of prim::Constant type is missing…解决方法