安装netron

pip install netron

使用

netron对各种模型的支持可谓各有其优点和缺点，netron可以显示pytorch模型定义在主类中的各部分的整体结构，忽略细节。但是看不到整个模型的整体结构和细节。netron可以显示onnx模型的整体结构，但全部都是基于conv,relu,maxpool等基础层的细节。事实上这两种需求都是有的。

pytorch模型转onnx

import torch
import torch.onnx


def convert2onnx(model_path=r"models/resnet101.pth", export_onnx_path="models/onnx/resnet101.onnx"):
    model = torch.load(model_path)
    model.eval()

    input_shape = (3, 224, 224)  # 输入数据,改成自己的输入shape #renet
    example = torch.randn(1, *input_shape)  # 生成张量
    torch.onnx.export(model, example, export_onnx_path, verbose=True)

对pytorch模型可视化

import netron
netron.start('models/resnet101.pth')

对onnx模型可视化

import netron
netron.start('models/onnx/resnet101.onnx')