Netron展示pytorch模型结构_netron显示pytorch模型

第一步：需要将pytorch模型转化为.ONNX格式，转化程序如下：

model_path = './best_model/best_model.pth'
Model = YL_Effientnat(7 + 1, False)
Model.load_state_dict(torch.load(model_path,map_location=lambda storage,loc:storage))
Model.train(False)
input= torch.randn((1,3,80,80))
torch.onnx.export(Model,input,"model.onnx",verbose=False)

第二步：安装Netron

1、通过pip install netron，安装netron库；然后，通过netron库加载pytorch模型的.onnx文件进行模型网络结构展示，代码如下：

import netron
modelpath = 'model.onnx'
netron.start(modelpath)

2、下载netron.exe可执行程序，进行安装（可执行程序可在我上传的资源中找的）。

然后，打开.onnx模型文件，即可显示如下图。