模型可视化 netron
目录
- 一、代码
- 二、netron
一、代码
可以先从yolov5s.yaml中可以看到每一层的结构
# parameters
nc: 20 # number of classes
depth_multiple: 0.33 # model depth multiple 控制模型的深度（BottleneckCSP个数）
width_multiple: 0.50 # layer channel multiple 控制Conv通道channel个数（卷积核数量）
# depth_multiple表示BottleneckCSP模块的缩放因子，将所有BottleneckCSP模块的Bottleneck乘上该参数得到最终个数。
# width_multiple表示卷积通道的缩放因子，就是将配置里面的backbone和head部分有关Conv通道的设置，全部乘以该系数。
# 通过这两个参数就可以实现不同复杂度的模型设计。
# anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8 wh stride=8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
# YOLOv5 backbone
backbone:
# [from, number, module, args]
# from表示当前模块的输入来自那一层的输出，-1表示来自上一层的输出
# number表示本模块重复的次数，1表示只有一个，3表示重复3次
# module: 模块名
# args:
[[-1, 1, Focus, [64, 3]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4 128个3x3的卷积核 stride=2
[-1, 3, C3, [128]], # CSP Bottleneck with 3 convolutions
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 9, C3, [256]],
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3, [512]],
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 1, SPP, [1024, [5, 9, 13]]],
[-1, 3, C3, [1024, False]], # 9
]
# YOLOv5 head
# 作者没有区分neck模块，所以head部分包含了PANet+Detect部分
head:
[[-1, 1, Conv, [512, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 6], 1, Concat, [1]], # cat backbone P4
[-1, 3, C3, [512, False]], # 13
[-1, 1, Conv, [256, 1, 1]],
[-1, 1, nn.Upsample, [None, 2, 'nearest']],
[[-1, 4], 1, Concat, [1]], # cat backbone P3
[-1, 3, C3, [256, False]], # 17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]],
[[-1, 14], 1, Concat, [1]], # cat head P4
[-1, 3, C3, [512, False]], # 20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]],
[[-1, 10], 1, Concat, [1]], # cat head P5
[-1, 3, C3, [1024, False]], # 23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
]
或者使用命名输出：
YOLOv5 2021-6-29 torch 1.7.0+cpu CPU
from n params module arguments
0 -1 1 3520 models.common.Focus [3, 32, 3]
1 -1 1 18560 models.common.Conv [32, 64, 3, 2]
2 -1 1 18816 models.common.C3 [64, 64, 1]
3 -1 1 73984 models.common.Conv [64, 128, 3, 2]
4 -1 1 156928 models.common.C3 [128, 128, 3]
5 -1 1 295424 models.common.Conv [128, 256, 3, 2]
6 -1 1 625152 models.common.C3 [256, 256, 3]
7 -1 1 1180672 models.common.Conv [256, 512, 3, 2]
8 -1 1 656896 models.common.SPP [512, 512, [5, 9, 13]]
9 -1 1 1182720 models.common.C3 [512, 512, 1, False]
10 -1 1 131584 models.common.Conv [512, 256, 1, 1]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
12 [-1, 6] 1 0 models.common.Concat [1]
13 -1 1 361984 models.common.C3 [512, 256, 1, False]
14 -1 1 33024 models.common.Conv [256, 128, 1, 1]
15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
16 [-1, 4] 1 0 models.common.Concat [1]
17 -1 1 90880 models.common.C3 [256, 128, 1, False]
18 -1 1 147712 models.common.Conv [128, 128, 3, 2]
19 [-1, 14] 1 0 models.common.Concat [1]
20 -1 1 296448 models.common.C3 [256, 256, 1, False]
21 -1 1 590336 models.common.Conv [256, 256, 3, 2]
22 [-1, 10] 1 0 models.common.Concat [1]
23 -1 1 1182720 models.common.C3 [512, 512, 1, False]
24 [17, 20, 23] 1 67425 Detect [20, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
Model Summary: 283 layers, 7114785 parameters, 7114785 gradients, 16.5 GFLOPs
二、netron
注意几点：
1、不能直接打开官方给的yolov5s.pt，很多细节都是没用的，可以试试。
2、不能直接用export.py文件转为onnx文件，很多细节是错的，可以试试。
3、至于使用tensorboard之类的打开网络结构，一般都很不直观。
之前的代码上观看可能并不直观，如果想要直观的看到每一层的结构，我建议使用netron这个工具。
但是要想使用netron之前还需要生成自己的pt文件，直接使用官网的yolov5s.pt是看不清的。生成代码如下：
import torch
from models.yolo import Model
cfg = "models/yolov5s.yaml"
# Create model
model = Model(cfg).to("cpu")
x = torch.randn(1, 3, 640, 640).to("cpu")
script_model = torch.jit.trace(model, x)
script_model.save("weights/m.pt")
再使用netron软件或者网站打开m.pt就可以看到yolov5s的网络结构了，虽然还是有点不好看，但是以及可视化的很不错了，大家可以试试。