深度学习可视化工具：Netron

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,
)

1. 添加特征图维度

如果仔细看下结果，可以发现图中没有特征图的维度，只有输入数据的维度 $(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)

相比之前代码，多了一个增加维度信息的步骤。此时可视化图中就能完整显示所有特征图的维度了。

2. 显示 BN 层

可以看见上图中并没有显示 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)

再次查看网络结构：