有时候为了更好的分析现有模型结构,可能需要对模型进行可视化。这里推荐一个在线的可视化图片生成工具 —— netron。通过网页端简单上传 onnx 模型就能生成对应模型的可视化图片,具体细节精确到卷积核大小。
网站链接:https://netron.app/
网站页面:
模型可视化图片局部图:
具体模型可视化的网页操作很简单,接下来的部分介绍一下 pytorch 模型到 onnx 模型的转换方法。
Pytorch 模型到 onnx 模型转换方法
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 保存的整个模型,使用以下方式加载
model = torch.load("test.pth") # pytorch模型加载
'''
# 保存的模型参数使用以下方式加载
torch_model = torch.load("test.pth") # pytorch模型加载
model = models.resnet50()
model.fc = torch.nn.Linear(2048, 4)
model.load_state_dict(torch_model)
'''
batch_size = 1 #批处理大小
input_shape = (3, 244, 384) #输入数据,改成自己的输入shape
# #set the model to inference mode
model.eval()
x = torch.randn(batch_size, *input_shape) # 生成张量
x = x.to(device)
export_onnx_file = "test.onnx" # 目的ONNX文件名
torch.onnx.export(model
x,
export_onnx_file,
opset_version=10,
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={"input":{0:"batch_size"}, # 批处理变量
"output":{0:"batch_size"}})
转换中常见的问题:
- UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported —— 可以将 opset_version 设置为 11