基于 onnx 的 pytorch 模型结构可视化

有时候为了更好的分析现有模型结构，可能需要对模型进行可视化。这里推荐一个在线的可视化图片生成工具 —— 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"}})

转换中常见的问题：

1. UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported —— 可以将 opset_version 设置为 11