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)

在这里插入图片描述

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)  # 输出网络结构

执行上面代码后,会调用本地浏览器打开,形式和tensorboard差不多。
在这里插入图片描述
点击网址以后,就会出现下面的可视化界面。
在这里插入图片描述
但是我发现第二种方法和第一种方法有一个差异地方,那就是会有 identity 模块。
在这里插入图片描述
不知道是由于什么原因,我看其它人貌似并没有这种情况发生,有大佬知道的话可以评论区指教一下,谢谢~

猜你喜欢

转载自blog.csdn.net/m0_63007797/article/details/133578327