本文介绍从零开始搭建YoloV5模型并训练教程

从零开始搭建YoloV5模型并训练教程

本文由林大佬原创，转载请注明出处，来自腾讯、阿里等一线AI算法工程师组成的QQ交流群欢迎你的加入: 1037662480

YoloV5出来有很长一段时间了, 大家都只知道这个版本似乎比Yolov4更好, 但很少有人去深挖里面的代码, 也少有人真正的进行过对比 (倒是有一些人分析过一些, 但我感觉还不够深入). 其实深入YoloV5网络构建, 训练过程, 最简单的方法就是去看他们的代码.

但是问题来了, 当我们想深入代码去探究网络的构建方法的时候, 里面通过cfg文件一堆parse_model方法的网络构建方式, 让我都有点懵逼, 更别说新手了. 因为我就想写下这篇文章, 来教大家如何去从零构建一个YoloV5模型, 并训练. 之前有大佬分析过YoloV5的训练过程, 和其工程上的优化, 这篇文章就不侧重一点来讲, 我们这篇教程着重讲一下如如何构建YoloV5的模型结构.

事实上, 我有点避重就轻, 因为这部分应该是构建一个模型最简单的部分, 但其实也是最重要的最核心的一个部分, 考虑到很多社区的朋友经常问我YoloV5怎么改Backbone, 怎么改Head, 怎么修剪通道, 怎么部署等等, 相信看完这篇文章,你就有了答案. 总之, 看完这篇文章并转发点赞, 你就成为了深度学习大师!

写这篇文章也是机缘巧合, 无意间我发现有一位大佬开源了一个这么一个库:

https://github.com/jinfagang/nb

不知道作者为啥取这个名字,先叫做牛逼库(Neuralnetwork Builder??)吧.

这个库把一些常用的blocks整合在了一起, 比如CSP, RFB, SPP,等, 接口统一之后构建模型的方法就会很简单. 我们就用这个库, 根据Yolov5的模型结构来构建一个简单版本的Yolov5. 当你熟悉这个构建过程之后, 修改backbone等的操作就会很简单.

构建Yolov5的backbone

我们知道, Yolov5实际上分为s, m, l, xl 几个不同的尺寸,它的区别就在于channel数目和CSP里面的n的宽度不同. 先来看看原始的Yolov5里面的yml的定义:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37


# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, BottleneckCSP, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, BottleneckCSP, [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, BottleneckCSP, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, BottleneckCSP, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, BottleneckCSP, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  

可以看到, 在backbone里面, 我们可以用nb库的相关blocks来构建:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


# divid by
cd = 2 # 1 for l, 2 for s, 1.5 for m, 0.5 for xl
wd = 3

self.focus = Focus(ch, 64//cd)
self.conv1 = ConvBase(64//cd, 128//cd, 3, 2)
self.csp1 = SimBottleneckCSP(128//cd, 128//cd, n=3//wd)
self.conv2 = ConvBase(128//cd, 256//cd, 3, 2)
self.csp2 = SimBottleneckCSP(256//cd, 256//cd, n=9//wd)
self.conv3 = ConvBase(256//cd, 512//cd, 3, 2)
self.csp3 = SimBottleneckCSP(512//cd, 512//cd, n=9//wd)
self.conv4 = ConvBase(512//cd, 1024//cd, 3, 2)
self.spp = SPP(1024//cd, 1024//cd)
self.csp4 = SimBottleneckCSP(1024//cd, 1024//cd, n=3//wd, shortcut=False)

对照这张图 (图来自某CSDN大佬), 以及对应的原始的yaml配置文件, 可以知道大概就是这么个构建, 然后我们定义一个函数来做forward:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


def _build_backbone(self, x):
    x = self.focus(x)
    x = self.conv1(x)
    x = self.csp1(x)
    x_p3 = self.conv2(x)  # P3
    x = self.csp2(x_p3)
    x_p4 = self.conv3(x)  # P4
    x = self.csp3(x_p4)
    x_p5 = self.conv4(x)  # P5
    x = self.spp(x_p5)
    x = self.csp4(x)
    return x_p3, x_p4, x_p5, x

需要注意的是, 我们从backbone拿出来的, 实际上是不同层抽出来的featuremap, 然后后面会送入到PANet (改进版本的FPN) 进行head的操作.

这里就引入了Yolov5里面的两个很大的改进点:

Focus模块: 这个模型的作用其实就是把3通道的原始图片, 将宽和高的像素填充到channel里面去, 将channel直接扩展到64维, 有人要问了, 这么做的目的是什么? 这么还用问? 问就是保证信息不丢失啊! 这样改变channel的方法比你直接用conv来的好吧? 用conv是必然会带来信息损失或者计算的增加的, 而这个, 就直接就是矩阵的reshape操作罢了. 很多人跟我扯什么空间转换啦, 什么颜色相关性啦, 在我看来都是bullshit.
CSP模块: 这个模型是干啥的? 个人觉得和resblocks没啥区别, 但是不同的是 CSP 最开始分叉的地方是直接在输入除以二的,不是直接将输入分叉, 因此其实计算量会减少很多.

构建Yolov5的PANet

接着就是构建PANet了, 还是这张图:

我们先定义我们需要的blocks:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


# PANet
self.conv5 = ConvBase(1024//cd, 512//cd)
self.up1 = nn.Upsample(scale_factor=2)
self.csp5 = SimBottleneckCSP(1024//cd, 512//cd, n=3//wd, shortcut=False)

self.conv6 = ConvBase(512//cd, 256//cd)
self.up2 = nn.Upsample(scale_factor=2)
self.csp6 = SimBottleneckCSP(512//cd, 256//cd, n=3//wd, shortcut=False)

self.conv7 = ConvBase(256//cd, 256//cd, 3, 2)
self.csp7 = SimBottleneckCSP(512//cd, 512//cd, n=3//wd, shortcut=False)

self.conv8 = ConvBase(512//cd, 512//cd, 3, 2)
self.csp8 = SimBottleneckCSP(512//cd, 1024//cd, n=3//wd, shortcut=False)

self.detect = Detect(self.nc, self.anchors, [256//cd, 512//cd, 1024//cd])

请注意, 这些内容都写入 __init__ 函数里面, 完整的代码我会在文章末尾贴出, 麻烦大家看完点个赞再走.

用这些模块来构建Head:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


def _build_head(self, p3, p4, p5, feas):
    h_p5 = self.conv5(feas)  # head P5
    x = self.up1(h_p5)
    x_concat = torch.cat([x, p4], dim=1)
    x = self.csp5(x_concat)

    h_p4 = self.conv6(x)  # head P4
    x = self.up2(h_p4)
    x_concat = torch.cat([x, p3], dim=1)
    x_small = self.csp6(x_concat)

    x = self.conv7(x_small)
    x_concat = torch.cat([x, h_p4], dim=1)
    x_medium = self.csp7(x_concat)

    x = self.conv8(x_medium)
    x_concat = torch.cat([x, h_p5], dim=1)
    x_large = self.csp8(x)
    return x_small, x_medium, x_large