赞
踩
Netron是一个用于神经网络、深度学习和机器学习模型的可视化工具。
Netron支持 ONNX、TensorFlow Lite、Caffe、Keras、Darknet、PaddlePaddle、ncnn、MNN、Core ML、RKNN、MXNet、MindSpore Lite、TNN、Barracuda、Tengine、CNTK、TensorFlow.js、Caffe2 和 UFF。
它还实验性支持PyTorch、TensorFlow、TorchScript、OpenVINO、Torch、Vitis AI、kmodel、Arm NN、BigDL、Chainer、Deeplearning4j、MediaPipe、ML.NET 和 scikit-learn。
snap install netron
命令行输入 netron
即可打开界面,如下所示:
比如使用下面代码生成的一个模型:
import torchvision.models as models
import torch
# 定义样例数据+网络
data = torch.randn(2, 3, 256, 256)
net = models.resnet34()
# 导出为onnx格式
torch.onnx.export(
net,
data,
'model.onnx',
export_params=True,
opset_version=8,
)
如果仔细看下结果,可以发现图中没有特征图的维度,只有输入数据的维度 ( 2 , 3 , 256 , 256 ) (2, 3, 256, 256) (2,3,256,256)。在 Netron 中,如果想看到特征图的维度,需要在导出为 ONNX 的时候,同时加上特征图维度信息。比如像下面这段代码:
import torchvision.models as models import torch import onnx import onnx.utils import onnx.version_converter # 定义数据+网络 data = torch.randn(2, 3, 256, 256) net = models.resnet34() # 导出 torch.onnx.export( net, data, 'model.onnx', export_params=True, opset_version=8, ) # 增加维度信息 model_file = 'model.onnx' onnx_model = onnx.load(model_file) onnx.save(onnx.shape_inference.infer_shapes(onnx_model), model_file)
相比之前代码,多了一个增加维度信息的步骤。此时可视化图中就能完整显示所有特征图的维度了。
可以看见上图中并没有显示 BN 层。如果想将 Conv 和 BN 分开显示,可以在 torch.onnx.export
时添加参数 training=2
,否则 ONNX 默认将 BN 层融合到 Conv 层。
import torchvision.models as models import torch import onnx import onnx.utils import onnx.version_converter # 定义数据+网络 data = torch.randn(2, 3, 256, 256) net = models.resnet34() # 导出 torch.onnx.export( net, data, 'model.onnx', export_params=True, opset_version=8, training=2, //---->>>>>>改动处 ) # 增加维度信息 model_file = 'model.onnx' onnx_model = onnx.load(model_file) onnx.save(onnx.shape_inference.infer_shapes(onnx_model), model_file)
再次查看网络结构:
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。