netron并不支持pytorch通过torch.save方法导出的模型文件,因此在pytorch保存模型的时候,需要将其导出为onnx格式的模型文件,可以利用torch.onnx模块实现这一目标。
整体的流程分为两步,第一步,pytorch导出onnx格式的模型文件。第二步,netron载入模型文件,进行可视化。接下来按照这两步进行实践。
一、pytorch导出onnx格式的模型文件
1、在cmd中下载onnx
pip install onnx
2、在代码中加入如下代码:
第一种:
#定义了要保存的ONNX模型的路径和名称
onnx_path = "myobjectSSDModel.onnx"
#括号中第一个是模型名字,第二个是模型的输入,将模型和输入数据导出为ONNX格式,并保存到定义的路径。
torch.onnx.export(mymodeSSD, myinputSSD, onnx_path)
#从定义的路径加载ONNX模型。
onnx_model = onnx.load(onnx_path)
#对ONNX模型进行形状推理,并保存推理后的模型到原路径。
onnx.save(onnx.shape_inference.infer_shapes(onnx_model), onnx_path)
#使用Netron来打开和显示ONNX模型的可视化界面。
netron.start(onnx_path)
第二种:
onnx = torch.onnx.export(net, x, 'test.onnx')
netron.start(onnx)
二、netron载入模型文件
运行后,点击生成页面中的地址,或直接打开网页netron
进入如下页面,点击open model
在文件夹中找到你导出的onnx模式的模型文件,和代码文件存储的地方相同
打开即可看到可视化结果