基于改进YOLOv4-tiny的零件目标检测*

殷宇翔 徐顺清 何坚强 蒋成晨 陆 群 唐乾榕

（盐城工学院电气工程学院 盐城 224000）

1 引言

机械制造业是我国经济发展的支柱性产业，其整体生产规模也在不断提升。近年来，柔性生产技术迅速发展，它是以数控机床为主的设备来生产多种类小批量的机械零部件［1］，比传统的批量生产更具灵活性，是促进先进制造业高新技术发展的重要途径。在柔性生产过程中会产生多种零部件的混合，往往需要人工进行归类分拣，生产效率不高。随着智能机器人的发展，自动抓取已经广泛运用于零件分拣领域，代替了传统的人工分拣，极大提高了工作效率［2］。但在抓取过程中，首先要对零件进行分类、定位［3］。国内对零件定位的研究大多停留于传统的模板匹配算法［4］，在识别的速度与精度方面还比较差，可移植性不强。且在实际的环境中，零件大多是散乱放置，无法进行较为准确的实时抓取。

随着神经网络的快速发展，更为高效精准的深度学习算法相继提出，已逐步代替传统算法。在发展过程中分成两种类别［5］。一类是以Faster RCNN［6］为代表的双阶段检测算法，该算法首先生成候选区域，再经过网络提取特征，最后对候选区域进行回归和分类，虽然有较高的准确率但速度较慢。在零件检测研究领域中，唐茂俊［7］等提出一种改进的Faster-RCNN算法，改善零件焊缝缺陷目标框定位问题，精确度提高了2.4%，但每副图像检测时间为0.245s。仝京等［8］通过调整现有的Mask-RCNN网络结构，提出使用mask分支的输出信息进行提取与加工，并结合多种边缘检测算法来提高目标识别的准确率。另一类是以YOLO［9］为代表的单阶段检测算法，该方法直接由网络预测出类别与候选区域，此类方法优点是识别精度较高且快速。杨琳［10］等通过利用Canny边缘检测和Sklansky算法与YOLOv4算法相结合，使零件形状的最小外接矩形框更为精准，准确率达到93.37%。张浩［11］等通过改进YOLOv3算法，采用Inception网络结构，使用不同尺寸的滤波器对目标进行多尺度特征提取，检测的速度提高了9%。可见单阶段算法在复杂环境仍有较高的识别能力。虽然文献［7～11］已提出一系列高精度零件检测算法，但是这些算法的参数量大，移植性不强，难以运用于工业机器人。而YOLO系列的简化版本检测速度快并对硬件配置要求较低，在嵌入式设备上使用更为普遍，因此为了满足工业机器人对多角度无序零件识别以及精准定位的需求，我们对YOLOv4-tiny算法进行改进。通过引入注意力机制和K-means算法来提高算法的检测精度。在自制零件数据集上，我们对比改进算法与YOLOv4-tiny算法，实验结果证明了改进算法的优越性。

2 YOLOv4-tiny算法原理

Alexey等提出了YOLOv4［12］的精简版模型YOLOv4-tiny［13］。YOLOv4-tiny的网络结构更为简单，对硬件的要求相对低，并且大大减少了参数的计算量，检测的速度更快，但其检测的精度会小幅度下降。YOLOv4-tiny网络主要包括以下几个部分：主干特征提取网络、特征金字塔（FPN）和网络预测头（YOLO Head）。与YOLOv4不同，YOLOv4-tiny的主干特征提取网络为CSPdarknet53-tiny，主要由卷积块和Resblock-body构成，卷积块采用CBL模块，由卷积层（Convolutional）、激活函数（Leaky Relu）和批量标准化（BN）构成。Resblock-body用来加强特征金字塔的构建，解决增加网络深度而导致模型退化的问题。

特征金字塔FPN的思想是通过多尺度特征融合。在预测过程中，低层的特征语义信息相对较少，但是位置准确，高层的特征语义信息比较丰富，但是目标位置比较粗略。FPN通过对高层特征进行上采样和对底层特征进行自顶向下的连接，能有效表达出特征图上各种维度特征。在基本不增加原网络模型计算量的基础下，FPN有效地提高了物体检测的综合能力。

在网络预测头部分，YOLOv4会生成三个head，而YOLOv4-tiny会生成两个尺寸分别为13×13和26×26的特征层传递到网络预测头，来预测大物体与小物体的检测结果。YOLOv4-tiny网络模型如图1所示。

图1 YOLOv4-tiny网络模型

3 改进YOLOv4-tiny

3.1 注意力机制的引入

YOLOv4-tiny在进行特征传递的时候，不会自动筛选重要信息，可能会忽略小目标以及部分遮挡信息。为了解决这个问题，我们在YOLOv4-tiny的网络结构中嵌入注意力机制，注意力机制是将感兴趣区域的特征数据进行权重系数加权，把更多的注意力集中到这个区域。

SGE（Spatial Group-wise Enhance）为空间注意力机制［14］，其网络结构如图2所示。它主要关注目标的一些位置信息。该机制首先对大小为W×H×C的特征图I的通道维度进行最大池化（）和平均池化（），之后将这两个结果进行一个合并，通过卷积操作调整通道数，使用Sigmoid激活函数后，将得到输入特征层每一个特征点的权值，再乘上原模块的输入层得到最终特征图Fs(I)，以此来关注零件的位置特征。公式表示为

图2 空间注意力模型

SE（Squeeze-and-Excitation）为通道注意力机制［15］，网络结构如图3所示。它主要关注目标的一些特征信息。通道注意力机制会对输入进来的特征图I分别进行全局平均池化（）和全局最大池化（），然后送入共享的全连接层进行处理，将处理后的两个结果相加，通过Sigmoid函数获得输入特征层每一个通道的权值，每个通道的权值与原特征层做内积运算，得到最终输出特征图I′，这样使得输出层通道中含有零件的更多关键信息。该模块计算公式如下：

图3 通道注意力模型

CBAM［16］注意力机制它是一个轻量级即插即用的模块，其网络模型如图4所示。它是由空间注意力机制和通道注意力机制结合组成。该模块计算公式如下：

图4 CBAM注意力机制模型

3.2 改进的YOLOv4-tiny算法模型

当把CBAM注意力机制融入到YOLOv4-tiny的CSP结构时，为了不破坏主干提取网络的初步特征，把CBAM注意力机制添加到加强特征提取网络，在FPN层中进行上采样后，将输出的特征层经过通道注意力机制得到权值，将所得的权值乘以原输出的特征层，让所得的结果再经过CBAM中的空间注意力机制获得权值，将权值与原输出的特征层相乘即为最后的输出，并作为下一层的输入。在这个过程中，CBAM模块前后的特征图尺寸以及通道都不会发生变化。在主干提取网络与FPN连接的分支进行相同的操作。改进后的网络将会自动对特征图上重要的目标信息重新进行权重分配，使网络聚焦于特征图中零件的关键信息，还能够抑制无关及次要信息的干扰，从而提高网络模型的检测精度，其改进后的网络模型如图5所示。并在同样位置替换成SE注意力机制作对比实验。

图5 改进YOLOv4-tiny网络结构

3.3 K-means算法聚类

检测算法中的锚框尺寸的设计影响检测精度，它与具体应用场景有关。YOLOv4-tiny中的锚框尺寸是在Coco数据集上聚类得出的，且检测目标种类多尺寸差异较大，使锚框尺寸有所差异。而在自建的零件数据集中，各类零件在不同图像中形状比例大致相同，因此对所有零件的锚框进行重新聚类，使算法在训练时得到更好的训练效果。首先选择N个点做为初始聚集的簇心，分别算得每个样本点到N个簇核心的距离，获得离该点最近的簇核心，把它归属到对应的簇，所有的点都归属到簇，重新计算每个簇的中心，再次分配数据集样本，当准则函数达到收敛为止。我们采用的距离公式如下：

其中，IOU为预测框与实际标注框的交并比，anchor为真实先验框，center为聚类中心。通过反复迭代，最终得到预设锚框尺寸为（75，174）、（210，214）、（259，130）、（302，132）、（348，68）、（365，84）。

4 实验分析

4.1 实验数据准备

在训练样本标注的过程中，针对非标零件检测，目前国内外没有相关的公共数据集，且自建的零件数据集采集的数据量较少，采用数据增强的方法来进行扩充。主要方法有：对图片进行增加噪点、随机尺寸缩放以及变换图像的明暗程度。这样同一张图就有了不同形式的变化，以此扩充数据集的数量，提高训练效果。数据增强如图6。

图6 数据增强

数据集共有2076张图片，分为四大类：螺丝（标签为screw）500张，螺母（标签为nut）480张，轴承（标签为bearing）500张，轴承底座（标签为pedestal）596张，标注的软件为labelimg，且每张图片对应产生目标物体的中心点坐标以及预测框的大小，分别存为xml文件。把所得数据分成9比1，分别作为训练集与测试集。

4.2 实验环境与模型训练

深度学习对硬件环境有较高要求，通常需要使用GPU进行迭代训练。实验软硬件配置如表1所示。配合python汇编语言，PyCharm编译器进行训练。

表1 实验配置表

使用改进后的YOLOv4-tiny算法对数据集进行重训练，首先采用迁移学习的方法对模型进行预训练，先加载预训练模型，初始学习率设置为0.001，权值衰减项设置为0.0005，迁移学习先冻结主干提取特征层50个epoch后，对应训练批次为8，此时损失值从52下降到3.39。在经过解封50个epoch后，对应批次调整为4。最后损失值在1.64附近达到稳定。在训练过程当中，当模型的迭代次数增加时损失值会开始下降，当损失值达到平稳状态时，代表模型已经收敛。此时模型识别精度达到最佳。各算法损失函数如图7所示。

图7 网络模型损失曲线图

4.3 评价标准

在目标检测中常用的评价指标有准确率（Precision）、召回率（Recall）、平均精确度（AP）和各类别平均AP值mAP（Mean Average Precision）、评估速度的常用指标FPS。具体公式如下：

TP表示被正确地划分为正例的个数，即样本为正，预测结果为正。FP表示被错误地划分为正例的个数，既样本为负，预测结果为正。FN表示被错误地划分为负例的个数，即样本为正，预测结果为负。准确率代表评估捕获的成果中目标成果所占得比例，召回率代表从预测正确目标中占所有正例的比例。根据precision与recall绘制P-R曲线，AP就是该曲线下的面积，mAP就是所有类AP的平均值。它能够反映全局性能的指标。FPS是指画面每秒传输帧数，常作为目标检测评估速度的标准。

4.4 实验结果展示与分析

将训练好的网络与其他算法在同样实验环境下作对比实验，分别测试出各算法在零件数据集上的mAP值、准确率、召回率与检测速度。检测结果如表2所示。

表2 算法检测结果对比

在相同配置的环境下，共进行了3次对比实验，在表2中可以看出把CBAM注意力机制嵌入到YOLOv4-tiny的FPN网络结构，与YOLOv4-tiny相比较，召回率提高了2.38%，表明所提算法对目标漏检有所提高，准确率从90.34%提高到93.38%，表明降低了目标错检，mAP值提高了1.82%，说明总体检测性能得到提升。虽然检测速度略微下降，但对实时检测影响不大。可知所提算法可以更好地从空间和通道这两个维度去加强对零件特征的提取，使目标特征更为丰富，提高了检测的综合能力。虽然使模型参数量有所增加，但检测性能却有明显提高。

在YOLOv4-tiny中嵌入CBAM注意力机制与嵌入SE注意力机制相比，前者各项指标比后者均有提高，仅平均检测时间增加0.16ms，是由于其网络模型是在SE模块上引入了空间注意力模块。综合结果表明，所提算法定位更为精准，达到了速度与精度的平衡，可以更好地移植到设备平台。测试集中检测结果如图8所示。所提算法在散乱放置、光线较弱、叠加等复杂场景对零件的检测均取得了较好效果。

图8 检测结果示例图

5 结语

针对零件目标检测任务的检测精度低的问题，我们使用K-means算法对锚框重新聚类，并通过在YOLOv4-tiny中嵌入CBAM注意力机制得到改进YOLOv4-tiny算法，该算法结合了YOLOv4与CBAM注意力机制的优势，加强网络的感受野，提高网络的特征提取能力。根据多个检测指标的实验结果所得，改进后的YOLOv4-tiny算法精确率提高到89.84%，平均检测速度达到28帧/秒，对零件的定位精准度和识别率均有所提升。算法所得的最小外接矩形框与中心坐标可为工业机器人在零件的智能分拣工作中提供参考，但在更具挑战的零件叠加场景下的识别能力较弱，如何在该场景下进行零件检测是后续重点研究内容。