yolov4&v5训练代码理解

1 简介

本文依据https://github.com/Tianxiaomo/pytorch-YOLOv4对YoloV4的模型训练和验证代码进行全面解析，理论部分可以参考：
YoloV1&V2&V3：https://blog.csdn.net/cdknight_happy/article/details/91793142
YoloV4：https://blog.csdn.net/cdknight_happy/article/details/107883216
GIOU：https://blog.csdn.net/cdknight_happy/article/details/107861071
训练tricks：https://blog.csdn.net/cdknight_happy/article/details/105139999

yolov5的训练代码https://github.com/ultralytics/yolov5和本文重点解释的代码基本一致，区别在于yolov5中使用的模型结构相对简单，但基本是一致的。损失计算、模型推理及统计部分则完全一样。因为yolov5的代码持续在维护且在模型部署方面的解释很多，所以建议学习yolov5的训练代码。

本文分别从网络结构、训练数据增广、损失计算、训练tricks、模型验证五个方面进行整理，对应train.py和test.py两个主文件。

2 网络结构

models文件夹下提供了四个网络结构，分别是yolov4s-mish、yolov4m-mish、yolov4l-mish和yolov4x-mish，四者的基本组件是一致的，不同点仅为网络的深度和宽度不同：

depth_multiple: 1.0  # expand model depth
width_multiple: 1.0  # expand layer channels

这里选用最小的模型yolov4s-mish进行解析。

总体网络结构如下图所示，网络解析代码为yolo.py的parse_model函数：

2.1 Conv

卷积 + BN + 激活函数，激活函数用的是Mish

2.2 Bottleneck

输入数据$c_1$个channel，中间层的feature map有$c_{\_}$个channel，若$c_{\_}\le c_1$，就形成了Bottleneck的效果。

如果输出channel数$c_2$等于输入channel数$c_1$，且设置了进行add的操作，则进行逐点相加。否则直接输出第二个Conv的结果。

2.3 BottleneckCSP

CSPNet的论文请参考：https://blog.csdn.net/cdknight_happy/article/details/107962173

原始的含义是将输入x分为两部分，一部分进行transition操作，一部分直接后传，然后将两个结果concatenate之后再进行transition操作。这里实现的是一个类似的CSPNet，并没有对输入数据进行切分。

2.4 BottleneckCSP2

2.5 SPPCSP

SPPCSP和BottleneckCSP一样，是一个近似版本的CSP。

SPP的意思是对输入数据进行多个尺度的最大池化，然后将各池化结果进行concat操作，这里进行了kernel size等于5,9,13的池化。目的在于提取多个粒度的特征。

2.6 Detect

训练时，Detect模块的作用就是使用nn.Conv2d对输入数据进行卷积操作，使输出为3 * 85 = 255个通道。3表示每个网格有3个anchor box。85 = 4 + 1 + 80,4表示检测框相对于某anchor的位置信息，1表示包含目标的置信度，80表示coco数据集中80个类别的概率。对卷积的输出进行reshape操作，将其转换为(bs,3,h,w,85)的格式进行输出。

3 损失计算

3.1 anchors

按照YoloV3的思路，作者设计了三个尺度的目标检测，每个尺度中的每个cell又使用了三个anchor box，预设的anchor box大小为：

[[12,16],[19,36],[40,28]]
[[36,75],[76,55],[72,146]]
[[142,110],[192,243],[459,401]]

预设的anchors毕竟是人为估计的，作者使用了预设anchors和训练集中所有标注框的重合程度计算最佳召回率，如果最佳召回率小于阈值，证明预设anchors设计不合理，需要根据训练集聚类生成更加合适的anchors。

预设anchors的召回率计算：
对输入图像的宽高信息进行缩放，使各图像的长边均为640.读取各标注框的宽高，乘以前面的缩放系数，得到训练时各标注框的宽高信息。计算各标注框的宽高除以各anchor框宽高之后的比值，计算宽高两个比值中的较小值，然后计算每一个标注框和哪一个anchor的比值最接近，如果较接近的宽高比值大于了设定的阈值，则认为该标注框能够和对应的anchor进行匹配。求所有标注框是否匹配某个anchor的匹配均值，得到总体召回率。物理含义是标注框的宽高和某一个anchor的宽高比值中最接近的那个需要在一个范围内，比如0.25 到 4之间，如果预设的所有的anchor都和标注框的宽高差距很大，即预设anchor比实际目标高很多且宽很多或者是矮很多且窄很多，表明用这组预设的anchor可能无法准确的回归出目标的位置。

预设anchor不合适(预设anchors和训练集标注数据的最佳召回率小于0.99)时，对训练集的标注使用kmeans聚类获取更加合适的anchor：
计算过程也是先把图像的长边缩放到640，得到缩放后的各标注框的宽高，对这些宽高进行白化(使其方差为1)，然后进行kmeans操作，就是把宽高当成二维坐标系下的一个点，找到n个中心点，这n个中心点表示的是和训练集中所有标注框宽高总体最接近的框的宽高。对这些宽高进行随机的小范围的变化，计算每次变化后和训练集所有标注框之间的召回率，选取召回率最高的那组框的宽高作为拟合出的最佳anchors。

YoloV2中，作者提到利用标注框与anchor之间的IOU作为kmeans的度量标准，这里代码中却是按宽高的比值作为距离判断标准。这两种方式理论上都说的过去，后续可以都尝试下。

3.2 损失

训练过程中，目标可能出现在图像上的任何位置，训练的目的就是让每一个网格能够对其包含的或其临近网格包含的目标进行预测，让每一个网格能够判断其是否应该参与某个目标的预测的概率(置信度)、目标的位置与网格自身之间的偏移量及目标的宽高、目标的类别。

总的损失是分类损失+定位损失+置信度损失。

3.2.1 给每个target设置参考anchor

在标注数据中，每幅图像都标注了所有感兴趣目标的类别和位置信息，其中位置信息是归一化之后的xywh的值。在计算损失前，首先在utils.py的build_targets函数中，将各anchors和target进行关联，也就是指定让哪一个anchor负责某个目标的预测。首先就是保证各目标基于宽高差距不太大的anchor进行预测，所谓差距不太大，就是指目标的宽高除以anchor的宽高要位于[0.25,4]之内。选取基于宽高不太大的anchor进行某个目标的预测，意味着决定了每一个目标该由哪个尺度的anchor进行预测。

在决定了基于哪个anchor进行预测后，还要根据目标所在的位置，决定由哪个网格进行目标的预测。按照rect4规则，根据目标中心点的位置，中心点所在的网格要负责目标的预测；如果目标的中心点在当前网格中的位置偏左偏下，那么左边的网格和下面的网格也要负责目标的预测。左上、右上、右下也按照相同的规则决定该由哪些网格进行目标预测，这样有助于提高目标检测的召回率。如下图所示。

决定了网格也就得到了目标位置相对于其所在网格左上角的偏移量。

训练过程中，一个网格可能要对不同位置的目标进行预测，通过迭代，网格知道对于不同的输入图像应该给出不同的激活值，通过不停的迭代调整权重系数，网格能够对不同位置的目标都给出准确的预测结果，从而各个网格都有了较强的鲁棒性，也就能够对不同种类、不同位置的目标进行准确地检测。

3.2.2 损失函数

损失一共分为三部分，分别是判断是否存在目标的损失、目标定位损失和目标分类损失，使用的都是BCE loss(Binary Cross Entropy Loss)，对应的pyTorch函数有nn.BCELoss和nn.BCEWithLogitsLoss。

nn.BCELoss的作用是计算二分类的交叉熵损失，即计算预测概率和target之间的交叉熵损失，对应到公式上就是$l=-\log (p_{y_i})$，$y_i$表示某个样本的真实类别，$p_{y_i}$表示样本被预测为其真实类别的概率。

nn.BCEWithLogitsLoss = sigmoid + BCELoss，好处是nn.BCEWithLogitsLoss的数值稳定性更高。计算公式是:

3.3.3 定位损失

取出每幅图像包含目标的网格和对应anchor的预测结果，对目标位置的预测值进行sigmoid运算后乘以2再减去0.5得到box中心点的位置，再对目标宽高预测值进行sigmoid运算后乘以2再平方再乘以对应的anchor的宽高得到预测框的宽高值。乘以2是在yolov4中为了进行对靠近cell边界点的目标进行检测，如https://blog.csdn.net/cdknight_happy/article/details/107883216中4.2节的解释。

相当于这里模型学习的结果中，执行的计算公式是：
$$\begin{gathered} \Delta x=\sigma (t_x)\ast 2-0.5 \\ {\Delta }_y=\sigma (t_y)\ast 2-0.5 \\ {s}_w=(\sigma (t_w)\ast 2{)}^2 \\ {s}_h=(\sigma (t_h)\ast 2{)}^2 \end{gathered}$$

$t_x,t_y,t_w,t_h$是Detect模块的输出值，$\Delta x,\Delta y$是预测的目标中心点的位置相对于其网格坐标的偏移量，$s_w,s_h$是预测的目标框的宽高与对应的anchor的宽高的比值。

也就是说，训练过程中，输入图像经过backbone+head+Detect模块的卷积处理后得到的输出经过上面的公式处理后得到了$\Delta x,\Delta y,s_w,s_h$的值。

在计算预测框与真实框之间的定位损失时，使用了GIOU loss，如https://blog.csdn.net/cdknight_happy/article/details/107861071介绍的理论。

3.3.4 分类损失

分类损失就是计算目标属于各类别的预测概率和真实目标类别之间的BCE loss。

3.3.5 置信度损失

置信度损失，按照理论应该是存在目标的概率乘以预测框与真实框之间的IOU。各网格标注的置信度值(tobj)用的是1 - 权重系数 + 权重系数乘以GIOU的值，计算损失时用的也是BCE loss。置信度损失进行综合时，对于各个尺度的损失使用了不同的加权系数。

4 训练tricks

4.1 CmBN

这里的实现是设置nbs=64，只有当训练图像量达到64时才进行权重的更新，因此这里实现的更像是BN。

4.2 余弦退火

https://blog.csdn.net/cdknight_happy/article/details/109379102中的2.6节。

这里使用的余弦退火函数是：
$${\eta }_t=\frac{(1+\cos (t\ast \pi /epochs))}{2}\ast 0.8+0.2$$
$t$表示训练当前epoch值，epochs表示总的训练epoch数目。当$t=0$时，${\eta }_t=1.0$，当$t=epochs$时，${\eta }_t=0.2$。即训练过程中学习率按照余弦曲线从1.0减小到0.2。

4.3 modelEMA

对各个epoch的训练结果(模型参数)进行指数移动平均累加得到最终的模型。
参考：https://zhuanlan.zhihu.com/p/68748778

4.4 混合精度训练

参考https://blog.csdn.net/cdknight_happy/article/details/108238751

4.5 DDP

参考https://blog.csdn.net/cdknight_happy/article/details/108262595

https://github.com/ultralytics/yolov5/issues/475

4.6 SyncBN

参考https://blog.csdn.net/cdknight_happy/article/details/109366599

https://github.com/ultralytics/yolov5/issues/475

SyncBatchNorm could increase accuracy for multiple gpu training, however, it will slow down training by a significant factor. It is only available for Multiple GPU DistributedDataParallel training.

It is best used when the batch-size on each GPU is small (<= 8).

4.7 加权系数解决类别不均衡

训练样本集中各类别的目标出现的概率不一致，如果计算损失时统一对待，会造成模型训练结果受出现频率高的类别的影响更大，那么模型也就越容易检测高频出现的目标，越不擅长于处理低频出现的目标。为了解决该问题，可以计算各类别目标的出现频率，以其倒数作为类别的加权系数，这样在损失计算时起到平衡类别不均衡的目的，有助于提升模型对全体类别的处理效果。

左上角的图为coco数据集中各类别目标出现的次数；第二幅图为所有标注框中心点的位置分布；第三幅图为所有标注框的宽高分布。中心点和宽高的标注值都是进行了归一化处理的。

4.8 学习率warmup

学习率warmup的好处参考https://blog.csdn.net/cdknight_happy/article/details/105139999的3.1节。在网络训练初期，因为模型的权重系数都是随机初始化得到的，如果在训练初始就使用较大的学习率，有可能会造成网络参数沿着错误的方向移动拉不回来了，所以学习率warmup的含义是在训练的初始几个epoch内，以较小的学习率进行模型训练，等到几个epoch训练完成后，模型的参数已经找到了正确的前进方向，此时再使用大的学习率进行模型的训练可以取得更好的训练效果。

学习率warmup落实到实现上，一般会将学习率从0线性增加到初始学习率。当然这里的代码中也存在了学习率从0.1线性减小到初始学习率的情况，以及对momentum值也进行了线性增加，从0.9增加到了0.937.

4.9 多尺度训练

多尺度训练的好处是让网络可以处理多个尺度的输入目标，提高模型对目标尺度变化的适应性。这里的实现是每次训练中随机从$[320,352,\cdots ,992]$之间选取一个尺度值，将输入图像缩放到该尺度进行推理过程。因为前面记录图像上的标注框时是按照归一化的目标左上角坐标及目标宽高进行记录的，所以这里在输入图像缩放后无需调整标注框的位置。

5 数据处理

5.1 数据cache

为了避免频繁的去读取训练数据的label，作者设计了在第一次读取数据后进行cache的操作。作者以图片名为key值、以当前图像中所有目标的标注信息和图像的大小为value值构建了字典，然后使用torch.save保存了这个字典。

为了读取的时候进行校验，进行cache的时候生成了hash值，具体规则是使用os.path.getsize(f)计算了所有图像及其label文件所占的存储空间(以字节为单位)，将该存储空间值作为hash值一并cache到本地文件中。

读取cache的时候，首先使用torch.load加载硬盘上存储的cache文件，然后比较读取的hash值和当前数据集重新计算的hash值是否相同，如果相同则表明cache文件对应了当前训练数据集，无需重新读取。

数据cache对于大规模的数据集反复训练时，可以避免训练数据集中标注数据的重复读取，有助于提升训练的效率。

5.2 数据清洗

读取标注信息的cache后，遍历所有的标注信息，统计label为空、目标重复标注的图像的数量。统计完成后可以视数据的干净程度决定是否进行数据清洗。

5.3 训练图像增广

mosaic数据增广：
mosaic数据增广是将四副图像类似于四个象限一样拼接到一起，但原点坐标不一定是在大图的中心位置。以这里图像大小为640为例，拼接后的大图像大小为1280 * 1280，四副图像的中心坐标介于320到960之间。随机选取四副图像，将其分别放于大图的四个象限位置，有可能某副小图填不满对应的象限，也有可能象限内只能放下小图的一部分。

对各图像上的原始标注框也进行相似的处理，将小图上的标注框变换为其在大图上的位置。最后生成的图像如下图所示：

Mosaic数据增广之后，还需要对图像和标注框进行随机的旋转、缩放、平移和裁剪，对越界的边界框进行处理，完成mosaic数据增广。

传统的数据增广：
这里应用的传统数据增广包括 HSV空间的图像颜色变换、随机的水平翻转和竖直翻转，最后将label改为归一化的图像左上角坐标及图像宽高信息，BGR转RGB，改为CHW格式，设置为连续内存，转换为tensor后返回图像的变换结果。

6 推理及性能统计

6.1 推理

加载已训练的模型，如果使用半精度浮点数进行推理，则调用nn.Module.half()将模型转换为半精度，同样读取输入图像，使用Tensor.half()将输入图像转换为fp16，执行前向推理过程，得到Detect模块的三个尺度的预测结果。

训练完成后，每一个网格(feature map的每一个点)学会了自身是否该对某个目标进行预测，也学会了对不同位置的目标进行中心点偏差值及目标宽高的预测，也学会了预测目标的类别，那么就应该使用这些预测值得到目标的实际类别和位置信息。那么预测环节，就应该让每一个网格都进行目标的预测，预测有目标的置信度、目标的位置和大小及目标的类别。然后使用置信度过滤删除不足够可信的预测结果，再进行NMS消除重叠的预测框，得到高置信度、位置和类别都相对可靠且不相互重复的预测结果。

其中位置信息如3.3.3节所述，训练过程中，输入图像经过backbone+head+Detect模块的卷积处理后得到的输出经过上面的公式处理后得到了$\Delta x,\Delta y,s_w,s_h$的值。那么在目标位置预测环节执行下述公式得到目标的中心点位置$(b_x,b_y)$和目标的宽高$p_w,p_h$：
$$\begin{gathered} b_x=\sigma (t_x)\ast 2-0.5+c_x \\ {b}_y=\sigma (t_y)\ast 2-0.5+c_y \\ {b}_w=(\sigma (t_w)\ast 2{)}^2\ast p_w \\ {b}_h=(\sigma (t_h)\ast 2{)}^2\ast p_h \end{gathered}$$
$t_x,t_y,t_w,t_h$是Detect模块的输出值，$c_x,c_y$是所在的网格的坐标，$p_w,p_h$是目标对应的anchor的宽高信息。

6.2 性能统计

统计各个训练epoch后模型在验证集上的P、R、mAP@0.5、mAP@0.5:0.95等值。AP@0.5:0.95是单个类别AP@0.5、AP@0.55、…、AP@0.9、AP@0.95求均值的结果，mAP@0.5:0.95是各类别AP@0.5:0.95的均值。