模型可视化 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  models.common.C3 [128, 128, 3] 5 -1 1  models.common.Conv [128, 256, 3, 2] 6 -1 1  models.common.C3 [256, 256, 3] 7 -1 1  models.common.Conv [256, 512, 3, 2] 8 -1 1  models.common.SPP [512, 512, [5, 9, 13]] 9 -1 1  models.common.C3 [512, 512, 1, False] 10 -1 1  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  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  models.common.Conv [128, 128, 3, 2] 19 [-1, 14] 1 0 models.common.Concat [1] 20 -1 1  models.common.C3 [256, 256, 1, False] 21 -1 1  models.common.Conv [256, 256, 3, 2] 22 [-1, 10] 1 0 models.common.Concat [1] 23 -1 1  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,  parameters,  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的网络结构了，虽然还是有点不好看，但是以及可视化的很不错了，大家可以试试。