基于YOLOv5的轻量化交通标志检测方法

何鑫，陈辉

(安徽理工大学 计算机科学与工程学院，安徽 淮南 232001)

随着人工智能技术的快速发展,智能汽车的自动驾驶已经逐步走入人们的生活,并已开始投入商业应用。交通标志识别(traffic sign recognition,TSR)是自动驾驶的基础性技术,也是自动驾驶系统的关键技术[1],对其研究也逐渐成为目标识别领域的热点。交通标志识别方法可以分为两大类：传统识别方法和基于深度学习的识别方法[2]。传统识别方式主要基于颜色、形状等方法识别,如Romdhane等[3]结合颜色分割、HOG(histogram of oriented gridients)和SVM分类器以识别交通标志;Shang[4]基于HOG特征训练级联分类器的方法识别交通标志。近年来随着卷积神经网络的高速发展,深度学习的检测方法逐渐成为热门的研究方向。深度学习方法有两类,一类是先进行候选框的筛选再提取特征进行分类、预测回归的两阶段目标检测算法(two-stage),代表性的算法有RCNN[5]、Fast-RCNN[6]、Mask-RCNN[7];另一类是一次检测即可产生预测的单阶段目标检测算法(one-stage),典型的算法有YOLO系列[8]、SSD系列[9]。虽然传统依赖颜色、形状的交通标志识别方法鲁棒性较强,但实时检测中受限于天气情况、光照条件、交通标志损坏等因素的干扰,基于检测RoI(region of interest)进行分类的传统检测方法往往不能在其他因素影响下达到理想的识别效果,且算法计算量大,而基于深度学习的识别方法特征表达能力强,也不易受到复杂环境的影响而导致识别精度大幅下降。因此,深度学习用于交通标志识别在某些方面更具优势。如江金洪等[10]将深度可分离卷积((depthwise separable convolution,DSC)[11]引入YOLOv3[12],并更换损失函数MSE为GIOU来检测交通标志;吕禾丰等[13]替换YOLOv5损失函数为EIOU[14],并使用加权Cluster非极大抑制(NMS)改进YOLOv5中原NMS以提升检测精度。虽然这些改进算法相较于传统检测方法在精度方面有所进步,但依旧存在算法计算量大、模型尺寸大等问题。

综上所述,基于颜色、形状和基于深度学习的交通标志识别技术分别存在着检测精度低和计算量大、模型尺寸大、推理速度慢的问题,使得模型难以部署到移动设备中。鉴于此,本文以YOLOv5为基础网络结构,通过GhostNet和坐标注意力机制的有效位置融合以及SIOU损失函数优化使用,提出一种轻量化改进的YOLOv5交通标志检测模型。

1 基于 YOLOv5的轻量化改进算法

1.1 主干网络的改进

YOLO算法是Redmon等于2016年首次提出的目标检测算法,在不同领域得到广泛应用。YOLOv5是YOLO系列的第五个版本,主要由输入端、Backbone、Neck、Prediction四部分组成。输入端主要负责图像的输入,Backbone负责对输入图片进行不同程度的特征提取,Neck部分融合、组合这些不同尺寸特征提取后的特征图,最后由Prediction预测。YOLOv5依照深度宽度的不同划分为五个不同的版本,其中YOLOv5m、YOLOv5l以及YOLOv5x网络结构复杂、体积过大,难以部署到嵌入式设备中,YOLOv5n模型尺寸过小,精度难以达到一定的高度,故本文选用深度较小,特征图宽度较小的YOLOv5s作为基础算法加以改进。改进模型的整体框架如图1所示。

图1 改进模型框架

改进模型利用GhostNet网络对YOLOv5s整体架构进行轻量化,将结构复杂、参数较高的C3模块与Ghost Bottleneck融合为C3Ghost模块,替换模型中所有原Bottleneck模块,剩余Conv模块更换为Ghost模块,实现网络结构的精简,最后在模型的小目标尺寸特征提取区部署CA坐标注意力机制,即YOLOv5s中Backbone最后一层,增强模型小目标的识别能力。

1.2 GhostNet网络

鉴于YOLOv5s存在网络复杂度高的问题,使用构建好的C3Ghost模块和普通Ghost模块融合搭建新YOLOv5s网络框架,不仅能有效减小模型整体参数以压缩模型体积,且能保证其具有一定的准确性。GhostNet是由Han等[15]在2020年提出的一种轻量化神经网络架构,该网络不仅能大幅度减少模型参数,且具有较高的识别性能,可以在同等计算成本下拥有赶超Mobilenetv3[16]的准确性。主流卷积神经网络在训练过程中会产生大量特征映射冗余,GhostNet网络提出一种新的Ghost模块。图2为GhostNet网络结构图,可以看出,GhostNet主要分两步,首先进行常规卷积运算得到一半特征图,再通过廉价的线性变换将其中一半卷积生成具有更多相似特征的另一半特征图,最后两部分相结合输出,这种方式提高了卷积量,同时减少了参数量和计算量。

图2 GhostNet网络结构

GhostNet 参考 了ResNet先升维再降维的网络结构,如图 3 所示。Ghost Bottlenecks 由两个Ghost 模块构成,第一个 Ghost 模块主要作用是扩展并增加通道数量;第二个 Ghost 模块主要负责减少匹配通道数量以及链接两模块输入输出。GhostNet 借鉴 MobileNetv2[17]中的设计技巧,在第一个 Ghost 模块后不使用 ReLu 激活函数,其余层都进行 BN (batch normalization) 批归一化和 ReLu非线性激活。这种特殊的结构设计,不但能够压缩参数量和计算量,还可以优化特征图,提高模型的性能。

图3 GhostNet结构图

图3所示C3Ghost模块参考CSP模块的构建思路,利用两个CBS(由普通卷积加批归一化加silu激活函数构成)和一个Ghost Bottlenecks组成。输入特征信息分别通过Ghost Bottlenecks以及CBS处理拼接成新的特征图,这样的构成方式减少了整体参数,同时还能保证特征信息的最大保留。

1.3 CA坐标注意力

对于输入图像来说,背景复杂,交通标志空间占比小,模型往往会受到大量冗余背景信息干扰导致损失小目标的特征信息,致使识别精度下降。此外,经过轻量化后的YOLOv5s中存在大量Ghost模块,Ghost模块使用分组卷积,减少了通道的数量,使得通道特征仅与自己建立联系。这一方面令参数量降低、网络复杂度减小,但是另一方面会影响对全局特征的提取,导致交通标志检测精度降低。因此本文将位置信息嵌入到通道注意力的坐标注意力机制[18],引入CA注意力机制到已优化的YOLOv5s网络结构Backbone部分的最后一层的C3模块,帮助模型把握全局特征,抵消部署Ghost模块带来的特征缺失。

有别于其他主流注意力模块,如SE通过全局池化来计算通道注意力[19],CA将全局池分解为两个一维的特征编码以缓解二维池化造成的位置信息缺失,并使坐标信息和输入特征图形成交互。图4为坐标注意力机制CA的网络结构图,由坐标信息的嵌入和坐标注意力生成两部分构成。

图4 CA网络结构图

首先,坐标信息的嵌入,对于给定维度为C×H×W的输入特征图,使用空间范围分别为(H,1)或(1,W)的池化层沿X轴和Y轴两个不同方向坐标的每个通道进行编码,经过编码平均池化后得到尺寸为C×H×1和C×1×W的特征图;其次,坐标注意力生成,通过拼接和1×1卷积将两种尺寸的特征图聚合成一个C/r×1×(W+H)的特征图,维度也从C维降成了C/r维;再分别用1×1卷积进行维度提升,并使用Sigmoid激活函数激活;最后进行特征图重加权,即将生成的两个同纬度特征图与输入特征图相乘,完成对模型坐标注意力的添加。

图5(a)为主干网络没有添加CA的识别效果,5(b)为添加后的效果,可以看出未添加CA会出现对小目标漏检的情况,添加后漏检情况得到改善。

(a)未添加CA (b)添加CA图5 CA识别效果对比图

1.4 边框回归损失函数

损失函数能够度量模型的预测值与真实值差异程度,其设定也会直接影响模型优化的速度和定位的准确性,因此选择合适的损失函数对模型起着至关重要的作用。

YOLOv5原损失函数为CIOU损失函数[20],CIOU不仅吸收了GIOU[21]、DIOU[20]的优点,即在IOU的基础上加上对预测框和真实框的重叠面积、中心点距离的考虑,而且提出对回归框宽高比进行损失计算,使得目标框回归更加趋于稳定。

虽然CIOU考虑到了预测框和真实框的重叠面积、中心点距离以及宽高比,但忽略了其方向之间不匹配的问题,导致模型收敛速度不理想。本文选用SIOU损失函数[22]替换YOLOv5中的CIOU损失函数。SIOU吸收了CIOU的优点,并添加了对真实框和预测框之间角度的损失计算,提高了模型的训练效率和收敛速度。SIOU损失函数由角度损失(angle cost)、距离损失(distance cost)、形状损失(shape cost)、IOU损失(IOU cost)四部分组成。

角度损失的添加是为了最大程度地减少与距离相关的变量数量,在模型训练收敛过程中首先将预测带到最靠近的坐标轴(X轴或Y轴),再沿此轴继续进行训练,其计算公式为

(1)

式中：ch为真实框和预测框的高度差,σ为真实框和预测框中心点的距离。

考虑到角度损失的添加,需要实现角度损失和距离损失的平衡。距离损失公式为：

(2)

(3)

形状损失指的是两框中心位置的偏差,计算公式为：

(4)

(5)

式中：w和wgt分别为预测框和真实框的宽度,h和hgt分别为预测框和真实框的高度,θ体现了对形状损失的关注程度。

IOU损失表示真实框和预测框之间的交集和并集之比,计算公式为

(6)

最终损失函数表达式为

(7)

2 TensorRT

TensorRT是由NVIDIA公司提出的一个基于深度学习的高性能学习推理优化器,它支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架,将TensorRT和NVIDIA的GPU结合起来,能在几乎所有的深度学习框架中进行快速和高效的部署推理。本文利用TensorRT对已训练好的模型权重进行优化,以加速模型推理速度。

3 实验结果与分析

3.1 数据集

为了验证轻量化改进模型的效果,本文采用由长沙理工大学采集并制作的交通标志数据集CCTSDB2021(csust Chinese traffic sign detection benchmark 2021)[23]作为实验所用数据集,验证本文算法。CCTSDB2021由三大类交通标志构成,分别为指示标志、禁止标志和警告标志,数据集制作规范符合实验要求。随机选取其中6 000张图片,按照VOC2007数据集格式调整,按照4∶1的比例划分为训练集和验证集进行实验验证。

3.2 数据增强处理

为了提升模型的泛化性能,使用Mosaic和Mixup数据增强方法处理数据集。Mosaic数据增强方法通过随机排列拼接经过裁剪、缩放的四张输入图片形成新的图片输入训练模型,扩充小样本目标的同时提升模型训练速度,增强网络的鲁棒性;而Mixup数据增强方法则是通过随机叠加融合两张不同的图片丰富样本数量,实现对小目标检测能力和模型鲁棒性的增强。图6为使用Mosaic和Mixup数据增强后的效果图。

(a)Mosaic (b)Mixup图6 数据增强后的效果图

Mixup计算公式为：

(8)

(9)

式中：(xi,yi)和(xj,yj)是原始数据的随机样本对;λ是服从beta分布的一个参数,λ∈[0,1]。

为验证数据增强方法对YOLOv5s模型检测产生的影响,将未开启数据增强与开启数据增强的YOLOv5s模型进行实验对比,表1为对比实验结果。

表1 数据增强前后对比结果

从表1可以看出,采用数据增强后的YOLOv5s模型MAP值有了提升,说明在YOLOv5s模型中,上述两种数据处理方式均可以有效提升模型的平均精度。

3.3 实验配置

本实验使用Pytorch深度学习框架搭建、训练改进模型,实验环境为Ubuntu20.04,CPU采用Intel(R)Xeon(R)Platinum8255C,GUP为V-100SXM2,内存为32 G。实验输入的图像尺寸为640×640,批训练处理量为32,初始学习率为0.01,IOU阈值为0.2,权重衰减为0.000 5,共训练150轮。

3.4 评价指标

实验使用平均精度(mean average precision,MAP)、推理速度、参数数量、模型大小作为对改进模型的评价指标。其中MAP通过准确率(Precision)和平均正确率(average precision,AP)计算得出。

准确率即网络检测到正样本数量占总样本数量的比率,公式为

(10)

式中：TP(true positive)为被正确识别为正样本的正样本数量,FP(false positive)为被错误识别为正样本的负样本数量。AP、MAP分别为预测样本的平均正确率、平均精度,所有类别的平均精度为：

(11)

(12)

3.5 实验结果与分析

本文选用YOLOv5s、YOLOv3-tiny、YOLOv4-tiny、YOLOv5s+Shufflenetv2、SSD五种网络进行交通标志检测实验,选用模型涵盖了近年来目标检测领域内经典框架和轻量型改进算法,并将本文模型与之进行性能对比,以验证本模型的可行性,对比结果见表2。

表2 不同算法对比结果

从表2可以看出,本文模型与轻量化网络模型YOLOv3-tiny相比平均精度提升了6.7%,参数量和模型大小分别减少了63.8%和61.9%;与另一主流轻量化网络模型YOLOv4-tiny相比平均精度提升了3.7%,参数量降低了46.9%,模型大小降低了71.8%;和引入Shufflenetv2的YOLOv5s的轻量化模型YOLOv5s+Shufflenetv2相比精度有较大提升,参数量和模型大小均有所减小;相较于主流检测模型SSD,精度大幅提升,参数量和模型大小大幅减小;与原模型YOLOv5s相比,平均精度只有0.1%的损失,但参数量和模型大小大幅度减小。

改进YOLOv5和原模型YOLOv5s对交通标志数据集训练结果的平均精度如图7所示,可以看出改进模型与原模型相比平均精度几乎没有差距,在训练前期效果甚至略强于原模型。

图7 平均精度对比图

另外改进模型推理速度也有提升。实验结果表明,原模型的推理时间为9.4 ms,改进模型的为9.2 ms,提升幅度为2%;在经过TensorRT加速后,模型推理时间更是可以达到4 ms,比原模型提升了57.4%。

改进模型对比主流检测模型、轻量化网络模型,不仅在精度上更具优势,且参数量和模型大小也远远小于其他模型;在和原模型的对比上,虽然精度有些许损失,但模型尺寸和参数量减小,且推理速度远超原模型,综合来看改进模型的综合性能优于原模型。

为了验证所改部分的有效性,对改进模型进行消融实验。实验对比结果见表3。表3中改进模型1为融入GhostNet的YOLOv5s模型,改进模型2添加了CA注意力。从表3可以看出,利用GhostNet对YOLOv5网络架构进行优化使得参数量和模型尺寸减小,但模型精度过低;而在主干网络加入CA坐标注意力模块后,模型对重点信息定位能力有所加强,精度增加较大,且参数量和模型尺寸进一步减小;YOLOv5s将原损失函数CIOU优化为SIOU后,弥补了CIOU在真实框和预测框角度问题的不足,精度略微增加;可以看出本文模型在保证了较少精度损失的同时,有效地压缩了参数量和模型大小。

表3 消融实验对比结果

4 结论

1)融合CA注意力与GhostNet网络结构的YOLOv5模型,相较于原模型,参数量减少50%,模型大小减小53.5%,对嵌入式设备的有限计算资源占用较少,利于部署。

2)轻量化模型精度损失仅0.1%,并未因参数量、计算量的减少造成过大的精度损失,相对而言改进模型综合性能更强。

3)改进模型推理速度快于原模型,相较于原模型推理速度提升了57.4%。