Netron可视化深度学习网络结构

有时候，我们构建网络模型想要直观的查看网络详细结构图，但是苦于没有办法。但是有了Netron以后，我们就可以将对应的onnx模型直接可视化，这样不仅可以观察网络的详细结构图，还可以查看网络每一层的具体参数，可谓是方便至极，下面我来简单分享一下。

一、介绍

Netron 是一个开源的模型可视化工具，用于可视化深度学习模型的结构和参数。它可以加载和显示多种框架和模型格式，包括ONNX（Open Neural Network Exchange）、TensorFlow、Keras、Caffe、Core ML 等。通过图形界面，用户可以直观地查看模型的网络结构、层级关系、参数等信息。

Netron 的主要特点包括：

1. 多框架和多格式支持：Netron 支持常见的深度学习框架和模型格式，包括ONNX、TensorFlow、Keras、Caffe、Core ML 等，因此可以加载和可视化多种类型的深度学习模型。
2. 直观的可视化界面：Netron 提供一个直观的图形界面，以树状结构展示模型的网络层次和参数。用户可以浏览模型的结构、查看各层的输入输出尺寸、参数数量等，帮助理解模型的组成和特性。
3. 跨平台支持：Netron 可以运行在多种操作系统上，包括 Windows、macOS 和 Linux，方便用户在不同环境中使用。
4. 快速加载和渲染：Netron 的设计优化了模型的加载和渲染过程，使得大型模型的可视化也能够在短时间内完成。

Netron 是一个简单但功能强大的模型可视化工具，对于深度学习开发者和研究人员来说，它可以提供有价值的模型分析和理解工具，帮助他们更好地理解和调试深度学习模型。

二、使用方法

2.1 onnx 格式模型可视化

我们可以将 pytorch 模型导出为 onnx 模型，使用 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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

2.2 代码直接运行可视化

由于 Netron 不支持直接将模型训练好的 pth 权重文件进行可视化，所以我们需要将现有模型转化为 onnx 格式的模型。这种方法需要我们安装下 netron 库，即使用命令 pip install netron。
转化代码如下：

# 针对有网络模型，但还没有训练保存 .pth 文件的情况
import netron
import torch.onnx
from torch.autograd import Variable
import torchvision.models as models

model = models.resnet34()
x = torch.randn(1, 3, 224, 224)  # 随机生成一个输入
onnx_path = "model.onnx"  # 定义模型数据保存的路径

torch.onnx.export(model, x, onnx_path )  # 将 pytorch 模型以 onnx 格式导出并保存
netron.start(onnx_path)  # 输出网络结构
1
2
3
4
5
6
7
8
9
10
11
12

执行上面代码后，会调用本地浏览器打开，形式和tensorboard差不多。

点击网址以后，就会出现下面的可视化界面。

但是我发现第二种方法和第一种方法有一个差异地方，那就是会有 identity 模块。

不知道是由于什么原因，我看其它人貌似并没有这种情况发生，有大佬知道的话可以评论区指教一下，谢谢～