计算机工程与应用 ›› 2021, Vol. 57 ›› Issue (20): 263-270.DOI: 10.3778/j.issn.1002-8331.2105-0347

• 工程与应用 • 上一篇    下一篇

面向轻量化网络的安全帽检测算法

蒋润熙,阿里甫·库尔班,耿丽婷   

  1. 新疆大学 软件学院,乌鲁木齐 830046
  • 出版日期:2021-10-15 发布日期:2021-10-21

Safety Helmet Detection Algorithm for Lightweight Network

JIANG Runxi, Alifu·Kuerban, GENG Liting   

  1. School of Software, Xinjiang University, Urumqi 830046, China
  • Online:2021-10-15 Published:2021-10-21

摘要:

随着深度学习的发展,神经网络模型的体积越来越大,伴随而来的是参数量与计算量的增多,但实际安全帽检测环境下需要把网络模型部署在算力有限的移动端或嵌入式设备中,而这些设备无法支持复杂的计算量。针对这个问题,提出了一种适合部署在移动设备的轻量级目标检测网络HourGlass-YOLO(HG-YOLO)。以YOLOv5为基础模型,基于Inverted Resblock结构重构了新的主干特征提取网络HourGlass;并使用通道剪枝技术,对BatchNormalization(BN)层进行稀疏训练,将权值较小的通道进行删减,在保证精度的情况下,减少模型的参数;融合卷积层和BN层来加快在CPU上的推理速度。实验结果表明HG-YOLO在保证精度的情况下,将YOLOv5模型的体积压缩87%、浮点数减少86%、参数量降低89%,相比SSD在检测速度上快了8.2倍,更适合实际工业场景中的部署。

关键词: 深度学习, 目标检测, 轻量化模型, 剪枝

Abstract:

With the development of deep learning, the neural network model volume is becoming more larger, accompanied by the increase of parameters and computation. However, in helmet detection environment, the network model needs to be deployed in the mobile or embedded devices with limited computing power, and these devices can’t support complex computation. To solve this problem, this paper proposes a lightweight target detection network HourGlass YOLO(HG-YOLO) for mobile devices. Taking YOLOv5 as the basic model, a new backbone feature extraction network HourGlass is reconstructed based on Inverted Resblock structure. Channel pruning is used to sparsely train the BatchNormalization(BN) layer, and the channel with smaller weight is deleted to reduce the model parameters while ensuring the accuracy. Finally, the convolution and BN layer are integrated to speed up the reasoning speed on CPU. The experimental results show that HG-YOLO can compress the volume of YOLOv5 model by 87%, reduce the number of floating points by 86% and reduce the parameters number by 89%. Compared with SSD, HG-YOLO is 8.2 times faster in detection speed, which is more suitable for deployment in industrial scenes.

Key words: deep learning, target detection, lightweight model, pruning