Netron 神经网络可视化教程

本文介绍如何利用 Netron 进行深度学习模型的可视化。其优点在于适用性广，支持大部分深度学习框架的模型导出文件，本文以 Pytorch 为例讲解 Netron 具体的使用流程。

Netron 介绍

Netron 是一个用于深度学习模型可视化的库，它支持以下格式的模型存储文件：

• ONNX格式（.onnx，.pb）
• Keras格式（.h5，.keras）
• CoreML格式（.mlmodel）
• TensorFlow Lite格式（.tflite）

需要注意的是，Netron 并不支持 PyTorch 通过 torch.save 方法导出的模型文件。因此，在 PyTorch 保存模型时，我们需要将其导出为 ONNX 格式的模型文件，可以使用 torch.onnx 模块来实现这一目标。

总体而言，使用Netron进行深度学习模型可视化的流程分为两步。第一步是将PyTorch模型导出为ONNX格式的模型文件。第二步是使用Netron加载模型文件并进行可视化。接下来，我将按照这两个步骤来进行实践演示。

PyTorch模型导出为ONNX格式

要将PyTorch模型导出为ONNX格式的模型文件，我们只需要使用torch.onnx.export方法，并传入一个数据样例即可：

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

导出的模型文件为 model.onnx。

使用Netron加载模型文件

Netron 提供了两种方式来运行。第一种是通过安装 Netron 软件，并打开软件来加载模型文件。你可以在 Netron 的 Github 主页找到软件的下载地址：https://github.com/lutzroeder/netron。

第二种方式是将Netron作为Python库进行安装，并在Python代码中调用Netron库来加载模型并进行可视化。你可以使用pip install netron来进行安装。

如果你既不想安装Netron软件，也不想安装Netron库，那么Netron的作者还提供了一个在线演示网站，你可以直接上传模型文件并查看可视化结果。网站链接是：https://netron.app/。网站页面如下所示：

将刚才的 model.onnx 进行上传，结果如下

点击每个模块，你还可以在右侧查看模块的详细信息，例如卷积核的大小和具体参数等。不过，如果你仔细观察结果，你会发现图中并没有显示特征图的维度，只有输入数据的维度（3，256，256）。在 Netron 中，如果想要看到特征图的维度，我们需要在导出为 ONNX 模型时，同时添加特征图维度信息。这个操作需要使用 ONNX 库的帮助，你可以使用 pip install onnx 来安装 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. Netron 神经网络可视化教程