【netron】模型可视化工具netron

1、简介

在实际的项目中，经常会遇到各种网络模型，需要我们快速去了解网络结构。如果单纯的去看模型文件，脑海中很难直观的浮现网络的架构。这时，就可以使用netron可视化工具，可以清晰的看到每一层的输入输出，网络总体的架构，简单好用。

2、使用方式

这里介绍两种使用方式： 1、使用在线版； 2、下载本地版

1）使用在线版

浏览器访问：https://netron.app/
点击 “Open Model” 按钮，选择要可视化的模型文件即可

2）下载本地版

终端进行安装： pip install netron
安装完成后，在脚本中 调用包 import netron
运行程序 netron.start("model.onnx")， 会自动打开浏览器进行可视化 （最后有例子）

3、支持的网络框架 和 模型文件类型

（1）支持的框架及对应的文件

（2）实验式支持，可能不太稳定

我习惯用 pytorch，但是 netron 对 pytorch 的 .pt 和 .pth 文件不是很友好，所以，我都是先转换为 onnx 格式，再进行可视化，下面举例。

另外，netron 可以直接可视化 yolo （DarkNet 框架）的 .cfg文件，很是方便。这篇文章有我可视化出来的 yolov3-spp 网络结构。

4、举例

一般情况下，netron 只展示最初的输入尺寸 和 最后的输出尺寸，中间层的尺寸都是不展示的（如下）。

可以同过 onnx.save(onnx.shape_inference.infer_shapes(onnx.load("model.onnx")), "model.onnx") 进行处理。这样中间的每一层的输入输出就都会推理出 并可视化出来了。 （谁有简便方法，评论区留言）

import torch
import torch.nn as nn
import netron
import onnx
from onnx import shape_inference


class My_Net(nn.Module):
    def __init__(self):
        super(My_Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(16),
            nn.LeakyReLU(),
        )

        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=1, bias=False),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(),
        )

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x


net = My_Net()
img = torch.rand((1, 3, 224, 224))
torch.onnx.export(model=net, args=img, f='model.onnx', input_names=['image'], output_names=['feature_map'])
onnx.save(onnx.shape_inference.infer_shapes(onnx.load("model.onnx")), "model.onnx")
netron.start("model.onnx")