改进YOLOv5s的轨道障碍物检测模型轻量化研究

doi:10.3778/j.issn.1002-8331.2208-0045

摘要/Abstract

摘要： 针对传统列车轨道障碍物检测方法实时性差和对小目标检测精度低的不足，提出一种改进YOLOv5s检测网络的轻量化障碍物检测模型。引入更加轻量化的Mixup数据增强方式，替代算法中原有的Mosaic数据增强方式；引入GhostNet网络结构中的深度可分离卷积GhostConv，替代原有YOLOv5s模型中特征提取网络与特征融合网络中的普通卷积层，减小了模型的计算开销；在模型特征提取网络末端加入CA空间注意力机制，让算法在训练过程中减少了重要位置信息的丢失，弥补了改进GhostNet对检测精度的损失；将改进后的模型进行稀疏训练和通道剪枝操作，剪掉对检测精度影响不大的通道，同时保留重要的特征信息，使模型更加轻量化。实验结果表明，改进后的模型在自制的多样化轨道交通数据集上，相较于原始YOLOv5s算法，在模型大小减小9.7?MB，检测速度提高14?FPS的前提下，检测精度提升了1.0个百分点。同时与目前主流的检测算法对比，在检测精度与检测速度上也具有一定的优越性，适用于复杂轨道交通环境下的障碍物目标检测。

关键词: 目标检测, YOLOv5s, GhostNet, 注意力机制, 通道剪枝, 轻量化

Abstract: Aiming at the shortcomings of the traditional train track obstacle detection methods with poor real-time performance and low detection accuracy for small targets, a lightweight obstacle detection model based on improved YOLOv5s detection network is proposed. Firstly, a more lightweight Mixup data enhancement method is introduced to replace the original Mosaic data enhancement method. Secondly, the deep separable convolution GhostConv in the GhostNet network structure is introduced to replace the ordinary convolution layer in the feature extraction network and feature fusion network in the original YOLOv5s model, which reduces the computational overhead of the model. The CA spatial attention mechanism is added to the end of the model feature extraction network, which reduces the loss of important location information in the training process of the algorithm and makes up for the loss of detection accuracy caused by improved GhostNet. Finally, sparse training and channel pruning are performed on the improved model to prune away the channels that have little influence on the detection accuracy, while retaining important feature information to make the model more lightweight. The experimental results show that, compared with the original YOLOv5s algorithm, the model size of the improved model is reduced by 9.7 MB, the detection speed is increased by 14 FPS, and the detection accuracy is improved by 1.0 percentage point on the self-made diversified rail transit dataset. At the same time, compared with the current mainstream detection algorithm, the detection accuracy and detection speed also have some advantages, which is suitable for the obstacle target detection in complex rail transit environment.

Key words: object detection, YOLOv5s, GhostNet, attentional mechanism, channel pruning, lightweight

李昂, 孙士杰, 张朝阳, 冯明涛, 吴成中, 李旺. 改进YOLOv5s的轨道障碍物检测模型轻量化研究[J]. 计算机工程与应用, 2023, 59(4): 197-207.

LI Ang, SUN Shijie, ZHANG Zhaoyang, FENG Mingtao, WU Chengzhong, LI Wang. Research on Lightweight of Improved YOLOv5s Track Obstacle Detection Model[J]. Computer Engineering and Applications, 2023, 59(4): 197-207.

参考文献

[1] 王泉东，杨岳，罗意平，等.铁路侵限异物检测方法综述[J].铁道科学与工程学报，2019，16（12）：3152-3159.
WANG Q D，YANG Y，LUO Y P，et al.Review on railway intrusion detection methods[J].Journal of Railway Science and Engineering，2019，16（12）：3152-3159.
[2] 郭双全，董昱.基于雷达的列车直轨运行前方障碍物检测方法研究[J].铁道科学与工程学报，2020，17（1）：224-231.
GUO S Q，DONG Y.Research on detection method of obstacle in front of straight track operation of train based on radar[J].Journal of Railway Science and Engineering，2020，17（1）：224-231.
[3] 金炳瑞.基于图像处理的铁路轨道异物入侵的自动识别研究[D].兰州：兰州交通大学，2016.
JIN B R.Research on automatic recognition of foreign body invasion in rail based on image processing[D].Lanzhou：Lanzhou Jiaotong University，2016.
[4] REN S，HE K，GIRSHICK R，et al.Faster R-CNN：towards real-time object detection with region proposal networks[C]//Advances in Neural Information Processing Systems，2015，28.
[5] 史红梅，柴华，王尧，等.基于目标识别与跟踪的嵌入式铁路异物侵限检测算法研究[J].铁道学报，2015，37（7）：58-65.
SHI H M，CHAI H，WANG Y，et al.Study on railway embedded detection algorithm for railway intrusion based on object recognition and tracking[J].Journal of the China Railway Society，2015，37（7）：58-65.
[6] 徐岩，陶慧青，虎丽丽.基于Faster R-CNN网络模型的铁路异物侵限检测算法研究[J].铁道学报，2020，42（5）：91-98.
XU Y，TAO H Q，HU L L.Research on railway foreign body intrusion detection algorithm based on Faster R-CNN network model[J].Journal of the China Railway Society，2020，42（5）：91-98.
[7] BOCHKOVSKIY A，WANG C Y，LIAO H Y M.YOLOv4：optimal speed and accuracy of object detection[J].arXiv：2004.10934，2020.
[8] LIU W，ANGUELOV D，ERHAN D，et al.SSD：single shot multibox detector[C]//14th European Conference on Computer Vision.Cham：Springer，2016：21-37.
[9] 刘力，苟军年.基于YOLOv4的铁道侵限障碍物检测方法研究[J].铁道科学与工程学报，2022，19（2）：528-536.
LIU L，GOU J N.Research on detection method of railway intrusion obstacles based on YOLOv4[J].Journal of Railway Science and Engineering，2022，19（2）：528-536.
[10] 郭磊，薛伟，王邱龙，等.一种基于改进YOLOv5的小目标检测算法[J].电子科技大学学报，2022，51（2）：251-258.
GUO L，XUE W，WANG Q L，et al.A small object detection algorithm based on improved YOLOv5[J].Journal of University of Electronic Science and Technology of China，2022，51（2）：251-258.
[11] 邱天衡，王玲，王鹏，等.基于改进YOLOv5的目标检测算法研究[J].计算机工程与应用，2022，58（13）：63-73.
QIU T H，WANG L，WANG P，et al.Research on object detection algorithm based on improved YOLOv5[J].Computer Engineering and Applications，2022，58（13）：63-73.
[12] 杨晓玲，罗顺利，梁皓添.基于YOLOv5的交通路面障碍物目标检测[J].智能城市，2021，7（20）：121-122.
YANG X L，LUO S L，LIANG H T.Obstacle object detection on traffic pavement based on YOLOv5[J].Smart City，2021，7（20）：121-122.
[13] HAN K，WANG Y，TIAN Q，et al.GhostNet：more features from cheap operations[C]//Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition，2020：1580-1589.
[14] HU J，SHEN L，SUN G.Squeeze-and-excitation networks[C]//Proceedings of the 2018 IEEE Conference on Computer Vision and Pattern Recognition，2018：7132-7141.
[15] HOU Q，ZHOU D，FENG J.Coordinate attention for efficient mobile network design[C]//Proceedings of the 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition，2021：13713-13722.
[16] 彭冬亮，王天兴.基于GoogLeNet模型的剪枝算法[J].控制与决策，2019，34（6）：1259-1264.
PENG D L，WANG T X.Pruning algorithm based on GoogLeNet model[J].Control and Decision，2019，34（6）：1259-1264.
[17] 孙彦丽，叶炯耀.基于剪枝与量化的卷积神经网络压缩方法[J].计算机科学，2020，47（8）：261-266.
SUN Y L，YE J Y.Convolutional neural networks compression based on pruning and quantization[J].Computer Science，2020，47（8）：261-266.
[18] FAN Y，TANG X，MA Z.A weight-based channel pruning algorithm for depth-wise separable convolution unit[C]//2021 4th International Conference on Algorithms，Compating and Artificial Inteuigence，2021：22.
[19] HAN S，ZHAN Y，LIU X.Variational automatic channel pruning algorithm based on structure optimization for convolutional neural networks[J].Journal of Internet Technology，2021，22（2）：339-351.
[20] LI Z，XIN J.Channel pruning in quantization-aware training：an adaptive projection-gradient descent-shrinkage-splitting method[J].arXiv：2204.04375，2022.
[21] ZHAO X，YAO Y，WU H，et al.Structural watermarking to deep neural networks via network channel pruning[C]//Proceedings of the 2021 IEEE International Workshop on Information Forensics and Security，2021.
[22] LI Y，ADAMCZEWSKI K，LI W，et al.Revisiting random channel pruning for neural network compression[C]//Proceedings of the 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition，2022.
[23] WOO S，PARK J，LEE J Y，et al.CBAM：convolutional block attention module[C]//Proceedings of the 15th European Conference on Computer Vision，2018：3-19.
[24] WANG Q，WU B，ZHU P，et al.ECA-Net：efficient channel attention for deep convolutional neural networks[C]//Proceedings of the 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition，2020.
[25] KRIZHEVSKY A，SUTSKEVER I，HINTON G E.Image-
Net classification with deep convolutional neural networks[J].Communications of the ACM，2017，60（6）：84-90.