基于改进YOLO v4的热轧带钢表面缺陷检测

季娟娟，王 佳+，陈亚杰，卢道华，3

(1.江苏科技大学 机械工程学院，江苏 镇江 212100；2.中国船舶重工集团公司 上海船舶设备研究所，上海 200031；3.江苏科技大学 海洋装备研究院，江苏 镇江 212003)

0 引 言

基于深度学习的检测方法因对任意复杂函数都能实现较好地拟合且特征提取能力较强，被广泛地应用于产品缺陷检测领域[1]。程婧怡等[2]利用改进的YOLO v3[3]算法检测热轧钢条表面缺陷，通过使用DIoU[4]边框回归损失并利用K-Means++[5]聚类算法来提取更多小缺陷目标特征。孙永鹏等[6]针对冲压件缺陷检测提出了一种基于改进YOLO v4[7]的算法，使用改进的MobileNet v3[8]网络代替CSPDarknet-53[9]，并结合SE Block[10](sequeze and excitation block)来提高检测精度。

卷积神经网络通常提取的是局部特征，不能很好捕捉全局特征信息间的依赖关系，随着网络深度的增加，参数量与计算量也会随之增加，检测速度变慢。随着Transformer的快速发展，其凭借着优于CNN(convolutional neural network)的全局信息特征捕获能力，在CNN的应用领域中也取得了不错的成果。鉴于CNN与Transformer的不同优势，本文提出了一种基于改进YOLO v4的热轧带钢表面缺陷检测方法，主要改进有以下4点：①将YOLO v4的主干网络CSPDarknet-53替换成MobileViT[11]轻量化网络，在YOLO模型中嵌入Transformer，在不降低检测速度的同时，提高检测精度并大大减少参数量；②利用K-Means++聚类算法来修改先验框的尺寸，使定位更加准确，减少模型损失；③在PANet[12](path aggregation network)后面添加ASFF[13](adaptively spatial feature fusion)模块，增强多类型缺陷的特征提取能力；④使用改进的DIoU NMS代替NMS进行后处理操作，减少误删量。为了验证改进的有效性，本文做了多个对比实验以及对每一个改进方法做了消融实验。

1 YOLO v4算法

将大小为W×H×S的图片输入到YOLO v4中。首先利用CSPDarknet-53进行特征提取；接着通过Neck部分进行多尺度特征融合；然后在Yolo Head部分利用3个不同尺度的特征层进行多尺度检测并得到初步的预测结果；最后对初步的预测结果做置信度过滤、NMS等后处理操作得到最终的输出结果，输出的结果经过式(1)的换算，可得到实际预测框的宽(bw)、高(bh)、中心点坐标 ((bx，by))。 式(1)中的cx、cy为当前锚框所属的网格左上角坐标值，pw、ph为锚框的宽高，σ为sigmoid激活函数，将函数内的值约束在0～1内YOLO v4的网络结构如图1所示。

图1 YOLO v4网络结构

(1)

2 改进的YOLO v4算法

为了充分利用输入图片全局特征的空间位置与语义之间的关系，本文算法将YOLO v4的主干网络CSPDarknet-53替换成MobileViT轻量化网络，并将前两个MobileViT block(MVIT)和最后的3*CBL模块提出作为Neck部分的输入；Neck部分通过SPP结构利用4个不同Pooling Size的最大池化对特征图进行不同程度的感受野增大，生成4个特征图后进行特征拼接，接着通过PANet结构进行从下至上，从上至下的多次特征提取拼接，多尺度地融合语义信息和空间信息，最后通过ASFF模块自适应地融合多尺度特征信息，加强特征的提取；Yolo Head部分通过3个不同大小的特征层进行多尺度检测，每个特征层由1个CBL模块和普通卷积组成。本文算法的网络结构如图2所示。

图2 本文算法网络结构

2.1 替换特征提取网络

CSPDarknet-53是典型的CNN网络模型，有较好的空间局部表征信息提取能力，然而计算量大，网络深度较深，容易丢失重要特征。与CNN不同，Transformer有更好的全局特征学习能力，在识别分类任务中有更高的精度，但是参数量大，在实时检测中达不到低延迟的要求。而MobileViT则使用Transformer将卷积的局部建模换成全局建模，结合了CNN和Transformer的优势，在参数和精度方面大幅度超过了MobileNet v3等轻量级网络架构[11]。

MobileViT主要是由多个MobileViT block和MobileNet v2[14]block组成。MobileViT block的结构如图3所示。

图3 MobileVit block结构

Local representations部分主要是用来提取局部空间信息，利用n×n与1×1的卷积对输入特征进行局部特征提取并投影到高维空间；Global representations部分主要用来学习全局特征，将特征展开为N个non-overlapping flattened patches，再利用Transformer对多个patches之间的关系进行编码；Fusion部分利用n×n卷积来融合前两个部分得到的局部、全局特征。

MobileNet v2 block又称为倒残差。先通过1×1的卷积升维，进一步提取更丰富的特征；然后利用3×3的深度可分离卷积，在使用较少计算量的情况下来提取更大感受野的特征；再利用1×1的卷积将输入的特征图降维，减少参数量；最后再利用残差连接，使其在较少增加计算量与参数量的同时又可以捕获到更多的特征信息。结构如图4所示。

图4 MobileNet v2 block结构

2.2 采用K-Means++锚框聚类

YOLO v4设定的先验框尺寸是针对VOC数据集上的图片聚类得到的，并不适合热轧带钢表面的小尺寸缺陷检测，因此，要采用聚类算法针对本文所研究的数据集生成更利于模型收敛的先验框尺寸。聚类算法最常用的是K-Means[15]算法和K-Means++算法，聚类中心初始值设定的好坏对聚类结果有着一定的影响。K-Means算法的聚类中心是提前指定好的，而K-Means++算法是在聚类中心的选择过程中自动选择较优的聚类中心，相较于K-Means算法，能更加灵活的生成代表数据集的先验框尺寸，因此本文采用K-Means++算法对数据集进行聚类。聚类结果为(19，48)，(25，170)，(40，85)，(50，39)，(67，157)，(81，72)，(146，76)，(154，29)，(156，181)。

2.3 添加ASFF模块

ASFF模块可通过学习后的权重参数自动地融合多尺度特征信息，加强特征的提取。本文算法使用的是一个三特征层的ASFF结构，如图5所示。

图5 三特征层ASFF结构

每个ASFF的输入为Level 1、Level 2、Level 3。当Level 3特征图中既有大物体，又有小物体时，小物体的存在会对特征提取有一定的干扰性，此时应该过滤掉小物体的信息，并结合Level 1与Level 2特征图中含有的高级语义特征信息来丰富Level 3特征图中的大物体特征信息。以ASFF-3为例，3个层的特征经过变换后乘上相应的权重系数并相加，即可得到新的融合后的特征图ASFF-3，其计算公式如式(2)所示

ASFF-3=X1→3*α3+X2→3*β3+X3→3*γ3

(2)

式中：X1→3为Level 1特征层通过通道数调整、尺寸缩放，变换至同Level 3特征层相同大小的特征层；同理，X2→3为Level 2变换至同Level 3特征层相同大小的特征层；X3→3为Level 3特征层；α，β，γ为各自的权重系数，在训练中通过反向传播获得。

2.4 改进的DIoU NMS

在通过Yolo Head得到初步的预测结果后，还需进行一系列的后处理操作才能得到真正的预测结果。后处理操作中，对结果影响最大的就是非极大值抑制(NMS)，NMS能过滤掉冗余的预测框，提高检测的准确度。NMS的具体步骤如下：

步骤1 对所有的候选预测框按照置信度值进行降序排列；

步骤2 选择置信度值最高的预测框A添加到输出列表中，并将其从候选列表中删除；

步骤3 计算框A与其它候选预测框的IoU值，并将大于NMSthreshold的候选预测框视为重叠预测，将其从候选列表中删除；

步骤4 重复上述3个步骤，直至候选列表为空后，返回输出列表，即真正的预测框。

(3)

(4)

图6 DIoU

采用大小不同的λ值进行多次实验对比，作出不同λ值的mAP(mean average precision)对比，如图7所示，可知λ取0.83时检测精度值最高，mAP值为80.16%。

图7 不同λ值的mAP对比

3 实验结果与分析

3.1 实验数据

图8 6类缺陷

3.2 评价指标

目标检测的评价指标主要有准确率(Accuracy)、检测速度(FPS)、召回率(Recall)、精准率(Precision)、平均精度AP(Average Precision)、mAP等等。本实验主要研究的是工业生产环境下的表面缺陷，既要满足实时性又要有好的检测性能，所以采用FPS与mAP作为评价指标来比较各算法。FPS是每秒可检测完的图片数量，mAP是各类缺陷平均精度的平均值。mAP、Precision、Recall的计算公式如下

(5)

(6)

(7)

其中，n为缺陷类别数；AP为每类缺陷的平均精度值，可由P-R曲线下的阴影面积获得；TP为被模型预测为正类的正样本；FN为被模型预测为负类的正样本；FP为被模型预测为正类的负样本。

3.3 实验结果

本实验是在ubuntu20.04操作系统下基于pytorh深度学习框架实现的，GPU选用内存12 G的NVIDIA GeForce GTX 3060，CUDA11.1，CUDNN11.1。训练150个轮数，冻结与解冻时的batchsize设为4，使用Adam优化器进行参数处理，冻结时的学习率为le-3，解冻时的学习率为le-4，训练集在解冻训练时采用在线数据增强。

本实验训练所得的损失函数曲线如图9所示。

图9 损失函数曲线

其中，横坐标是训练轮数，纵坐标是平均损失，由图9可看出，本文算法收敛速度较快。前50个训练轮数为冻结训练部分，后100个训练轮数为解冻训练部分，验证集损失在前50个训练轮数中的趋势从下降变为开始上升，在第50个训练轮数处由于模型解冻且学习率改变，验证集损失与训练集损失的损失曲线发生轻微的突变，但又迅速下降并开始收敛，尤其是验证集损失不仅停止了上升的趋势且变得下降收敛，因此损失函数曲线的整体走向趋势是下降的。

本实验算法检测得到的各类P-R曲线如图10所示。

图10 各类P-R曲线

P-R曲线是由横轴为召回率，纵轴为精准率构成的曲线，由图10可知，各类的P-R曲线往右上外凸，越接近右上角，P-R曲线下的面积越大，AP值就越高，模型的性能也就越好。

3.4 对比分析与消融实验

3.4.1 聚类算法对比

将K-Means、K-Means++两种算法在本实验数据集上聚类生成的先验框与原始的YOLO v4先验框，分别应用在YOLO v4算法上进行对比，对比结果见表1。

表1 3种先验框的mAP值

由表1可知，K-Means++算法聚类生成的先验框在YOLOv4上有更高的精度值，分别比原YOLO anchor、K-Means anchor高了2.04%、1.41%，对热轧钢表面小缺陷的检测准确度更高，可提高边框回归位置预测的精确度。

3.4.2 YOLO v4算法与本文算法结果对比

将YOLO v4算法与本文算法分别对数据集进行检测，所得到的AP值的对比如图11所示。

图11 两种算法mAP值对比

由图11可知，本文算法相较于YOLO v4算法，夹杂类与斑块类的平均精度均提升了5%，麻点类提升了10%，划痕类提升了11%，压入氧化皮类提升了12%，裂纹类更是提升了26%。6类缺陷的检测精度都得到了不同程度的提升，尤其是裂纹类。

从6类缺陷的检测结果中，分别抽出一张进行对比，对比的结果如图12所示。

图左列的图(a)是YOLO v4算法的结果，右列的图(b)是本文算法的检测结果。由图12对比可知，YOLO v4算法更容易出现漏检的情况，且部分缺陷无法识别或者识别出的置信度值较低。而本文算法则能够尽可能检测出图上应有的缺陷。

3.4.3 不同算法对比

为验证本文算法的综合性能，用Faster R-CNN、YOLO v3、YOLO v4、YOLO v5、文献[2]、文献[6]算法对实验数据集进行检测，并与本文算法的结果进行对比，对比的结果见表2。

由表2可知，本文算法的性能均优于Faster R-CNN与YOLO v4，检测速度虽然比YOLO v3与YOLO v5略低，但精度值最高。文献[2]采用四特征层输出，相较于YOLO v3精度值有所提高，但由于使用的是全卷积神经网络，特征提取不是很全面，因此对检测精度值的提高低于本文算法。文献[6]采用MobileNet v3作为主干网络，由于MobileNet v3能在保证一定检测精度的前提下拥有较强的提速能力，因此在检测速度方面超过了本文算法，但由于跟文献[2]一样，均采用全卷积神经网络，所以检测精度值依然低于本文算法。

本文算法由于结合了Transformer和卷积神经网络的优势，具有较强的全局、局部特征提取能力，再加上使用了ASFF模块，K-Means++算法与改进的DIoU NMS算法来进一步地提高检测精度，所以本文算法的检测精度最高。综合两个评价指标可知，本文算法的综合检测性能最佳，做到了准确度更高的实时检测。

3.4.4 消融实验

将本文算法对原YOLO v4算法的每一个改进做消融实验，结果见表3。

表3 消融比较

由表3可知，仅更换主干网络，检测的精度值与速度就能得到提高；利用K-Means++重新聚类先验框，精度提升的幅度更大；再利用ASFF模块增强对特征的利用率，精度得以提升，但速度相较于之前有所下降；最后再加上改进的DIoU NMS的替换，进一步提高检测精度。最终，整个改进的算法使热轧带钢表面缺陷检测精度值得到了11.57%的提升。

4 结束语

本文针对热轧带钢表面缺陷检测难度大，识别准确率低等问题提出了一种基于改进YOLO v4的检测算法，该算法在YOLO v4框架下，结合Transformer技术重构出一个新的深度学习网络。采用新主干网络MobileViT并增加ASFF模块来增强对微小缺陷的特征提取能力；使用K-Means++聚类和改进的DIoU NMS来增强算法对本实验数据集检测的鲁棒性。总体上是，结合一个工业检测应用场景，组合两种主流技术形成的一种解决方案。最后通过多个对比实验与消融实验验证了本文算法改进的有效性。之后，将针对检测精度与检测速度进行进一步提升。