计算机系统应用  2024, Vol. 33 Issue (2): 265-275   PDF    
基于Efficient-YOLO的轻量化轴承缺陷检测
娄瑶迪, 岳俊峰, 周迪斌, 刘文浩     
杭州师范大学 信息科学与技术学院, 杭州 311121
摘要:针对现有深度模型在工业轴承外观缺陷检测领域, 存在模型参数量大、特征融合不充分以及对小目标检测精度低等问题, 提出了一种轻量化自适应特征融合检测网络(Efficient-YOLO). 首先, 该网络采用嵌入CBAM注意力机制的EfficientNetV2结构进行基本特征提取, 便于确保模型精度同时显著优化模型参数量; 其次, 设计了一种自适应特征融合网络(CBAM-BiFPN), 用来增加网络对有效特征信息的提取; 接着, 在下游特征融合网络引入Swin Transformer机制, 同时配合上游网络引入的Ghost卷积, 大幅度提高模型对轴承外观缺陷的全局感知能力; 最后, 在推理阶段运用改进的非极大值抑制方法(Soft-CIoU-NMS), 加入距离有关的权重评价因素, 减少了重叠框的漏检. 实验结果表明: 与现有主流检测模型相比, 此方法在轴承表面缺陷数据集上, mAP达到了90.1%, 参数量降低至1.99M, 计算量为7 GFLOPs, 对轴承缺陷小目标识别率显著提升, 满足工业现场轴承外观缺陷检测需求.
关键词: 轴承缺陷检测    深度学习    EfficientNetV2    YOLOv5    特征融合    
Lightweight Bearing Defect Detection Based on Efficient-YOLO
LOU Yao-Di, YUE Jun-Feng, ZHOU Di-Bin, LIU Wen-Hao     
School of Information Science and Technology, Hangzhou Normal University, Hangzhou 311121, China
Abstract: Since the existing deep model faces many problems such as a large number of model parameters, insufficient feature fusion, and low detection accuracy for small targets in the field of industrial bearing appearance defect detection, a lightweight adaptive feature fusion detection network (Efficient-YOLO) is proposed. First of all, the network uses the EfficientNetV2 structure embedded in the CBAM attention mechanism for basic feature extraction to ensure model accuracy and significantly optimize the model parameters. Secondly, an adaptive feature fusion network (CBAM-BiFPN) is designed to strengthen the network’s extraction of effective feature information. Then, the Swin Transformer mechanism is introduced in the downstream feature fusion network, and the Ghost convolution introduced by the upstream network is used to greatly improve the model’s global perception of bearing appearance defects. Finally, the improved non-maximum suppression method (Soft-CIoU-NMS) is applied in the inference phase, with distance-related weight evaluation factors added, so as to reduce missed detection of overlapping frames. The experimental results show that compared with the existing mainstream detection models, the method has a mAP of 90.1% on the bearing surface defect dataset. The number of parameters is reduced to 1.99M. and the calculation amount is 7 GFLOPs. The recognition rate of small targets with bearing defects is significantly improved, which meets the needs of industrial bearing appearance defect detection.
Key words: bearing defect detection     deep learning     EfficientNetV2     YOLOv5     feature fusion    

轴承是机械设备中的重要组成部件, 其外观缺陷往往影响自身品质. 然而, 在生产和保存过程中, 轴承易受制造工艺和环境因素影响, 出现各类缺陷, 如锈斑、擦伤、划痕和缺口等. 不同类型的缺陷特征之间存在显著差异, 因此难以设计一个适用于所有缺陷评估的通用算子, 这一问题造成了该领域进展缓慢的现象. 同时, 外观缺陷可能出现在正反端面、内外侧面或倒角等位置, 给图像采集和检测增加了复杂性.

随着人工智能技术的不断进步, 深度学习技术在工业零件表面缺陷检测领域得到广泛应用. 主流的深度学习目标检测算法可分为单阶段和两阶段检测算法. 两阶段目标检测网络通常需要多次候选框生成、再分类和目标定位过程, 导致较长的计算时间. 相比之下, YOLO算法是一种单阶段目标检测算法, 仅需一次前向传播, 大大降低了计算复杂度, 提高了检测速度. 这使得YOLO在工业轴承缺陷检测中更适用于快速检测目标信息的需求. 此外, 工业缺陷数据通常包含多个角度和尺度的图像, 以涵盖不同方向和大小的缺陷表现. 对于这种多样化的缺陷数据, YOLO算法通过引入多尺度的特征图来检测不同大小的目标. 多尺度特征表示有助于更好地捕捉小目标的细节和上下文信息, 提高对小目标的检测准确性. 为进一步提升模型性能, 研究人员提出了一系列基于YOLO算法的改进, 例如: Zhao等[1]提出了一种轴承缺陷检测算法(GRP-YOLOv5), 在特征提取阶段将ResC2Net与残差状结构相结合, 在融合部分增加了PConv卷积, 提高了模型捕获缺陷的能力, 在化工设备缺陷轴承数据集上精度达到93.5%; Li等[2]提出了基于改进YOLOv5算法的深度学习缺陷检测方法, 结合了卷积注意力机制, 并在颈部网络引入双向特征金字塔网络(BiFPN)模块, 在金属轴数据集上精度达到93.6%; Niu等[3]提出了一种改进的YOLOv5模型, 用于检测印刷电路板缺陷. 首先使用K-means++方法改进了先验锚框的位置匹配, 然后引入了ECA-Net模块, 以增强模型对通道特征的敏感性. 在公开的PCB数据集上, 该模型达到了99.1%的精度; Chen等[4]提出一种Transformer-YOLO模型, 使用Transformer作为特征提取网络, 模型在电路板数据集上精度达到97.04%. 上述模型虽然在各自的数据集上达到了较高的精度, 但是还存在模型参数量大以及检测实时性较差的问题. 因此, 为降低模型的体积并提高检测速度, Qian等[5]提出了LFF-YOLO模型用于多尺度缺陷检测, 首先使用ShuffleNetv2[6]作为特征提取网络, 其次提出了轻量级特征金字塔网络(LFPN), 实现了检测速度的提升, 简化后的模型参数减少74.6%; Zhao等[7]提出了一种LSD-YOLOv5模型用于钢带表面损伤识别, 在骨干网络中使用MobileNetV2, 并且引入坐标注意力机制, 颈部网络使用了一个更小的双向特征金字塔网络(BiFPN-S), 成功减少了模型的体积, 参数量降低了61.5%. Fan等[8]构建了高精度低延迟的缺陷检测算法(GCB-Net), 模型基于改进的YOLOV4目标检测网络, 引入了卷积注意力模块和Ghost轻量级网络, 完成对缺陷类别的预测, 模型参数量压缩了82.3%. Xue等[9]提出一种基于改进YOLO算法和ResNet18骨干特征提取网络的煤矸石检测算法, 研究了模型的特征尺度缩减和非结构化修剪, 模型体积可以压缩到原来的28.5%. 上述算法虽然实现了模型的轻量化, 但是由于模型参数和结构的减少, 导致了检测精度的下降, 以及对小目标检测能力的降低.

针对上述问题, 结合轴承缺陷检测工业场合低成本的目标实时检测需求, 本文基于YOLOv5基模型提出了一种轻量化自适应特征融合的Efficient-YOLO模型用于轴承外观检测, 并从以下方面进行改进: (1)为解决目前轴承检测中模型参数量较大且占用资源较多的问题, 采用改进的EfficientNetV2作为骨干网络, 在保证算法精度的情况下实现算法的轻量化; (2)设计了一种自适应特征融合网络CBAM-BiFPN, 用于提高模型特征融合能力, 增强对轴承缺陷中边缘目标和不明显目标的检测能力, 同时减少模型参数量; (3)为提高对轴承缺陷小目标的检测能力, 使用C3STR模块替代C3模块, 并且用GhostConv[10]替代传统的深度分离卷积; (4)为减少轴承缺陷图像中重叠缺陷的漏检情况, 在后处理阶段, 通过使用Soft-CIoU-NMS替代非极大值抑制算法(non maximum suppression, NMS)[11], 解决了NMS存在的重叠框选择难题.

1 YOLOv5模型结构

YOLO是一种目标检测算法, 且不断迭代更新. 本文基础网络部分使用的是YOLOv5 6.0版本, 除输出端外, 主要包括3个部分: Backbone、Neck、Head.

输入端主要由Mosaic图像增强、自适应锚框计算以及自适应图片缩放等模块组成, 用于增加训练集的多样性, 提高模型的泛化能力.

Backbone结构负责特征提取, 主要部分为卷积层, C3结构和SPPF (spatial pyramid pooling fast)结构. 标准的卷积层结构用于对特征图的压缩和下采样; C3结构拼接两部分特征图, 作为下层的输入; SPPF结构对不同大小的特征图进行池化, 增强模型对不同尺度目标的感知能力.

Neck颈部结构包括特征金字塔网络(feature pyramid network, FPN)[12]路径聚合网络(path aggregation network, PAN)[13], 以实现特征融合和尺度适应. FPN结构通过自上而下的上采样操作, 传递语义信息. PAN结构则通过自下而上的下采样操作, 传递定位信息. 这两种结构相互补充, 聚合不同层次的特征信息, 实现对不同尺寸和不同大小目标的适应, 提高模型的检测效果.

Head头部结构包含3个检测头, 后处理中使用非极大值抑制来消除冗余框.

2 本文方法 2.1 Efficient-YOLO模型结构

Efficient-YOLO的整体结构如图1所示, 输入的图片调整为640×640×3的统一形状, 首先通过EfficientNetV2[14]骨干网络进行特征提取, 生成10×10×512的特征图, 然后输入到自适应特征融合网络(convolutional block attention module-bidirectional feature pyramid network, CBAM-BiFPN)中, 最后将输出的3种特征图传递给耦合头部网络, 生成检测框.

图 1 Efficient-YOLO结构图

2.2 改进的 EfficientNetV2模块

在工业轴承检测任务中, 传统的YOLOv5骨干网络使用大量卷积操作, 导致高度的计算资源消耗, 需要花费较长的时间来执行检测任务, 检测实时性能较差. 因此, 本文使用轻量化EfficientNetV2结构替换原有的骨干网络, 降低参数量, 减少计算复杂度, 提高模型的检测速度.

EfficientNetV2是在EfficientNet[15]的基础上进行了改进, 解决了原始结构中存在的复杂度大和计算成本高的问题, 包括2个组成部分: MBConv模块和Fused-MBConv模块, 如图2(a)和图2(b)所示. 其中, MBConv模块是EfficientNet中的基础结构, 实现了深度神经网络的高效训练; 而Fused-MBConv由升维卷积和降维卷积块组成, 可以在很小的计算量及很低的参数量条件下, 提高训练和推理的效率.

原MBConv模块引入了通道注意力模块(squeeze-and-excitation, SE)[16], 该模块对每个通道进行加权, 提高网络对重要的特征通道的关注度, 可在学习特征时更加灵活, 相比于原模型中的SE注意力模块只关注于通道域, 卷积注意力模块(convolutional block attention module, CBAM)[17]包括两个注意力分支, 分别是针对通道和空间维度的注意力. 特征图被输入之后, 先通过通道注意力获得新的特征图, 再通过空间注意力, 得到最终的特征图, 此过程可以表示为式(1)和式(2), 其中$F'$表示通过通道注意力获得的新特征, $F''$为最后输出的特征, 符号$ \otimes $表示元素相乘.

$ F' = {M_c}\left( F \right) \otimes F\; $ (1)
$ F'' = {M_s}\left( {F'} \right) \otimes F'\; $ (2)
图 2 MBConv、Fused-MBConv结构图

CBAM模块如图3所示.

图 3 CBAM模块图

为了在大特征图提升小目标的关注度, 本文在浅层网络将CBAM注意力模块嵌入到MBConv模块中, 替换了原有的SE注意力模块. 在降维Conv之前嵌入CBAM, 可有效筛选升维后深度分离卷积输出的高维度语义特征, 少量增加参数和计算量, 并且在保留重要细节特征图的同时, 额外为网络提供筛选重要度较高的语义信息, 将浅层的强定位特征传递到深层网络, 改进后的MBConv结构如图4(a)所示.

由于浅层特征提取网络可保留更多的小目标特征信息和位置信息, 在稍微增加参数量以及计算量的情况下, 本文将MBConv模块中的深度分离3×3卷积替换为Ghost 5×5卷积, 结构如图4(b)所示. 使用Ghost卷积的方法, 可以仅通过少量的线性变换就能生成新的特征图, 在轻量化的网络中能够保留更多冗余特征. Ghost卷积分为3个部分, 第1部分是普通卷积操作, 先生成少部分通道数的特征图, 第2部分是分组卷积操作, 通过对第1部分生成的特征图进行线性运算生成剩余特征图以保证通道数不变. 第3部分就是将前两部分得到的特征图拼接. 替换为Ghost 5×5卷积可以在有限的感受野范围内保留更多边缘信息, 以此增强对小目标缺陷的定位准确性, 改善图像边缘部分缺陷的遗漏情况, Ghost卷积结构如图5所示.

图 4 CBAM-MBConv、Ghost-MBConv结构图

图 5 Ghost卷积结构图

2.3 CBAM-BiFPN模块

YOLOv5模型通过建立双向特征融合网络增强了特征等级结构的局部建模能力. 然而, FPN和PAN模型未考虑特征融合层之间的语义信息差异, 存在多尺度信息丢失的问题, 缺乏对边缘小目标和密集目标的特征利用. 为了补偿信息损失, EfficientDet检测算法提出了加权双向特征金字塔网络BiFPN (bidirectional feature pyramid network, BiFPN), 在同一层的输入节点和输出节点之间新加一条边, 并且在双向网络中简化单一节点, 以便在不增加太多成本的情况下融合更多的特性.

YOLOv5在FPN和PAN结构的基础上, 把P4特征层也加入到特征融合中, 如图6(a)所示. 此结构可以额外融合中层提取网络的特征, 但是由于PAN是作为整个Neck部分重要的信息融合金字塔结构, 因此对输入分辨率和目标大小敏感且不够灵活. 额外的加权操作不仅给PAN部分增加了计算量, 而且容易使模型产生过拟合的情况.

图 6 YOLOv5-BiFPN和CBAM-BiFPN结构图

为适用于本文的轻量化网络设计目标, 并且提高对轴承细节特征的关注度, 设计了一种自适应特征融合网络CBAM-BiFPN, 如图6(b)所示. 通过引入CBAM注意力机制于加权边以自适应抽取重要度较高的特征, 这种特性能够有效抑制不重要特征, 为加权操作提供了更为丰富的语义信息流. 此外, 引入CBAM模块的加权边可有效规制模型的过拟合倾向, 使得模型的泛化性能更强, 能够更好地适应不同样式和种类的轴承缺陷检测任务. 与YOLOv5提出的BiFPN改进方法相比, 本文提高模型对特征细节感知能力的同时, 简化了上采样过程的通路节点, 优化了卷积窗口大小, 实现了节点之间的计算量大幅度减少、整体网络参数量的下降.

2.4 C3STR模块

在深层网络的高级特征图中, 轴承图像中小目标的特征信息逐渐丢失, 因此本文借鉴Swin Transformer[18]的思想, 在原模型的大尺寸检测头中, 引入C3STR结构, 借助窗口自注意力模块增强小目标的语义信息和特征表示, 为底层带来更丰富的语义信息, 尤其是全局的语义信息.

基于Swin Transformer改进的C3STR结构如图7所示. 其中STR模块是Swin Transformer block的封装, 由两种Swin Transformer block成对组合而成. 这两种结构一个使用窗口多头自注意力模块W-MSA, 另一个使用滑动窗口多头自注意力模块SW-MSA, 两者的配合使用使得相邻窗口间的信息得到有效传递.

图 7 C3STR 结构图

本文使用C3STR加强卷积块替换了C3结构中的Bottleneck堆叠部分, 不仅能够通过自注意力机制使全局建模图像依赖关系, 提高模型对全局信息的特征提取能力, 而且C3STR结构支撑于小检测头部分, 能够在增加少量参数和计算量的情况下提高对全局大目标的检测性能.

2.5 Soft-CIoU-NMS模块

轴承缺陷图像中的部分缺陷往往以高密度和相互重叠的方式出现, 传统目标检测使用NMS方法进行后处理, 常常导致相邻检测框的分数强制设置为零, 会将重叠部分大于重叠阈值的检测框移除, 导致漏检的情况. 改进的Soft-NMS算法[19]的流程和NMS算法类似, 但是对于IoU大于阈值的检测框不是直接删除, 而是使用函数运算降低得分数. 本文使用的Soft-CIoU-NMS结构, 是基于Soft-NMS和CIoU算法, 旨在解决传统Soft-NMS方法对距离度量不敏感的问题, 增加对检测框的考量, 减少对密集目标的错误抑制, 从而提高检测精度. Soft-CIoU-NMS的公式可以表示为:

$ S_i^{{\mathrm{Soft {\textit{-}} CIoU}}} = {S_i}\left( {1 - \mathop {\max }\limits_{j \ne i}^N \frac{{CIoU\left( {{b_i}, {b_j}} \right)}}{{\rho \left( {i, j} \right)}}} \right) $ (3)

其中, $ {S}_{i} $表示第i个框的检测分数, $ {b}_{i} $表示第i个框的边界框, $ \rho \left(i.j\right) $表示第i个框和第j个框的距离度量, $CIoU\left( {{b_i}, {b_j}} \right)$表示基于CIoU度量的重叠度量. 公式中的$\mathop {\max }\limits_{j \ne i} $表示对于第i个框, 降低最大IoU对应框的分数, $\dfrac{{CIoU\left( {{b_i}, {b_j}} \right)}}{{\rho \left( {i, j} \right)}}$表示基于距离度量的重叠度量.

3 实验结果 3.1 实验数据与设置 3.1.1 数据集采集

本次实验数据集为自制的轴承表面缺陷数据集, 从不同条件共采集3802张分辨率为1280×1024的轴承图片, 经过数据增强扩充至9505张, 并将缺陷分为划痕、擦伤、缺口和锈斑等4个类别. 使用LabelImg软件对数据进行标注, 按照8:1:1的比例分为训练集、验证集和测试集, 其中训练集包含7604张图片, 验证集包含950张图片, 测试集包含951张图片, 部分数据如图8所示.

图 8 数据集抽样

3.1.2 模型训练与推理

实验基于PyTorch深度学习框架, 所有模型训练从零开始, 使用GPU进行训练, epoch设置为200轮, 其中前30轮用于热身阶段. 优化器使用SGD随机梯度下降法. 初始学习率为0.01, 动量为0.937, 使用余弦退火策略将学习率降至0.00001, 权重衰减为0.0005. 推理验证测试统一使用CPU进行.

实验环境为Ubuntu 20.04, PyTorch 2.0, CUDA 12.2. CPU型号为Intel I7 13700k, 显卡为NVIDIA RTX4090. 使用COCO检测和评估标准中的指标进行性能评估, 包括平均精度(AP), IoU阈值为0.5的平均准确度(mAP), 不同阈值(从0.5至0.95, 步长为0.05)下的平均准确度(mAP@.5:.95), 参数量(Params), 浮点运算(GFLOPs)和每秒帧数(FPS). 其中APmAP计算公式为式(4)和式(5), 其中P表示精确率, R表示召回率.

$ AP = \int ^1_0 P\left( R \right){\rm{d}}R $ (4)
$ mAP = \frac{{\displaystyle\sum\limits_{i = 0}^n {AP\left( i \right)} }}{n} $ (5)
3.2 实验结果对比 3.2.1 各种目标检测算法对比

为验证Efficient-YOLO算法性能, 将其与主流目标检测算法的进行对比, 如表1所示.

表 1 各种目标检测算法对比实验

表1中可看出, 相较于其他模型, 本文提出的改进模型在几个方面展现了卓越的高效性. 首先, 该模型在精度上取得了更高的表现, 即在目标检测任务中能够更准确地定位和识别目标. 其次, 该模型的检测速度更快, 即在相同的计算资源下能够更快地完成目标检测任务. 此外, 该改进模型的体积更小, 占用的计算资源也更少, 因此能够更容易地在嵌入式设备或有限计算资源的工业环境中部署和应用. 对比于传统两阶段检测算法Faster-RCNN[20]以及SSD模型[21], 改进后的模型在精度、参数量方面都有大幅提升、计算量也大幅减少, 检测速度提高了数倍. 与原始模型YOLOv5s相比, Efficient-YOLO在平均精度上提高了2.5%, 达到了90.1%, 参数量降低了5.03M, 仅有1.99M, 计算量降低了8.8 GFLOPs, 检测速度提高了10f/s, 在各类缺陷的检测精度上也都有不同程度的提升.

图9(a)可得, 与同系列YOLO算法模型相比, 本文模型的各类缺陷检测精度处于最高点, 说明在各类缺陷的检测精度均高于其他模型. 由图9(b)可得, 改进模型处于图的最左侧, 比YOLOv4-tiny的参数量还小, 在高度方面处于最高点, 说明本文模型在mAP上表现最好.

图 9 YOLO系列检测网络效果对比

图 9 YOLO系列检测网络效果对比(续)

为进一步证明本文提出的改进算法在轴承表面细微小目标检测中的性能效果, 参照于YOLOv5s基准模型测试结果, 分别从AP, Recall及特征图可视化等方面进行评价分析.

在本文的轴承表面缺陷数据集中, 缺口类缺陷作为显著性的小目标, 在阈值(threshold)都为0.6的情况下, 本文模型比YOLOv5s基线模型在缺口类breach上的Recall值提升了3.5%, AP值提高了3.5%; 在细微划痕类上的Recall值提升了4.0%, AP值提高了4.4%. 图10中分别展示了在小目标缺口和细微划痕类中, IoU阈值随着步长为0.1变化时, 本文模型(橙线)与原模型(蓝线)的AP值及Recall率曲线对比.

图 10 APRecall对比

图11显示了原模型和改进模型的特征图可视化对比效果. 由图可知, 改进后的模型比原模型的热力点面积更大, 说明对目标的关注度更高; 改进后的模型比原模型的热力点更多, 说明目标漏检情况得到了改善, 小目标的检测能力更强.

3.2.2 基于YOLOv5s的不同骨干网络对比

本文首先以YOLOv5s为基线进行轻量化改进, 分别使用EfficietLite[15]、MobileNetV3[22]、ShuffleNetV2、GhostNet[10]以及EfficientNetV2等主流特征提取网络重构YOLOv5s的骨干部分, 模型训练后对轴承表面缺陷数据集进行检测推理, 实验结果见表2.

图 11 特征图可视化

表 2 不同骨干网络对比实验

表2可知, 以EfficietLite、MobileNetV3、ShuffleNetV2、GhostNet等作为骨干的模型, 对比于原YOLOv5模型在参数量以及计算量方面都有所下降, 但是牺牲了精确率, 而继承了EfficientNetV2作为骨干的模型, 在保证精度的情况下, 模型体积较小, 检测速度较快, 降低了资源占用率, 对比于其他算法性能更优.

3.2.3 嵌入不同注意力机制的MBConv模块对比

为体现本文所用的CBAM注意力机制的性能效果, 将其与SE、CA[23]、ECA[24]、CoT[25]等几种经典注意力机制进行对比, 如表3所示.

由实验可得, 对比于其他注意力机制, 本文所使用的CBAM注意力机制结合卷积和注意力机制, 能够充分利用空间和通道的信息, 关注不同空间位置的重要性. 在mAP@.5上表现要好于其他注意力机制, 且在参数量以及检测速度方面均处于最优.

表 3 嵌入不同注意力机制的MBConv模块对比实验

3.2.4 各种特征融合网络对比

为验证CBAM-BiFPN作为特征融合网络的性能表现, 本文以EfficientNetV2为骨干网络的YOLOv5s为基准, 将CBAM-BiFPN与FPN、FPNPAN、BiFPN这3种主流特征融合网络进行对比实验.

表4所示, 在mAP表现上, 原模型的精度更高, 但也导致参数和计算量较大. 在参数量方面, 使用CBAM-BiFPN结构替换后的网络模型, 仅拥有1.89M参数量, 且整体网络的计算量大幅下降, 相较于使用FPN结构的模型, 参数量降低了40.5%. 在测试速度方面, 本文的模型也稍快于其他模型.

表 4 各种特征融合网络对比实验

在第2.2节中, 阐述了通过作用于加权边的CBAM注意力模块可增强BiFPN网络的泛化性能, 为了证明这一特性, 本文列出了对比于YOLOv5提出的BiFPN改进方法的训练误差轮次下降结果, 分为cls_loss, box_loss, obj_loss这3部分, 如图12所示, 在YOLOv5训练loss计算方法中, 其中cls_loss用于监督类别分类, box_loss用于监督检测框的回归, obj_loss用于监督grid中是否存在物体, loss的计算在每层feature mAP上进行. 由表4可知, 在相同mAP@.5精度表现的两种加权特征融合网络, 本文所提出的CBAM-BiFPN训练误差结果下降和收敛均略高于YOLOv5-BiFPN水平, 可见轻量化模型过拟合倾向得到了有效抑制, 不难预见此种训练得到的模型泛化性能相比其他更强, 可期待在工业应用中得到更好的鲁棒性.

3.3 消融实验

为确认本文所提出的Efficient-YOLO的性能优势, 在轴承表面缺陷数据集上, 设计了5组消融实验. 消融实验中不使用迁移学习, 以准确比较设计模块的性能. 如表5所示.

图 12 CBAM-BiFPN、YOLOv5-BiFPN 的损失曲线

表 5 消融实验

表5可知, 实验①使用改进EfficientNetV2同时加入CBAM-BiFPN, 并作为后续实验的对比基线, 对比于原模型大幅度降低了参数量以及计算量, 精度也得到小幅度改善. 由实验分析可见, 用Ghost Conv替换MBConv-CBAM模块的深度分离3×3卷积, 获得了更多的位置信息, 保留了更多边缘信息, 使得小目标的检测精度提高, 在scratch类上表现良好. 使用C3STR模块替代原先的C3结构, 由于Swin Transformer在深层网络处的语义信息获取能力, 在scratch类、breach类以及abrade类中均有明显提升, 在增加少量参数的情况下, 增强了模块对各类缺陷的关注度. 使用Soft-CIoU-NMS替换原先的NMS结构, 解决了重叠框的选择难题, 提升了对距离以及检测框的关注度, 使得检测精度得到了提高. 综合原模型和最终改进的模型的比较结果, 改进后的模型, 在mAP、GFLOPs和FPS均优于原先的YOLOv5s模型.

针对小目标、特征不显著的breach和scratch缺陷类, 以及广泛分布的rust缺陷存在难检测、漏检、确信度不高等问题, 开启Soft-CIoU-NMS情况下, 本文测试了加入GhostConv+C3STR改进组合前后的模型实际检测效果, 如图13所示, 第1行为改进前检测效果图, 第2行为改进后检测效果图. 得益于Ghost Conv在特征提取层收获更大的感受野, 图像边缘处的缺陷检出率提高, 配合Swin Transformer在大尺寸整张图片的全局注意力运算时移动窗口的全局特征抽取的能力, 模型的下游部分在密集型预测型任务上置信度得到改善.

图 13 消融实验对比图

为更直观地展示Efficient-YOLO模型目标检测能力, 选取了部分实验结果, 如图14所示.

图 14 模型检测效果展示

4 结论

本文针对轴承缺陷检测在工业应用中的需求, 提出了一种轻量化自适应特征融合网络Efficient-YOLO. 首先采用改进的EfficientNetV2作为骨干网络, 大幅减少了参数量, 其次在特征融合网络使用了嵌入CBAM注意力机制的双向加权特征融合网络, 有效提高模型对浅层和深层语义信息的融合能力, 然后引入C3STR配合Ghost卷积, 提高对小目标的检测能力, 最后在后处理阶段使用了Soft-CIoU-NMS, 提高了检测精度. 实验证明, 本文方法在保持轴承外观缺陷检测精度的同时实现了检测模型的轻量化, 并且性能表现优于现有基准模型. 对比于原模型, 改进后的模型的参数量下降了71.7%, 计算量减少了55.7%, 检测速度提升了34.5%, 精度提高了2.5%, 进一步满足了对轴承表面缺陷检测的需求. 下一步的研究方向将聚焦于如何在嵌入式平台上部署低精度运算, 并探索边缘计算的应用.

参考文献
[1]
Zhao Y, Chen BL, Liu BS, et al. GRP-YOLOv5: An improved bearing defect detection algorithm based on YOLOv5. Sensors, 2023, 23(17): 7437. DOI:10.3390/s23177437
[2]
Li B, Gao QJ. Defect detection for metal shaft surfaces based on an improved YOLOv5 algorithm and transfer learning. Sensors, 2023, 23(7): 3761. DOI:10.3390/s23073761
[3]
Niu J, Li HY, Chen X, et al. An improved YOLOv5 network for detection of printed circuit board defects. Journal of Sensors, 2023, 2023: 7270093.
[4]
Chen W, Huang ZT, Mu Q, et al. PCB defect detection method based on Transformer-YOLO. IEEE Access, 2022, 10: 129480-129489. DOI:10.1109/ACCESS.2022.3228206
[5]
Qian XH, Wang X, Yang SY, et al. LFF-YOLO: A YOLO algorithm with lightweight feature fusion network for multi-scale defect detection. IEEE Access, 2022, 10: 130339-130349. DOI:10.1109/ACCESS.2022.3227205
[6]
Ma NN, Zhang XY, Zheng HT, et al. ShuffleNet v2: Practical guidelines for efficient CNN architecture design. Proceedings of the 15th European Conference on Computer Vision (ECCV). Munich: Springer, 2018. 122–138.
[7]
Zhao H, Wan F, Lei GB, et al. LSD-YOLOv5: A steel strip surface defect detection algorithm based on lightweight network and enhanced feature fusion mode. Sensors, 2023, 23(14): 6558. DOI:10.3390/s23146558
[8]
Fan BB, Li WX. Application of GCB-Net based on defect detection algorithm for steel plates. https://www.researchsquare.com/article/rs-1550068/v1. (2022-04-21).
[9]
Xue GH, Li SX, Hou P, et al. Research on lightweight YOLO coal gangue detection algorithm based on ResNet18 backbone feature network. Internet of Things, 2023, 22: 100762. DOI:10.1016/j.iot.2023.100762
[10]
Han K, Wang Y, Tian Q, et al. GhostNet: More features from cheap operations. Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020. 1580–1589.
[11]
Neubeck A, van Gool L. Efficient non-maximum suppression. Proceedings of the 18th International Conference on Pattern Recognition (ICPR2006). Hong Kong: IEEE, 2006. 850–855.
[12]
Lin TY, Dollár P, Girshick R, et al. Feature pyramid networks for object detection. Proceedings of the 2017 IEEE Conference on Computer Vision and Pattern Recognition. Honolulu: IEEE, 2017. 936–944.
[13]
Liu S, Qi L, Qin HF, et al. Path aggregation network for instance segmentation. Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City: IEEE, 2018. 8759–8768.
[14]
Tan MX, Le QV. EfficientNetV2: Smaller models and faster training. Proceedings of the 38th International Conference on Machine Learning. PMLR, 2021. 10096–10106.
[15]
Tan MX, Pang RM, Le QV. EfficientDet: Scalable and efficient object detection. Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Seattle: IEEE, 2020. 10778–10787.
[16]
Hu J, Shen L, Sun G. Squeeze-and-excitation networks. Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City: IEEE, 2018. 7132–7141.
[17]
Woo S, Park J, Lee JY, et al. CBAM: Convolutional block attention module. Proceedings of the 15th European Conference on Computer Vision (ECCV). Munich: Springer, 2018. 3–19.
[18]
Liu Z, Lin YT, Cao Y, et al. Swin Transformer: Hierarchical vision Transformer using shifted windows. Proceedings of the 2021 IEEE/CVF International Conference on Computer Vision (ICCV). Montreal: IEEE, 2021. 9992–10002.
[19]
Bodla N, Singh B, Chellappa R, et al. Soft-NMS-improving object detection with one line of code. Proceedings of the 2017 IEEE International Conference on Computer Vision. Venice: IEEE, 2017. 5562–5570.
[20]
Ren SQ, He KM, Girshick R, et al. Faster R-CNN: Towards real-time object detection with region proposal networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017, 39(6): 1137-1149. DOI:10.1109/TPAMI.2016.2577031
[21]
Liu W, Anguelov D, Erhan D, et al. SSD: Single shot multibox detector. Proceedings of the 14th European Conference on Computer Vision. Amsterdam: Springer, 2016. 21–37.
[22]
Howard A, Sandler M, Chu G, et al. Searching for MobileNetV3. Proceedings of the 2019 IEEE/CVF International Conference on Computer Vision. Seoul: IEEE, 2019. 1314–1324.
[23]
Hou QB, Zhou DQ, Feng JS. Coordinate attention for efficient mobile network design. Proceedings of the 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Nashville: IEEE, 2021. 13708–13717.
[24]
Wang QL, Wu BG, Zhu PF, et al. ECA-Net: Efficient channel attention for deep convolutional neural networks. Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Seattle: IEEE, 2020. 11531–11539.
[25]
Li YH, Yao T, Pan YW, et al. Contextual Transformer networks for visual recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022, 45(2): 1489-1500.