安装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')