基于改进YOLOv5的超声钢轨缺陷检测方法

Ultrasonic Rail Defect Target Detection Based on Improved YOLOv5

摘要: 钢轨内部缺陷检测是保证列车运行安全非常重要的安全保证，为了增加对钢轨内部缺陷的识别准确率，减少人为不确定因素的干扰，提出了一种用YOLOv5目标识别的方法来辅助识别，并对YOLOv5算法进行改进，改进了数据增强，提升数据鲁棒性，添加NAM注意力模块来捕获局部信息，同时引入GhostNet，对模型进行轻量化处理。改进的YOLOv5_GNM算法的AP50值到达了87.7%，参数量减少了37%，帧数达到了66帧，提高了缺陷识别的准确率。

Abstract: Rail internal defect detection is a very important safety guarantee to ensure the safety of train operation. In order to increase the recognition accuracy of rail internal defects and reduce the interference of human uncertainties, a YOLOv5 target recognition method is proposed to assist the recognition, and YOLOv5 algorithm is improved to improve data enhancement, improve data robustness, add NAM attention module to capture local information, introduce GhostNet, and lightweight the model. The AP50 value of the improved YOLOv5_GNM algorithm reaches 87.7%, the number of parameters decreases by 37%, and the number of frames reaches 66, which im-proves the accuracy of defect recognition.

1. 引言

随着国家的工业发展，钢轨等材料产品的质量和安全也越发引起人们的重视。钢轨等零部件的内部缺陷检测识别，已有许多方法，但是考虑不损坏材料本身等方面，无损检测技术表现出了独特的优势。在无损检测领域，超声检测操作便捷、成本低，同时超声相控阵以其高精度、无损、高安全性等优点，被广泛应用于交通运输等领域 [1] 。超声相控阵技术是超声检测非常重要的一部分，在朝着高实时性，高智能化方向发展 [2] ，然而以往的超声相控阵图像多为由技术人员进行判断，需要一定的技术门槛。基于深度学习的机器视觉技术因高效性和准确性，用于缺陷检测中可以在一定程度上提升检测的效率 [3] 。

深度学习的目标检测框架被广泛划分为两种类型：两阶段检测框架和一阶段检测框架。两阶段检测算法是一种基于检测框和分类器的二阶段算法，最早的二阶段网络是2014年Girshick等人 [4] 提出的R-CNN，在VOC系列数据集上达到了53.7%的准确率。随后，Fast R-CNN [5] 在前者基础上增设了ROI pooling，在主干网络上提取特征图进行分类回归，减少了计算量。2015年提出的Faster R-CNN [6] 是最典型的二阶段网络，它能够自动提取目标区域，计算性能大幅度提升。一阶段算法则包括SDD [7] 和YOLO [8] [9] [10] [11] 系列算法等，这些算法的速度很快，可以实现实时检测。当然，学者们也对YOLOv5进行了很多改进方案。仇娇慧等设计了一种基于通道和空间的注意力机制的齿轮表面缺陷检测方法，在YOLOv5s网络模型的基础上引入卷积注意力模块，增强缺陷区域特征并抑制背景区域特征，提高小目标的检测精度 [12] 。王淑青等对YOLOv5模型输入端改进了一种动态反馈多尺度训练的数据增强方法以提高特征训练精度，然后用ELU函数替换主干网络中的激活函数，加快训练速度 [13] 。将超声相控阵图像缺陷检测技术与深度学习紧密联系起来，可以准确地检测各种缺陷的大小和位置，减少一些不稳定因素的干扰，以提高超声相控阵图像缺陷识别的准确率，使设备的可靠性得到进一步加强 [14] 。吉春生提出了基于深度学习的方法来识别相控阵超声检测图谱中的缺陷，基于Faster R-CNN网络框架，根据焊缝缺陷的特点，使用带有特征金字塔网络的ResNet50作为特征提取的主干，以此来更好地检测缺陷图谱中的小缺陷 [15] 。综上所述，本文提出一种改进的YOLOv5网络框架，对钢轨缺陷的超声相控阵图像进行精准识别。

2. YOLOv5算法概述

YOLOv5算法是基于YOLOv4的改进版本，经过多次改进后，YOLOv5的检测准确度、实时性上提升很大，这对于钢轨的实时成像检测具有非常重要的意义。YOLOv5网络模型分为4个部分，其网络结构如图1所示。主要结构从Input输入到Detect输出，下方是对各个结构的细分。

Figure 1. YOLOv5s structure

图1. YOLOv5s结构图

2.1. Input端

Input端包含Mosaic数据增强、自适应锚框计算和自适应图片缩放三个部分。Mosaic数据增强是使用裁剪缩放等方式将4张图片拼接；自适应锚框计算则是针对不同的数据集设定初始的锚框，其初始框基础上输出预测框，计算差距后再反向更新进行迭代；自适应图片缩放则是将图像都缩放成统一的尺寸。

2.2. Backbone

Backbone包括Focus结构和CSPNET结构。Focus将输入的608 × 608 × 3的图像进行切片操作，然后再经过32个卷积核的卷积后得到304 × 304 × 32的特征图。

2.3. Neck

Yolov5采用FPN (Feature Pyramid Networks，特征金字塔网络) + PAN (Pyramid Attention Network，金字塔注意力网络)的结构，其结构如图2所示。

Figure 2. Neck structure

图2. Neck结构图

2.4. 输出端

Prediction包括Bounding box损失函数和NMS非极大值抑制。YOLOv5使用CIOU_Loss作为Bounding box的损失函数，提高了预测框回归的速度和精度。

3. YOLOv5算法改进

3.1. 数据增强改进

许多的神经网路的参数都是非常多而复杂的，要使得这些参数可以正确工作运行起来则需要大量的数据进行模型训练。获取数据这项工作通常是要花费很多成本，但是通过使用数据增强，对数据训练样本进行缩放、平移、翻转、色彩变换等可以降低模型对其中一些元素的依赖，从而提高模型的泛化能力。

Mosaic数据增强的方法，丰富数据集的同时，也在一定程度上提升了网络的训练速度，如图3所示。

Figure 3. Data enhancement flowchart

图3. 数据增强流程图

由于钢轨实验样本数量少的原因，本文改进Mosaic数据增强方式，将原来的4张图片进行修改，从而对16张图片进行裁剪，并随机缩放，随机排列形成一张新图片。进一步增加了鲁棒性，效果如图4所示。

Figure 4. Mosaic-16 effect picture

图4. Mosaic-16效果图

3.2. 引入注意力机制

在钢轨超声检测中，由于需要提高对重要的缺陷信息的关注度，降低对一些杂波的等其他信息注意力，所以引入注意力机制。

之前的很多的研究都在考虑如何通过注意力算子来获取显著性的特征，但是很少有研究去考虑权值的贡献因子的影响，其中这个因子可以进一步抑制某些不被考虑的特征。本文通过在网络结构中使用一种基于归一化的注意力模块——NAM模块，用权值的贡献因子来提升注意力机制的效果。NAM是Normalization-Based Attention Module的缩写，是CBAM模块的变体 [16] ，它使用了Batch Normalization (BN)比例因子来表示权值的重要性，但是并不会像SE，BAM和CBAM这些结构增加网络的全连接层和卷积层，这样可以尽量少增加模型的计算量。

如图5所示，它包括通道注意子模块和空间注意子模块，可以通过调整通道和空间维度中训练权重的方差测量来重新加权注意力。

$B_{out}=BN\left(B_{in}\right)=\gamma \frac{B_{in}-{\mu }_b}{\sqrt{{\sigma }_b^2+\epsilon }}+\beta$ (1)

其中 $\gamma$ 和 $\beta$ 分别表示可训练的变换参数(尺度和位移)；而 ${\mu }_b$ 和 ${\sigma }_b$ 分别是每个batch b中的平均值和标准差；其中 ${\mu }_b$ 是均值， ${\sigma }_b^2$ 是BN中的方差，方差越大，意味着变化越大，信息越丰富。对于通道注意力模块，基于归一化相关权重 $W_{\gamma }$ 可以将更多的注意力集中在重要通道，同时抑制信息量较小的权重。假设 $F_1\in R^{H\times W\times C}$ 是输入特征图，其中H、W和C表示高度，宽度和通道数。 $\gamma$ 是每个通道的比例因子，权重为 $W_{\gamma }$ ，则通道注意力的输出 $M_C$ 可以表示为：

$M_c=sigmoid\left(W_{\gamma }\left(BN\left(F_1\right)\right)\right)$ (2)

Figure 5. Normalization-based attention module

图5. NAM模块

Figure 6. Channel attention module

图6. 通道注意力模块

Figure 7. Spatial attention module

图7. 空间注意力模块

通道注意力模块如图6所示，空间注意力模块如图7所示。以通道注意模块的方式，其空间注意模块将BN应用于空间维度中的像素，进行像素归一化(PN)。它根据比例因子 $\lambda$ 将更多的注意力集中在在信息量更大的像素上，并调整相关的权重W。同样， $F_2\in R^{H\times W\times C}$ 是输入特征图，输出空间注意力模块的输出 $M_s$ 是：

$M_s=sinmoid\left(W_{\lambda }\left(PN\left(F_2\right)\right)\right)$ (3)

同时为了抑制不太重要的权重，向损失函数中增加正则化项。

$Loss={\displaystyle {\sum }_{\left(x,y\right)}l\left(f\left(x,W\right),y\right)}+p{\displaystyle \sum g\left(\gamma \right)}+p{\displaystyle \sum g\left(\lambda \right)}$ (4)

其中， $l(\cdot )$ 和 $g(\cdot )$ 分别表示损失函数和范数惩罚函数；x和y分别为输入和输出；W是网络权重；p是平衡惩罚因子。

在每个网络block的末端嵌入一个NAM模块，它重新处理并合理化从主干中提取的重要特征，以方便对象在下一步中预测。这样可以使模型更加关注图中的关键信息，提升计算效率，但由于NAM的轻量化，增加的计算成本并不多。

Figure 8. Ghost module

图8. Ghost模块

3.3. 模型轻量化

在实际的超声检测设备中，为了满足三防等要求，硬件的配置不是很高，这就要求后续附加的目标检测模型的计算量有一定的要求。因此对模型进行轻量化处理，引入GhostBottleneck模块，它来源于GhostNet [17] ，如图8所示，能够做在最大限度减少网络计算损耗的同时提高网络的检测准确率假设正常卷积是为32个卷积核的卷积计算，在Ghost模块的第一阶段会采用其原有卷积的一半进行计算，在第二阶段中 $\phi$ 是3*3的小卷积，通过3*3的这些小卷积对第一阶段所得的特征图进行逐层卷积，来进一步优化之前卷积的计算损耗。

GhostBottleneck模块作为一个即插即用模块，可以保证充分的增值输入数据。此模块包含两种结构，一种是Stride为1的GhostBottleneck模块，由两个堆叠的Ghost模块组成，另一种是Stride为2的GhostBottleneck模块，中间连接了一个DW卷积层，如图9所示。其中第一个Ghost模块的作用是增加输入特征图的通道数，为后面的操作来做扩展处理，而第二个Ghost模块的作用是减少输出特征图的通道数使其与网络中的直径结构相匹配。本文使用GhostBottleneck来替换原本YOLO网络主干部分中的 Bottleneck CSP结构来优化网络的参数量和计算规模，减少模型的计算量。

Figure 9. Two GhostBottleneck module

图9. 两种GhostBottleneck模块

4. 实验

4.1. 实验环境

本文采用自制钢轨缺陷超声相控阵数据集进行训练，包含1500张图片，搭建PyTorch深度学习平台，具体测试环境：显卡为NVIDIA RTX 2070SUPER，显存为8 G，CUDA版本10.2，编译语言为Python3.7，批处理量batch size设置为16，初始学习率设置为0.01，动量系数设置为0.937，权值衰减系数设置为0.0005。一共训练300个epochs。训练集80%，验证集10%，数据集10%。

4.2. 评价标准

对于该算法的评价标准，本文使用了结合召回率和精准率的mAP (mean Average Precision)对改进前和改进后的算法性能进行评价。其中召回率表示为在所有正类样本中的检出率，其计算公式如下

$R=\frac{TP}{TP+FN}\times 100\%$ (5)

其中的R代表召回率，TP表示算法将样本正类预测为正类的个数；FN则是表示了将样本正类预测为负类的个数。

精准率P可以表示未预测的正样本占所有正样本的比例，计算公式如下：

$P=\frac{TP}{TP+FP}\times 100\%$ (6)

PR曲线是以P为纵轴，每种类别通过PR曲线积分

$AP={\displaystyle {\int }_0^{1}p\left(r\right)\text{d}r}$ (7)

得到AP，对所有类别的AP值求平均值即为mAP

$mAP=\frac{{\displaystyle {\sum }_{i=1}^{N}AP}}{N}$ (8)

4.3. 对比实验

对原YOLOv5s和改进后的YOLOv5s进行实验对比，做4次实验，如下所示。

1) 原YOLOv5s

2) YOLOv5s_M：改进数据增强。

3) YOLOv5s_G：引入GhostBottleneck模块。

4) YOLOv5s_N：引入NAM模块。

5) YOLOv5_GNM：引入NAM模块、GhostBottleneck模块和改进数据增强。

实验结果如表1和表2所示：

Table 1. Comparison table of model complexity

表1. 模型复杂度对比表

Table 2. Model performance comparison table

表2. 模型性能对比表

由表可知，YOLOv5s_M算法对mosaic数据增强进行改进后，虽然模型大小有了略微的增加，但mAP增长了0.8%，证明了改进的可行性。YOLOv5s_G算法进行GhostBottleneck结构改进后，模型的参数数量减少41%，平均计算时间少了5 ms，帧率从52.6帧提升到71.4帧，提升了35.8%。但mAP只减少了1.2%，用少量的性能损失换来大量的参数量的降低，实现了轻量化，这是由于GhostNet使用了线性运算替代部分卷积过程，因此在其得到近似于原始特征图的同时减少了大量的计算量。YOLOv5s_N算法增加了NAM模块，增加了注意力机制，检测速度虽然增加了1 ms，但是mAP提升了1.7%，表现了注意力机制对于YOLOv5目标检测精确度有加强作用。YOLOv5s_GNM算法则是在检测精度上提升1.1%的同时，模型参数减少了37%，帧率也从52.6帧提升至66.7帧，实现了模型的轻量化。

改进前的YOLOv5s网络模型和改进后的YOLOv5s_GNM网络模型的检测结果如图10所示，Interference的置信度从0.85提升到0.88，缺陷Hole的置信度从0.84提升到0.89，改进后的检测精度有明显提升。

(a) (b)

Figure 10. Comparison of test results. (a) Before improvement; (b) After improvement

图10. 检测结果对比。(a) 改进前；(b) 改进后

5. 结论

本文提出了一种改进YOLOv5算法，并将其融合到超声相控阵钢轨缺陷检测中，可以辅助专业人员通过超声相控阵仪器检测钢轨缺陷。结果表明：

1) 通过改进Mosaic数据增强丰富了数据集，增加小样本目标，实验发现，模型的检测精度有了小幅度的上升。

2) 在YOLOv5s的网络结构中引入注意力机制，融入NAM模块，以此来强化重要的缺陷特征，在少量增加模型计算量的基础上增加了模型检测的准确度。

3) 在YOLOv5s网络结构中引入GhostBottleneck结构，模型的参数量和计算量大幅度减少，实现了模型的轻量化，方便后续模型部署至工控机或嵌入式设备。

基金项目

QNTD202108轨道交通智能运维青年创新团队；12172211非饱和复合地基固结理论构建与解析模型研究。