Netron visualizes deep learning network structure

Sometimes, when we build a network model, we want to intuitively view the detailed structure diagram of the network, but there is no way. But with Netron, we can directly visualize the corresponding onnx model. In this way, we can not only observe the detailed structure diagram of the network, but also view the specific parameters of each layer of the network. It is extremely convenient. Let me briefly share it below.

1. Introduction

Netron is an open source model visualization tool used to visualize the structure and parameters of deep learning models. It can load and display a variety of frameworks and model formats, including ONNX (Open Neural Network Exchange), TensorFlow, Keras, Caffe, Core ML, etc. Through the graphical interface, users can intuitively view the network structure, hierarchical relationships, parameters and other information of the model.

Netron’s key features include:

1. Multi-framework and multi-format support: Netron supports common deep learning frameworks and model formats, including ONNX, TensorFlow, Keras, Caffe, Core ML, etc., so multiple types of deep learning models can be loaded and visualized.
2. Intuitive visual interface: Netron provides an intuitive graphical interface that displays the network level and parameters of the model in a tree structure. Users can browse the structure of the model, view the input and output sizes of each layer, the number of parameters, etc., to help understand the composition and characteristics of the model.
3. Cross-platform support: Netron can run on a variety of operating systems, including Windows, macOS and Linux, making it convenient for users to use it in different environments.
4. Fast loading and rendering: Netron's design optimizes the model loading and rendering process, allowing the visualization of large models to be completed in a short time.

Netron is a simple but powerful model visualization tool that provides valuable model analysis and understanding tools for deep learning developers and researchers, helping them better understand and debug deep learning models.

2. How to use

2.1 onnx format model visualization

We can export the pytorch model as an onnx model and use the netron website to open it for visualization. The specific operations are as follows:

onnx model conversion code:

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 Code runs directly for visualization

Since Netron does not support direct visualization of the pth weight file trained by the model, we need to convert the existing model into an onnx format model. This method requires us to install the netron library, that is, use the command pip install netron.
The conversion code is as follows:

# 针对有网络模型，但还没有训练保存 .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)  # 输出网络结构

After executing the above code, the local browser will be called to open it, which is similar to tensorboard.

After clicking on the URL, the following visual interface will appear.

But I found one difference between the second method and the first method, that is, there will be an identity module.

I don’t know what the reason is. I see that this doesn’t seem to happen to other people. If anyone knows, you can give me some advice in the comment area. Thank you~