基于注意力机制的接触轨表面缺陷检测模型

孙 明

（上海地铁维护保障有限公司供电分公司，上海 201106）

0 引言

随着我国城市轨道交通的飞速发展，对城市轨道交通的安全运行也提出了更高的要求，稳定、安全和高效的电源供应是其中至关重要的一环。随着城市轨道交通供电技术的不断发展，目前城市轨道交通供电系统主要可以分为接触网供电和接触轨供电两种供电方式。接触轨供电方式因其美观、建设成本低、供电稳定、检测方便、便于在隧道中使用等优点在线路上得到了越来越广泛的应用，但是由于接触轨长期与集电靴高速摩擦，极易出现磨耗、擦伤、疤痕等病害，会给轨道交通的安全供电造成极大的隐患，因此对接触轨表面缺陷的准确、快速、智能化检测也成为当今城市轨道交通维护保障工作的研究热点问题。近年来，基于计算机视觉的故障检测技术在轨道交通维护保障领域得到了越来越广泛的研究与应用，为了实现对接触轨表面缺陷的快速、准确、自动化检测，本文针对接触轨表面缺陷复杂、精细等特点，在目标检测算法模型YOLOv5 的基础上加入了CBAM（Convolutional Block Attention Module）和CA（Coordinate Attention）注意力机制，设计了基于注意力机制改进的YOLOv5 接触轨表面缺陷检测模型，提升了模型的检测速度与定位精度。

1 相关模型

1.1 YOLOv5 算法模型

为实现对接触轨表面缺陷的快速检测，本文选择使用YOLOv5 算法模型作为接触轨表面缺陷检测基准模型。YOLO算法是一个单阶段目标检测算法，可以端到端地完成对特征图像的训练与预测。YOLO 系列目标检测算法经过不断的发展与创新，现已推出了更加轻量化的YOLOv5 算法模型。YOLOv5 算法模型由输入端、骨干网络、颈部网络和头部网络4个部分组成。根据模型宽度、深度和参数量等不同，YOLOv5 也可以分为5 个版 本：YOLOv5x、YOLOv5l、YOLOv5m、YOLOv5s、YOLOv5n，模 型随着宽度、深度的增加，模型的参数量、层数和每秒执行的浮点运算量（GFLOPs）也在增加，检测精度会随之越高，但所需训练时间也相应增长。本文使用模型为了获取更轻量化的模型，实现更快速、精确的检测，选择了YOLOv5s 作为基准模型。

1.2 CBAM 注意力机制

针对接触轨表面缺陷精细，分布范围有限的特点，在模型中加入了CBAM 注意力机制，提高模型收敛能力与速度。CBAM 主要由通道注意力机制和空间注意力机制两部分组成，其结构见图1。

图1 CBAM 结构示意

在CBAM 中，特征图像输入后会沿着堆叠的两个机制模块分别进行特征提取，依次推断出一个一维的通道注意力特征图和一个二维空间的空间注意力特征图，最后将通道注意力特征图、空间注意力特征图与原图三者相乘完成自适应，完成对图像特征的精炼和对关注面积的扩大，有效提升了卷积神经网络的表示能力。总体过程可概括为原理公式：

其中，⊗表示为元素方向上的乘法，F′是通道注意力机制进行图像特征精炼后的输出，F′则是最终CBAM 对图像特征精炼后的输出。

1.3 CA 注意力机制

为了使模型更好学习到接触轨表面图像特征的长距离依赖关系，获取精确的接触轨表面缺陷特征坐标定位信息，本文在模型中加入了CA 注意力机制。CA 注意力机制针对通道注意力机制难以获取到图像准确坐标定位信息，通过将位置信息嵌入到通道注意力中提出的，其结构如图2 所示。

图2 CA 结构示意

整个CoordAtt 主要可概括为两个处理步骤：坐标信息嵌入与协调注意力生成，总体可概括为公式：

其中，yc（i，j）代表最终CoordAtt 的输出，xc（i，j）代表坐标信息嵌入处理后的输出（i）和（j）分别代表两个方向上生成的协调注意力参数。

1.4 改进模型

结合CBAM 和CCA 的特点，在YOLOv5s 基准网络框架的基础上，在网络模型第5、8、11、28 层分别加入CA 和在第21、25层加入CBAM，设计了结合CBAM 与CA 的基于注意力机制的CB-CA YOLOv5 接触轨表面缺陷检测模型，模型结构如图3 所示。

图3 CB-CA YOLOv5 网络结构示意

2 实验结果与分析

2.1 数据集及实验环境

本文实验所用的接触轨表面数据集来源于上海轨道交通16 号线，由550 张分辨率为3088×2320 接触轨表面图像组成，每张图像至少包含一个表面缺陷，最终放入模型训练、验证和检测的图片尺寸都为608×608。本文使用的接触轨表面缺陷检测模型在Linux-x86_64 系统，NVIDIA Quadro P500 GPU 的服务器上进行实验，每个模型实验训练方式都为单GPU 运算，训练迭代次数（Epoch）为300 次。

2.2 评价指标

本文使用了多个目标检测模型研究中常用的评价指标对模型进行定性、定量的分析与比较：

TP 代表人工标注的接触轨表面缺陷被正确识别个数，FP代表被错误识别为接触轨表面缺陷的个数，FN 代表未被识别出的接触轨表面缺陷的个数。

mAP_0.5：交并比设为0.5 时，计算每一类图片的AP（Average Precision）再对所有类别求平均，因为本文只有裂纹一类，所以本文使用的为AP_0.5。

mAP_0.5：0.95：设置步长为0.05，取交并比在0.5～0.95 之间的平均mAP，因为本文只有裂纹一类，所以本文使用的为AP_0.5：0.95。

FPS（Frames Per Second）：每秒检测图片帧数。

TS代表一张图片的检测时间。

2.3 实验结果分析

将改进后的CB-CA YOLOv5 模型对接触轨表面进行检测，能有效、准确对缺陷进行识别与定位，并且对三种类型的缺陷都有着较好的检测精度与置信度表现。为了进一步对本文使用的模型进行定性、定量的分析，使用式（4）、式（5）计算评价指标，选取了4 种目标检测模型与本文模型进行进一步的比较、分析，实验结果见表1。在YOLOv5s 基础上进行改进的模型，对比YOLOv5s，模型在不损失训练、检测速度的情况下，各项检测指标都实现了大幅提升，精度、AP_0.5：0.95 分别提高了7.3%和5.8%，同时只需要263 次迭代次数就能完成最优模型的训练，模型收敛能力也获得了大幅提高。因此，比较5 种方法，本文提出的基于注意力机制的CB-CA YOLOv5 模型最适合于进行接触轨表面缺陷检测。

表1 算法框架数据对比

3 结论

针对接触轨表面缺陷类型复杂、精细、分布范围有限等特点，本文提出基于注意力机制的CB-CA YOLOv5 接触轨表面缺陷检测模型，实现了对接触轨表面缺陷的准确检测与精确定位。一方面，在模型中加入了CBAM 注意力机制，提高模型的收敛速度与检测性能。另一方面，在模型组加入CA 注意力机制，能获取更精确的接触轨表面缺陷特征坐标定位信息。最终，在本文采集的接触轨表面缺陷图像数据集进行实验，最终模型精确率为0.958，召回率为0.999，AP_0.5 为0.968，AP_0.5：0.95 为0.874，FPS 为167。与其余版本的YOLOv5 模型进行试验对比，本文模型在维持较好的检测速度同时，大幅提高了各项检测性能，并且增强了模型的收敛能力，有较好的实用价值。