基于注意力机制的YOLOv5优化模型

YOLOv5 Optimization Model Based on Attention Mechanism

目标检测是机器视觉研究中的重要分支。目前在工业生态中应用广泛的YOLOv5模型经过版本迭代,在预测权重大小以及检测精度方面都有所优化,但模型的处理速度仍然较低,尤其是对于小目标及遮挡目标的检测效果有待改进。该文提出一种基于注意力机制的YOLO v5改进模型。首先,通过引入维度关联注意力机制模块进行特征融合,提升主干网络的特征提取能力,达到改善小目标与遮挡目标的检测效果;其次,采用SIoU损失函数代替CIoU损失函数,作为新的边界框回归参数的损失函数,提高边界框的定位精度以及检测速度。实验结果显示,优化模型的平均精度均值达到87.8%,相比于YOLOv5提高了4.7百分点,在单GPU上模型的检测速度达到83.3 FPS。

With the development of machine vision technology,target detection has become an important branch.At present,the YOLOv5 model,which is widely used in the industrial ecology,has undergone version iterations and has been optimized in terms of prediction weight and detection accuracy,but the processing speed of the model is still not high,especially for small targets and occluded objects.The detection effect needs to be improved.We propose an improved model of YOLO v5 based on attention mechanism.First of all,by introducing the dimension related attention mechanism module for feature fusion,the feature extraction ability of the backbone network is improved to improve the detection effect of small targets and occluded objects;secondly,the SIoU loss function is used instead of the CIoU loss function as a new bounding box regression parameter.The loss function improves the positioning accuracy and detection speed of the bounding box.The experimental results show that the average precision of the optimized model reaches 87.8%,which is 4.7 percentage points higher than that of YOLO v5,and the detection speed of the model on a single GPU reaches 83.3 FPS.