Netron 可视化

1、pytroch模型转化为onnx模型

netron并不支持pytorch通过torch.save方法导出的模型文件，因此在pytorch保存模型的时候，需要将其导出为onnx格式的模型文件

pip install netron

pip install onnx

import torch
import torch.nn
from lib.models import model_factory
import onnx
import onnx.utils
import onnx.version_converter

device = torch.device('cpu')

# 定义网络
net = model_factory['bisenetv1'](19, aux_mode="pred")

net.load_state_dict(torch.load("model_final_v1_city.pth", map_location='cpu'), strict=False)
net.eval()

input_names = ['input']
output_names = ['output']

# 定义输入
x = torch.randn(1, 3, 1024, 2048, device=device)

# 导出onnx
torch.onnx.export(net, x, 'bisenetv1.onnx', input_names=input_names, output_names=output_names, verbose='True')

# 最后三行目的是导入中间输出的维度信息
model_file = 'bisenetv1.onnx'                
onnx_model = onnx.load(model_file)
onnx.save(onnx.shape_inference.infer_shapes(onnx_model), model_file)

 2、可视化

# 使用netron进行模型可视化
import netron

onnx_path = "../bisenetv1.onnx"

# 启动模型
netron.start(onnx_path)