基于MobileNet的移动端列车图像故障检测算法
2023-06-25周鹏张龙信
周鹏 张龙信
摘 要:针对列车故障检测效率低的问题,提出一种基于MobielNet的移动端列车图像故障检测算法。首先,在MobileNet中引入注意力卷积块和Ghost模块,用以提升网络的学习能力。其次,使用残差聚合网络获取多层次的特征图。最后,将该模型移植到移动端设备上完成列车故障检测任务。实验结果表明,该算法的平均精度均值达到了85.35%,与YOLOv3-Tiny、YOLOv4-Tiny、YOLOX、YOLOv5相比,mAP分别提高了8.83%、5.49%、7.89%、5.31%,并且FED拥有更低检测延迟。
关键词:列车故障检测;目标识别;MobileNet;移动设备;注意力机制
中图分类号:TP391.4 文献标识码:A 文章编号:2096-4706(2023)06-0046-05
Mobile Terminal Train Image Fault Detection Algorithm Based on MobielNet
ZHOU Peng, ZHANG Longxin
(Hunan University of Technology, Zhuzhou 412007, China)
Abstract: To solve the problem of low efficiency of train fault detection, mobile terminal train image fault detection algorithm based on MobielNet is proposed. First, attention convolution block and Ghost module are introduced into MobileNet to improve the learning ability of the network. Secondly, residual aggregation network is used to obtain multi-level feature map. Finally, the model is transplanted to the mobile terminal equipment to complete the train fault detection task. The experimental results show that the average accuracy of the algorithm reaches 85.35%. Compared with YOLOv3-Tiny, YOLOv4-Tiny, YOLOX and YOLOv5, mAP improves 8.83%, 5.49%, 7.89% and 5.31% respectively, and FED has lower detection delay.
Keywords: train fault detection; target recognition; MobileNet; mobile device; attention mechanism
0 引 言
中國铁路高速发展,铁路安全问题一直是人们关注的焦点,当前对于列车零部件故障的检测主要依赖人工识别,检测效率低,易受检测人员自身状态和自然环境的影响。近年来基于机器视觉的列车故障检测方法在我国铁路系统中得到了广泛的应用,但由于列车故障种类多,识别难度高,传统的基于机器视觉的图像处理技术难以完成列车故障检测任务。随着计算机视觉技术的发展,使得基于深度学习的列车故障检测方法越来越受到人们的重视。相比于传统列车故障检测方法,基于深度学习的列车故障检测方法能自动的设计和提取特征,提高检测效率。针对这些因素,本文提出一种基于MobielNet[1]的移动端列车图像故障检测算法FED,旨在提高列车故障检测精度,完成实时检测任务。
本文工作的主要贡献总结如下:
(1)设计了残差聚合网络(Residual Pixel Aggregation Network, Residual-PAN),用于产生多层次的特征图,并进行特征融合,显著提高了网络特征提取能力。Residual-PAN使用1×1的卷积操作统一特征图的通道数,减少网络参数,并使用卷积核为5×5(或者7×7)的深度可分离卷积扩大感受野。
(2)提出了Enhanced MobileNet(EMNet)。在MobileNet的基础上,增加了卷积注意力机制(Convolution Block Attention Module, CBAM)和Ghost模块,可以在少量参数的情况下,获取更多的特征信息,提高网络的学习能力。
(3)将Fastand Effective Detection(FED)模型应用于移动端的列车零部件故障检测,并进行了一系列实验以评估其有效性。实验结果表明,FED具有更高的检测精度、较低的检测延迟和较少的参数量,能在硬件性能较差的移动设备完成实时检测任务。
本文其余工作安排如下:在第1节回顾了目标检测算法的相关工作,第2节中介绍了提出的FED模型。第3节描述了实验和结果分析,第4节对所做的工作进行了总结。
1 相关工作
1.1 移动端目标识别
移动端目标识别是指使用轻量级的移动设备实现目标识别任务,其流程如图1所示。首先,使用列车图像训练检测模型。然后,将检测模型部署到移动设备,并使用自带的摄像头收集实时图像和视频进行列车故障检测。最后,展示检测结果。
1.2 目标检测方法
Mao等人提出了Min-YOLOv3[2]模型,在YOLO9000的基础上,YOLOv3重新设计了Darknet-53残差网络,通过引入ResNet网络中的残差结构,解决网络结构层数太多而出现的梯度消失问题,降低训练深层网络的难度。并使用特征金字塔网络(Feature Pyramid Network, FPN)进行多尺度融合预测。FPN网络用于融合不同尺度的特征信息,高层特征具有较多的语义信息,而低层特征有较丰富位置信息,将不同层的特征进行融合,可以最大化利用特征信息。Bochkovshiy等人设计了YOLOv4[3]模型,在YOLOv3基础上,YOLOv4将Darknet-53替换为CSPDarknet-53,并加入了路径聚集网络、Mish激活函数和Mosica数据增强等方法,使得YOLOv4模型的检测精度和检测速度达到了一个全新的高度。Zhu等人提出了YOLOv5[4],在YOLOv4基础上,YOLOv5使用Mosaic和图片自适应放缩法对数据集进行优化,并使用自适应锚框计算法设置初始锚框的大小,提升检测速度。YOLOv5在主干网中加入Focus结构,并加入特征金字塔和PAN提高检测进度。Panboonyuen等人提出了YOLOX[5],YOLOX添加Mosaic和MixUp两种更加强大的数据增强方法,并提出了SimOTA标签分配算法,极大地提高了模型的检测精度。YOLOX与YOLOv5相比,在参数量的基本相同情况下,YOLOX能在COCO-2017上mAP达到50.00%(比YOLOv5高出1.8%的mAP),且YOLOX的检测速度达到68.9 FPS。
1.3 注意力机制
注意力机制是机器学习中一种数据处理方法,广泛用于在自然语言处理、图像识别及语音识别等各种不同类型的机器学习任务中。Li等人提出了通道注意力机制Squeeze-and-ExcitationNetwork(SENet)[6]。SENet分为3个步骤。第一步输入一个H×W×C(H表示特征圖的高,W表示特征图的宽,C表示特征图的通道数)的特征图F进行全局平均池化,得到1×1×C的特征向量;第二步输入1×1×C的特征向量,在通道维度进行经过压缩和扩张,得到1×1×C的权重系数,并使用sigmoid激活函数将权重系数的值都限制在0到1的范围;第三步输入特征图F与权重系数,两者相乘得到加权后的特征图。Woo等提出了混合域注意力机制CBAM(Convolutional Attention Module),在SENet的基础上增加了空间注意力机制,加强了特征图局部间的联系。CBAM是一个轻量级的通用模块,可以无缝地集成到任何CNN架构中,并且可以与CNN一起进行端到端的训练。
2 推荐框架
2.1 整体框架
本节展示了所提出的FED模型,其结果如图2所示(其中Conv表示卷积操作,Upsample表示上采样,DSConv表示深度可分离卷积)。FED模型由EMNet、Residual-PAN和YOLOHead组成。EM是轻量化的网络,能高效的提取图片特征。Residual-PAN用于获取多层次特征映射,并对相邻的特征图进行融合。YOLOHead对目标进行回归和分类。
2.2 Enhanced MobileNet
本文提出了EMNet,如图3所示,其中pw conv表示点卷积,dw conv表示深度卷积。EMNet主要由EM块组成,EM块集成Ghost模块和CBAM模块,EMNet块结构如图3所示。Ghost模块由点卷积和深度卷积组成,Ghost可以用少量的参数生成大量的特征信息,提高网络的学习能力。CBAM是轻量级注意力机制,能轻易地嵌入到神经网络。CBAM通过在特征图的通道维度和空间维度进行加权,以获取图片中更重要特征信息。CBAM使用的激活函数是Relu和H-Sigmoid。
2.3 Residual-PAN
本文使用Residual-PAN获取多层次的特征图,并对相邻特征图进行融合。传统的PAN中,输出特征图的通道数与骨干网产生特征图通道数保持一致,需要昂贵的计算成本,不适用于移动设备。Residual-PAN通过1×1的卷积操作使所有特征图的通道数都等于特征图中最小的通道数,并使用残差结构融合相邻的特征图,使得计算成本大幅降低。此外,Residual-PAN额外增加了一个特征图,以检测更多的对象。Residual-PAN通过调整DSConv卷积核(5×5/7×7)的尺寸以获取更大的感受野。Residual-PAN使得模型的参数更少,检测精度更高。
2.4 标签分配策略与损失函数
正样本与负样本的标签分配策略对识别模型的效果影响巨大。FED采用SimOTA策略进行标签分配。SimOTA策略会随着训练过程的不断变化而变化,自动分析正样本的数量,且速度更快。对于分类,ZP使用Varifocal损失函数。对于回归,ZP使用GIoU和Distribution Focal损失函数。计算公式为:
loss=lossvfl+αlossgiou+βlossdfl
其中,lossvfl表示Varifocal损失函数,lossgiou表示GIoU损失函数,lossdfl表示Distribution Focal损失函数,α、β表示常量参数。
3 实验与结果
本节首先描述了实验中的数据集,然后给出了网络结构设置和性能评价指标。最后展示了实验结果。
3.1 数据集描述
为验证算法的有效性,本文采用大量列车零部件图像来构建实验的Train TD(Train Test Dataset),所有图像的大小为930×680像素。训练集拥有13 581张图片,测试集有8 366张列车零部件图像。列车典型零部件主要有3种,分别为截断塞门把手(Truncated Plug Door Handle, TPDH)、上拉杆(Upper Lever, UL)和紧锁板(Locking Plate, LP),如图4所示,左侧为故障的零部件,右侧为正常列车零部件。列车故障检测需要对以上3种零部件进行精确的定位,然后再根据这些零部件的状态进行判别,分析是否出现故障,得到最终的检测结果。
3.2 模型参数设置
本文的训练实验环境为Ubuntu 20.04系统,显卡为4块NVIDIA GeForce GTX 1080Ti,显存为22 GB。网络结构使用Pytorch框架实现。学习率初始化为0.001,会随着余弦退火函数发生变化,训练周期为300。测试环境为高通骁龙865。
3.3 评价指标
主要评价指标如下:
(1)Precision表示模型预测出的所有目标中正确的比例,体现模型对负样本的区分能力。其计算公式为:
其中TP(True Positives)表示被正确识别的正样本的数量,FP(False Positives)表示被错误识别为正样本的负样本数量。Precision越高,模型对负样本的区分能力越强。
(2)Recall表示被模型正确识别的目标占总目标数量的比例,体现模型对正样本的识别能力。其计算公式为:
其中FN(False Negatives)表示被错误识别为负样本的正样本数量。Recall越高,模型对正样本的识别能力越强。
(3)mean Average Precision(mAP)平均精度均值,即AP的平均值。它是目標检测算法的主要评估指标。AP是Precision-Recall曲线下面积,其计算公式为:
其中P(r)表示Precision-Recall曲线,mAP的计算公式为:
其中N表示检测故障类别总数,APs对应目标的AP值。mAP值越高,表明该目标检测模型在给定数据集上的检测效果越好。
(4)FLOPs(Floating Point Operations)浮点运算数,表示模型计算量,用于衡量模型的复杂度。FLOPs值越低,模型的计算量越小,检测速度越快。
3.4 实验结构与分析
在COCO-2017数据集上对FED模型进行消融实验,实验结果如表1所示。当FED只使用EM和PAN,mAP为44.63%,参数量为2.98 M。当FED用Residual-PAN替换PAN时,mAP增加到48.42%,我们推测是因为Residual-PAN能获取多层次的特征图,并利用残差块融合相邻特征,防止梯度消失,提升检测精度。当FED增加CBAM后mAP增加了4.71%,CBAM是一种轻量级注意力机制,能提升关键特征的权重,是模型更加关注图像中的重要信息。当FED增加Ghostblock后mAP提升至47.81%,Ghostblock只需要少量的参数就能得到大量的特征信息,提高网络的学习能力。
从表2可以看出,FED的检测精度和检测速度远超其他检测模型。我们认为主要有以下几个原因:(1)FED的主干网EM参数更小,在引入CBAM和Ghost模块后,网络会聚焦的关键信息,学习能力更强。(2)Residual-PAN能取得多层次的特征检测更多目标,并使用残差块进行特征融合,能有效地防止反向传播过程中梯度消失的问题。Residual-PAN还会使用大尺寸的卷积核进行深度可分离卷积,扩大感受野。(3)FED使用SimOTA动态分配标签,正负样本分配更准确,速度更快。
在Traindataset测试FED和其他先进目标检测模型,实验结果如表3所示。FED的mAP为85.35%,比YOLOv3-Tiny、YOLOv4-Tiny、YOLOX、YOLOv5分别提升了8.83%、5.49%、7.89%、5.31%。FED拥有更快、更高效的特征提取网络,Residual-PAN能丰富特征信息,使得FED取得更优的检测效果。实际的检测效果如图5所示,FED可以快速地找出图片中故障的零部件,并对故障的零部件位置和故障类别进行标注,给出该故障的置信度,方便工作人员进行检测和修复工作。
4 结 论
相比于传统的列车故障检测算法,基于深度学习的列车故障检测方法检测精度更高,检测速度更快,但需要强大的算力支持,需要部署在高性能的服务器上,对检测环境要求苛刻。本文实现了一种基于MobielNet的移动端列车故障检测算法FED,在Train TD上,mAP达到了85.35%,检测延迟为17.87 ms,满足实时性和精确性的要求。
参考文献:
[1] EL-SAADAWY H,TANTAWI M,Shedeed H A,et al. A Two-Stage Method for Bone X-Rays Abnormality Detection Using MobileNet Network [C]//AICV: The International Conference on Artificial Intelligence and Computer Vision.Cairo:Springer,2020,1153:372-380.
[2] MAO Q C,SUN H M,LIU Y B,et al. Mini-YOLOv3: Real-Time Object Detector for Embedded Applicatio [J].IEEE Access,2019,7:133529-133538.
[3] BOCHKOVSKIY A,WANG C Y,LIAO H Y M. YOLOv4: Optimal Speed and Accuracy of Object Detection [J/OL]. arXiv:2004.10934 [cs.CV].[2022-09-28].https://arxiv.org/abs/2004.10934.
[4] ZHU X K,LYU S C,WANG X,et al. TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios [C]//2021 IEEE/CVF International Conference on Computer Vision Workshops(ICCVW).Montreal:IEEE,2021:2778-2788.
[5] PANBOONYUEN T,THONGBAI S,WONGWEERANIMIT W,et al. Object Detection of Road Assets Using Transformer-Based YOLOX with Feature Pyramid Decoder on Thai Highway Panorama [J/OL].Information,2022,13(1)[2022-09-28].https://doi.org/10.3390/info13010005.
[6] LI Y,LIU Y,CUI W G,et al. Epileptic Seizure Detection in EEG Signals Using a Unified Temporal-Spectral Squeeze-and-Excitation Network [J].IEEE Transactions on Neural Systems and Rehabilitation Engineering,2020,28(4):782-794.
作者简介:周鹏(1997—),男,汉族,湖南常德人,硕士在读,研究方向:基于深度学习的列车识别方法;张龙信(1983—),男,汉族,湖南株洲人,副教授,博士,研究方向:高性能计算、机器学习。
收稿日期:2022-10-17