计算机工程与应用 ›› 2021, Vol. 57 ›› Issue (21): 165-173.DOI: 10.3778/j.issn.1002-8331.2007-0230

• 模式识别与人工智能 • 上一篇    下一篇

改进YOLOv3-SPP的无人机目标检测模型压缩方案

黄文斌,陈仁文,袁婷婷   

  1. 南京航空航天大学 机械结构力学及控制国家重点实验室,南京 210016
  • 出版日期:2021-11-01 发布日期:2021-11-04

Compression of UAV Object Detection Model Based on Improved YOLOv3-SPP

HUANG Wenbin, CHEN Renwen, YUAN Tingting   

  1. State Key Laboratory of Mechanics and Control of Mechanical Structures, Nanjing University of Aeronautics and Astronautics, Nanjing 210016, China
  • Online:2021-11-01 Published:2021-11-04

摘要:

无人机设备算力低下,深度模型计算量过大不适合直接部署,航拍图像目标小并且密集,模型对目标识别分类效果也不佳。为了提高深度模型航拍目标检测的精度和速度,降低计算量。对YOLOv3-SPP模型进行改进,将GIoU代替平方和用作定位损失,提高定位精度。提出了一种数据集优化和数据增强方法。再针对特定类别按照权值进行采样处理均衡化类别数量。随机组合不同场景样本组成批训练,提高模型训练效率和检测鲁棒性。再对模型进行压缩,在BN层添加缩放因子进行稀疏训练和通道剪枝的基础上,通过缩放因子衡量模型残差层重要性,修剪不重要残差,进一步减小前向推理层数和参数。实验表明,模型参数量减小了95.7%,模型大小减小95.82%,同等算力下模型推理速度提高为原来3倍。且精度和速度均高于最新YOLOv5系列轻量模型。

关键词: 目标检测, 无人机, 通道剪枝, 模型压缩, YOLOv3

Abstract:

Due to the limited memory and computing power of UAV equipment, it is very challenging inrunning deep-learning models on such devices for object detection. Dense and small sizes objects in aerial image and datasets increases the difficulty on aerial target recognition and classification. To deal with these challenges, the YOLOv3-SPP model is improved by using GIoU instead of the mean-square error, which raise the positioning accuracy if objects. A method of data enhancement is also proposed. It balances categories number by weighted sampling specific categories. During training, different scene pictures are randomly combined to form batch training, which increases detection robustness. Then, the model is compressed. On the basis of sparse training and channel pruning by adding scaling factor to BN layer, the importance of residual layer is measured by scaling factor, and unimportant residual is pruned to further reduce the number of forward reasoning layers and parameters. In experiment, model parameters are reduced by 95.7%, the module size is reduced by 95.82%, and the speed is increased by 3 times. The accuracy and speed are higher than the latest YOLOv5 series lightweight model.

Key words: object detection, UAVs, channel pruning, model compression, YOLOv3