【netron】模型可视化工具netron

1、简介

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


2、使用方式

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

1)使用在线版

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

【netron】模型可视化工具netron_第1张图片

2)下载本地版

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


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

(1)支持的框架及对应的文件
【netron】模型可视化工具netron_第2张图片

(2)实验式支持,可能不太稳定
【netron】模型可视化工具netron_第3张图片

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

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


4、举例

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

【netron】模型可视化工具netron_第4张图片
可以同过 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")

【netron】模型可视化工具netron_第5张图片

你可能感兴趣的:(pytorch,深度学习,python,计算机视觉,目标检测)