基于改进YOLO v4的轻量化烟梗识别方法

郑银环, 林晓琛, 吴 飞, 金圣洁, 吴傲男

(武汉理工大学 机电工程学院,湖北 武汉 430070)

0 引 言

在烟叶制丝工艺中,首先需要将大片烟叶进行打叶复烤,通过打叶设备使叶片和烟梗分离,将整片烟叶破碎成小块。经打叶复烤后虽绝大多数烟梗已分离,但其中仍夹杂着一定量的烟梗。烟叶打叶复烤后应符合严格的质量标准[1],因此需将复烤后的破碎烟叶进行松散回溯处理,即通过120 ℃热蒸汽提升烟叶的温度与水分,达到烟叶松散和舒展的效果,随后通过振动筛装置使烟叶铺开减少烟叶堆叠情况,再进一步去梗。由于生产质量标准严格,该工序一直都是人工完成,带来生产效率低、耗时长、人工费用高等生产实际问题。

针对烟叶中烟梗识别与剔除的问题,国内外学者进行了大量的研究。文献[2]研究了低能X射线透射成像,结合形态学滤噪、灰度阈值分割、归属判断等图像实时分析处理手段,实现了打叶片烟中烟梗的在线识别及剔除;文献[3]设计了一种以基于现场与编程门阵列(field programmable gate array,FPGA)的高速图像处理设备为核心的烟梗在线检测系统,该系统采用 X 光源透射皮带上的烟叶以及烟梗,通过探测器接受X射线并通过高速图像处理设备对图像进行算法处理,进而实现烟梗识别;文献[4]利用 X 射线的透射性原理,结合 X 射线穿过物质后的衰减规律,对片烟中烟梗质量进行拟合,从而计算出叶中含梗率;文献[5]采用高功率红外线作为透射光源,以高清连续摄录机为图像获取平台,以灰度算法为核心程序,实现烟梗面积自动化测算系统的开发。

上述研究利用烟叶和烟梗在不同光源下反映出的特性不同作为判据依据进行烟梗识别及分离,而在实际生产中,经松散回溯后的破碎烟叶通常为散乱堆叠状态且烟梗特征不一,造成上述方法在实际生产中的识别精度不高。

深度学习技术在计算机视觉领域中的应用发展迅速。根据算法原理不同主要分为2类:① 先进行区域生成,再通过卷积神经网络(convolutional neural network,CNN)进行样本的分类,其中最具代表性的算法有R-CNN、SPP-Net、Fast R-CNN、Faster R-CNN、R-FCN等[6-10]; ② 不进行区域生成,直接在网络中提取特征来预测物体分类和位置,常见算法有OverFeat[11]、SSD[12]、YOLO系列算法[13-14]等,此类算法因具备可同时预测物体种类和位置的特性,相比于第①类算法在目标检测过程中更具快速性。而YOLO v4作为检测速度和检测精度兼备的检测算法,在工业现场应用广泛。

基于深度学习的目标检测技术虽然识别精度高、可移植性好,但往往计算量大,模型结构复杂,占用空间大,而工业现场的工控机等主机设备往往储存空间有限,使得检测算法的现场部署实施困难。

本文以YOLO v4为基础网络,建立了烟叶回溯振动铺开后烟梗的轻量化检测模型。利用通道剪枝和层剪枝相结合的方法对模型进行压缩,通道剪枝大大减小了模型参数量和计算量,降低了模型对资源的占用。层剪枝可以进一步减小计算量,并大大提高模型推理速度。两者相结合,可以大幅度压缩模型的深度和宽度。

1 数据集

1.1 数据采集

本文在数据集制作过程模拟真实工业场景,综合考虑了真实生产线上烟叶可能出现的情况。实验对象为松散回溯工序后的破碎烟叶,且已由振动筛装置铺开无明显堆叠情况,人工将破碎叶片随机摆在黑色背景的实验拍摄平台,使用英特尔RealSense d435i型号相机采集烟叶数据集。实验拍摄平台如图1所示。

图1 图像采集平台

实验采用设备灯光进行拍照,光照条件设置为曝光度640、对比度50、饱和度64。图像输出尺寸为1 920像素×1 080像素,帧率为8,共拍摄1 000张烟叶照片。

1.2 数据预处理

为丰富数据集,提高模型泛化能力,进一步提升模型性能,本文采用Mosaic数据增强方法,随机读取4张烟叶照片,分别对4张照片进行翻转、缩放等操作,随之将4张照片按照左上、左下、右下、右上的顺序排布,最后利用矩阵的方式将4张图片在固定位置截取下来并拼接成1张新的图片。原始烟叶照片及Mosaic数据增强后的照片如图2所示。

图2 破碎烟叶图像

最终采集到烟叶数据集共2 000张,其中训练集与测试集比例划分为7∶3。

2 烟梗快速识别方法

松散回溯后振动铺开的烟叶形状不规则,烟叶中烟梗特征较不明显。将标注好的烟叶烟梗数据集输入到YOLO v4网络中进行基础训练,得到基础模型。针对YOLO v4基础训练模型进行稀疏训练,随后将稀疏训练模型进行通道剪枝和层剪枝,修剪掉模型中的冗余通道、Shortcut层和卷积层,大幅度减小了模型的参数量和模型大小,提高了模型推理速度,但模型精度会有所下降。最后对剪枝后的模型进行微调,使模型精度重新回升至较高水平。

2.1 基于YOLO v4的烟梗检测模型

YOLO v4以CSPDarknet53作为主干网络实现图像特征的提取。CSPDarknet53将YOLO v3的主干网络Darknet53与CSPNet进行结合形成新网络主干,其中CSPNet将基本层的特征图分成两部分,然后通过一个跨阶段的层次结构合并它们,使新的网络主干能够实现更丰富的梯度组合,同时减少计算量;以PANet(path aggregation network)代替了YOLO v3中的特征金字塔网络(feature pyramid network,FPN)作为参数聚合,针对不同主干层进行参数聚合;同时继承了YOLO v3的多尺度预测模块,输出层的锚框机制与YOLO v3相同,主要改进了训练时的损失函数CIoU-Loss以及预测框筛选的DIoU-nms,进一步提升了算法的检测精度。

检测烟叶中烟梗的YOLO v4网络模型结构如图3所示。

图3 YOLO v4网络模型结构图

该检测模型主要由CBM(convolution,batch normalization and mish)、CBL(convolution,batch normalization and Leaky ReLU)、CSPX(center and scale prediction)、SPP(spatial pyramid pooling)和ResUnit等基础组件组成。

CBM是YOLO v4网络结构中的最小组件,由Conv+BN+Mish激活函数组成,Conv是卷积层,用来提取特征,BN用于归一化处理,Mish为激活函数,允许更好的信息深入神经网络,从而得到更好的准确性和泛化;CBL由Conv+BN+Leaky_ReLU激活函数组成,与CBM不同的是采用了Leaky-ReLU激活函数,在反向传播过程中,其输入小于0的部分,也可以计算得到梯度,避免了梯度方向出现锯齿的问题;ResUnit借鉴网络中的残差结构,可以构建更深的网络,解决了深度神经网络中给网络叠加更多层后,网络深度在增加性能却快速下降的问题;CSPX借鉴CSPNet网络结构,由3个卷积层和X个ResUnit模块Concat而成,在减少计算量的同时实现更丰富的梯度组合;SPP采用1×1、5×5、9×9、13×13的最大池化方式进行多尺度融合,大大增加了感受野。

2.2 损失函数

损失函数在模型中起到关键性作用,可以衡量模型预测的好坏,用来表现预测与实际数据之间的差距程度。基于YOLO v4的烟叶中烟梗检测模型的损失函数由回归框损失函数(lossCIoU)、置信度损失函数(lossconf)和分类损失函数(lossclass)组成。损失函数计算公式为:

loss=lossCIoU+lossconf+lossclass

(1)

回归框损失值为:

lossCIoU=1-IoU+d2/f2+αν

(2)

(3)

(4)

置信度损失值为:

(5)

分类损失值为:

(6)

2.3 基于改进YOLO v4的烟梗检测模型

基于YOLO v4的烟梗检测模型可实现松散回溯后破碎烟叶中烟梗的识别,但基础训练模型储存空间较大,模型参数量较多,且模型推理速度较慢。基于上述问题,本文提出了基于通道剪枝和层剪枝相结合的模型轻量化剪枝方案,将YOLO v4训练后得到的模型进行裁剪,该操作能够大大减小模型存储量和计算量,显著提高模型推理速度,且能保证模型精度。

通道剪枝本质上是指在模型经过稀疏训练后,消除BN层γ值较小的通道,即输入通道中贡献率较低的通道。通道剪枝在保证网络结构完整的情况下降低了模型参数量,减小了模型存储空间,且不需要特定的网络框架或硬件支持,使其更方便部署在移动设备上。

通道剪枝算法原理如图4所示。烟梗识别模型采用通道修剪γ值来评估其通道的重要性。假设γ值接近于0,输出通道也接近于0,因此γ值所在的通道对整个网络的影响很小,可以进行修剪。

图4 通道剪枝算法原理图

在对烟梗特征提取网络进行通道修剪后,进一步去除多余的残差块,如图5所示。在修剪冗余残差块时,首先对所有Shortcut层的前一个CBL的γ值的均值进行排序。若它接近于0,则该残留块没有学习到有用的信息,那么对其进行剪枝不会影响网络性能,因此γ值均值较小的残差块将被修剪。为保证网络的完整性,每剪一个Shortcut结构,会同时剪掉一个Shortcut层及其前面的2个卷积层。本文一共剪掉16个Shortcut,共计剪掉了48个层,精度下降较少。

图5 修剪冗余残差块流程

稀疏训练后,许多γ值收敛到0,因此对模型的冗余通道和网络层进行剪枝,减少了模型的参数、大小和推理时间。

2.4 模型剪枝具体步骤

本研究所用计算机型号为GPU 1080Ti、CPU E5-2678v3、内存大小62 GiB、显存大小11 GiB,开发环境为Python 3.6+PyTorch 1.9.0+openc v4.5.5+CUDA 11.4。模型剪枝具体步骤模型如下:

1) 基础训练。剪枝前,先对数据集进行基础训练,利用YOLO官网给出的初始权重训练数据集,初始学习率(learning rate)设置为0.002 3,批样本量(batch size)为8,迭代次数(epoch)为200,基础训练情况如图6所示。

从图6a可以看出,随着迭代次数的推进,基础训练趋于150次时损失值趋于平缓,最终模型训练200次loss值稳定在0.5左右,模型收敛。从图6b可以看出,基础训练模型精度在前40次迭代稳步上升,到达40个epoch时达到0.850,在接下来的160次迭代中平均精度均值(mean average precision,mAP)平缓上升,模型最终收敛精度达到0.941。

2) 稀疏训练。烟梗识别模型经过基础训练后包含大量冗余通道,且通道数量在不同BN层间存在显著差异。本文采用稀疏训练方法来解决通道数量差异较大的剪枝问题,为后续模型剪枝作准备。通过将BN层γ值的L1归一化项添加到损失函数中,使γ值变得稀疏,L1归一化消除最终没有为输出提供有用信息的通道。稀疏因子s设置为0.001,批样本量(batch size)为8,迭代次数(epoch)为500,初始学习率(learning rate)为0.002 3,在稀疏训练迭代次数达到70%和90%的2个阶段进行γ值为0.1的学习率衰减,使模型精度有所回升。

从图6a可以看出,基础训练得到的模型进行稀疏训练后,整体loss曲线呈先上升后下降趋势,整体存在较大波动,在迭代次数为220时达到峰值,随后下降且最终趋于平缓,loss值最终停靠在1.0左右且存在细微波动。从图6b可以看出,稀疏训练模型精度继承基础训练的最终精度0.940且在稀疏训练的前120次的迭代中趋于稳定,在迭代次数为120后模型精度发生了较大跨越,最终在迭代次数为360后趋于平缓,mAP值在0.700左右且存在细微波动。

模型训练过程中权值分布直方图如图7所示,直方图表示模型训练过程中神经网络中每一层的权重。图7中:纵坐标值的增加表示模型BN层的递增,用来反映模型训练过程的推进;横坐标γ反映模型的稳定性和稀疏性,随着训练迭代次数即纵坐标值的增加,γ值分布范围逐渐相似,说明模型趋于收敛,γ值趋于0说明模型结构逐渐稀疏。

图7 YOLO v4模型基础训练、稀疏训练γ值直方图

图7a为YOLO v4基础训练模型各BN层γ值分布中心的变化趋势图。从图7a可以看出,γ值随着训练迭代次数的增加始终分布在1左右,γ值始终趋于稳定,说明基础训练一直处于稳定状态。

图7b为YOLO v4稀疏训练模型各BN层γ值分布中心的变化趋势图。从图7b可以看出,γ值在开始迭代时分布在1附近,随着训练迭代次数的增加逐渐趋于0,说明γ值逐渐变得稀疏,当迭代次数达到500左右时γ值趋于稳定,说明此时稀疏训练已处于稳定状态。

3) 模型剪枝及微调。对稀疏训练后的模型设置80%的通道剪枝全局比例和0.01的每层最低保持通道数比例,剪掉12个Shortcut及其前面24个CBL,总共剪掉36层。

剪枝结果见表1所列。

表1 剪枝前后模型参数

从表1可以看出,模型经过通道与层剪枝后参数量下降了95.49%,存储空间减少了234.44 MB,模型检测时间减少了约0.005 s,模型mAP降低了21.8%。鉴于模型剪枝后精度出现下跌,对剪枝模型进行100次迭代基础训练的微调,批样本量(batch size)设置为8,以补偿模型性能的下降。微调后模型存储空间略有增加,但控制在较小范围内。

模型精度回升到与剪枝前相近,且检测速度仍能维持在模型剪枝前水平。结果表明,对模型进行剪枝及微调可在保证模型精度及检测速度的前提下大幅简化模型。

3 实验结果与分析

3.1 烟梗检测效果分析

基于本文剪枝微调后的YOLO v4优化模型烟梗检测效果如图8所示。

图8 模型检测效果图

从图8可以看出,在实验背景下能够准确地识别出不同破碎烟叶片中存在的烟梗,识别精度高,检测速度快,适用于卷烟自动化工厂中识别剔除带梗烟叶工艺步骤,能代替现阶段烟厂工人经验化识别剔除步骤,减少劳动力,提高识梗去梗准确率和成品烟丝品质。

3.2 模型性能评价指标

选用mAP、精确率P、召回率R、F1值、模型大小以及模型检测时间等指标对不同模型进行性能评估。

3.3 不同模型烟梗检测效果比较

基于YOLO v4剪枝模型对烟叶图像进行测试,精确率为90.98%,召回率为95.46%,F1值为93.17%,mAP为93.00%,检测单张图片时间达到0.021 00 s。模型识别精度高、识别速度快,且占用空间相比同类模型得到极大压缩,为烟厂片烟精选工作提供有利的技术支持。

与同类别的目标检测模型相比,本文使用的方法更具备优势。常用的目标检测模型有SSD[12]、YOLO v3[13]、YOLO v4[14]、YOLO v4-tiny[15]、YOLO v5[16]等,在同样的硬件及开发环境下,使用这5种模型分别进行测试,但得结果见表2所列。

表2 不同模型性能评估

比较几种重要的模型评价指标可以看出,本文YOLO v4剪枝模型的召回率和F1值均最高,其精确率较Sota方法的YOLO v5模型虽下降0.21%,但模型mAP相较剪枝前下降了1.10%,变化较小,但仍保持在较高的精度水平,模型检测速度略有提升,同时模型存储大小得到了极大压缩,甚至比YOLO v4-tiny轻量化模型还要小7.25 MB,应用于工业现场可极大减小部署设备储存空间。剪枝模型检测速度虽相较于SSD、YOLO v3、YOLO v4-tiny略有下降,但最大检测时间差控制在0.008 00 s左右,对工业应用现场影响不大,且与此3种模型相比,本文方法的模型精度分别提升了9.31%、4.06%、13.64%。

综合比较,本文模型剪枝方法与同类目标检测模型相比,在极大压缩模型存储空间的基础上,仍能保持较高的模型精度以及模型检测速度,表明本文方法在片烟精选工业现场中的烟梗识别剔除应用中具有很强的应用前景。

3.4 不同剪枝率和层数模型效果对比

本文方法的模型剪枝过程中,涉及到人为设置通道剪枝全局比例以及层剪枝Shortcut层数这2个重要参数,不同参数下模型剪枝效果存在差异。本文比较了不同剪枝参数下模型的性能,结果见表3所列。以通道剪枝全局比例85%和80%以及层剪枝层数12和16分别对YOLO v4基础模型进行剪枝操作,分析模型剪枝结果。

表3 不同剪枝率和剪枝层数模型性能对比

在模型精度方面,剪枝后精度均出现一定幅度的下降,但整体降幅不大,全局比例为80%,层剪枝层数为12时,模型精度下降幅度最小,相比于原模型,mAP仅下降1.10%,说明较小的通道剪枝比例和层剪枝层数对模型精度影响更小。

在模型存储空间方面,剪枝后相比于原模型均得到大幅度压缩,且通道剪枝比例和层剪枝层数越大,对模型的压缩力度也越大,模型最小能压缩到9.88 MB,仅占原模型空间的4.0%。在模型检测速度方面,不同参数组合下的剪枝,模型检测均比原始模型有所提升,相比原始模型检测时间最快缩短了0.080 00 s。最终本文选用80%的通道剪枝全局比例及12个Shortcut的层剪枝层数作为YOLO v4基础模型的剪枝参数,在实现93.00%模型高精度的同时仍能维持较高的检测速度,且极大降低了模型存储空间。

4 结 论

本文基于通道剪枝和层剪枝相结合的模型压缩方案,提出了一种基于改进YOLO v4的轻量化烟梗识别方法。在保证模型精度和检测速度的前提下,对模型进行大幅度压缩,简化了模型结构,减小了占用空间。

本文使用英特尔RealSense d435i型号相机采集原始烟叶数据集100张,利用Mosaic数据增强的方法拓展数据集到200张,运用Labelimg对数据集进行标注,随后运用YOLO v4算法对数据集进行基础训练,得到基础模型后进行稀疏训练、模型剪枝、微调,最终得到优化后的烟梗识别模型。经实验验证,优化模型的存储空间下降了93.77%,模型精度达到93.00%,检测速度达到47.62 帧/s。实验结果表明,剪枝算法在保证检测速度快、精度高的同时极大地压缩了模型的大小,为算法模型在工业现场的部署提供了优势,有效解决了传统工控机存储空间有限的问题,适用于片烟精选中松散回溯后烟梗识别的工艺流程。