基于YOLOv5s的金属制品表面缺陷的轻量化算法研究

贾云飞， 郑红木， 刘闪亮

(中国民航大学 电子信息与自动化学院， 天津 300300)

当今，计算机视觉中目标检测技术飞速发展，这一技术在工业产品缺陷检测中的应用已经十分广泛，它不仅可以帮助企业进行产品质量把关，还可以降低人力成本，对企业的发展十分有利。目标检测网络主要分为two-stage网络和one-stage网络，其中two-stage网络中的经典网络如SPP-NET[1]、R-CNN[2]、Faster R-CNN[3]等。这些网络通过算法生成一系列预选框，同时进行区域分类，然后对预选位置进行修正，精确目标位置。而one-stage网络如YOLOv3[4]、EfficientDet[5]、SSD[6]等网络则可以联合区域提案和分类预测，对输入网络的图像特征进行提取并直接输出目标的类别和位置信息[7]。现阶段，由于one-stage网络的高效性，该类网络发展迅速。郝用兴等[8]通过对Faster R-CNN中提取网络，并在特征提取网络中对VGG16和Resnet50进行模型融合，在瑕疵大小分布不均匀的铝材表面瑕疵检测上更具优势；
程婧怡等[9]针对金属表面目标尺寸小和特征不清晰等问题，通过在YOLOv3网络中加入DIOU边框回归损失和K-means++聚类分析筛选最优回归框，该算法在定位金属板表面小缺陷目标的精度上更具优势；
Li等[10]通过引入MobileNet网络对SSD的特征提取网络进行改进，在电子产品表面缺陷检测上取得了不错的准确率以及检测速率。

上述研究分别对不同的目标检测网络进行改进并对工业产品或电器设备进行缺陷检测，达到了可观的效果，但没有针对企业内需要体积更小、计算参数更少的模型进行研究，这样的模型在成本预算低、算力相对不足的设备中也能达到不错的检测速度。基于此，本文结合深度学习目标检测的方法，在原有的轻量级YOLOv5s网络的基础上，通过结构裁剪的思想，基于网络的BN(batch normalization)层进行模型裁剪，得到更轻量级模型的同时也具备了优异的检测效果，并对裁剪后的模型进行训练，并使用TensorRT进行推理加速，进一步提高了模型的检测速率。

1.1 YOLO系列算法简介

目标检测算法YOLO由Redmon等[11]在2016年提出。其结构包含了24个卷积层和2个全连接层，输出层的激活函数为线性函数，其他激活函数为ReLu函数。其核心思想在于将整张图输入网络，输出层会直接输出对目标边界框的位置预测和类别回归。之后在此基础上提出了采用Darknet-19作为主干网络的YOLOv2算法模型，YOLOv2相比YOLOv1改进了准确度、预测对象广度并具有更快的速度[12]。而之后的YOLOv3则更是成为了YOLO系列的经典算法，其采用Darknet-53作为主干网络，包含53个卷积层，且层间残差结构连接，组成残差层[4]。YOLOv4和YOLOv5则是在YOLO系列原作者退出研究之后由后人继续研究完成，对比之前的YOLO系列网络改动相对较小，其中YOLOv5s模型大小仅有27 MB，相比之前的YOLO模型大幅度减小。

1.2 YOLOv5s网络模型

YOLOv5s的特征提取网络为CSPDarknet(cross stage partial darknet)，由CSP结构和Darknet网络相结合构成。其网络结构主要分为3部分，即Backbone、Neck和Head,其中Backbone部分主要进行输入图像特征提取；
Neck部分对图像特征进行融合；
Head部分对图像特征进行结果预测并输出[13]。网络模型中的CBL结构主要由提取特征的卷积层、BN层和Leaky ReLu激活函数构成，如图1所示。YOLOv5s还设计了2种CSP网络结构，其中CSP1_N用于Backbone中，CSP2_N用于Neck部分，并且采取交叉结构连接，可以在减少计算量的同时保证运算的精度。

图1 YOLOv5s中的2种CPS结构Figure 1 Two CSP structures in YOLOv5s

2.1 BN层概述

在卷积神经网络中，网络通过前向传播和反向传播共同作用来进行权值更新，在进行一次反向传播时，网络会对之前的前向传播训练得到的参数进行修正，此时各层网络参数权重会同时更新。由于每层的输入都会受到前面所有层输入的影响，网络参数微小的变化都会随着网络深入而放大，所以需要较低的学习率和初始参数，从而降低了学习效率。2015年，Google提出了BN层，通过这个算法对网络各层的输入进行归一化，将输入的批数据进行均值为0、方差为1的归一化处理[14]，对各层输入数据分布进行了控制，但对网络上一层的输出数据经过归一化处理之后送入下层网络的数据，在上层网络学习的特征会因为数据的分布变化而遭到破坏，因此BN层添加了缩放参数γ和平移参数β，来恢复上层网络学习的特征。因此，BN层的总体计算过程大致如下：首先，对输入数据进行求均值和方差；
其次，对数据进行标准化处理；
最后，训练缩放参数和平移参数，输出通过参数线性变换得到新值。通过这种方法使得每层数据分布相同，能提高网络学习效率，具体公式如下：

(1)

(2)

(3)

(4)

2.2 网络模型裁剪

一般来说，神经网络中层数越深，网络结构越精细，通过训练网络得到的输出结果也越精准。但与此同时,也带来了计算参数更多、计算量更大的弊端，而且在训练数据量较少的情况下，模型相对数据集复杂度过高，容易产生过拟合现象，虽然产生的网络更深更精细,但最终训练模型的效果反而不佳。与ImageNet数据集中1 400多万张样本数据相比，工业产品的表面瑕疵检测面临样本小、图像样本效果不佳等问题，所以在工业检测领域，通过模型剪枝不仅能够解决网络计算量大的问题，还可以解决网络过拟合现象，提升计算速度，在降低训练参数量大小的同时，提高网络模型的精确度。

神经网络经过训练后得到模型权重，通过前向传播和反向传播生成了一系列的权重参数，在使用验证数据进行推理时，输入数据进入神经网络并通过上述参数进行推理计算，而其中的参数有一部分权重占比小，对输出结果不起作用或是仅仅起到微小的作用，而这样的参数存在对计算结果虽然影响很小，但却加大了计算量，因此对其进行裁剪可以使得输入数据通过最有效率的参数空间路径得出输出结果，在提高计算效率的同时，也几乎不会降低网络推理效果，且对数据集规模小的数据还能降低训练过拟合带来的影响，提高网络的推理效果。神经网络的裁剪方法具体可分为非结构裁剪和结构裁剪。非结构裁剪是一种细粒度的裁剪方法，精度相对较高，但对硬件平台和特定算法依赖程度较高；
而结构裁剪是针对层级的裁剪，是一种粗粒度的裁剪方法，对硬件平台和特定算法依赖程度低，可行性强，因此，本文采用结构裁剪的方法对YOLOv5s进行裁剪。

对原始YOLOv5s网络进行训练后，模型BN层的γ值一般情况下呈现近似正态分布，难以对训练后的网络模型进行裁剪，所以在模型裁剪之前需要先对YOLOv5s网络BN层中的损失函数的γ值添加L1正则化约束。通过L1正则化约束可以使得网络模型稀疏化，通常也称作稀疏训练，此时损失函数：

(5)

对于本文所研究的工业瑕疵检测这类问题的小样本数据集来说，稀疏训练可以使得网络中很多参数分布接近于0，降低了训练后的模型复杂度，从而可以减少过拟合现象在训练数据集上的表现，提升模型在验证数据集上准确率。由文献[15]中的实验结果可以看出，当采用复杂度更高的YOLOv5模型时会在数据集上表现出过拟合现象，导致在验证集上模型预测准确率降低。文献[16]中的实验表明了稀疏训练的模型可以学习到更泛化的特征，具备鲁棒性和易用性。在对网络模型进行稀疏训练后，BN层中的一部分γ值会趋向于0分布，这部分权值对推理结果影响甚微，此时可以对其所在的BN层以及相邻上层的卷积核和向下层输出的通道进行裁剪，以此降低模型参数数量，实现模型的轻量化，γ裁剪过程如图2所示。

图2 γ裁剪过程Figure 2 γ pruning process

TensorRT(以下简称TR)是NVIDIA基于深度学习加快推理的高性能平台。TR可应用于图像分类、目标检测和图像分割等计算机视觉领域的相关推理任务中。TR作为可编程的推理器，能够为训练好的模型提供推理加速[17]。TR包含了构建和部署2个阶段，构建阶段是TR针对网络配置进行优化，产生一个优化决策并用于网络的前向传播，其对神经网络进行了一系列转换和优化，其中优化包括删除未使用的网络层以避免不必要的计算，对多个层进行融合形成单层。当训练好的模型进行网络推理时，模型网络的各个层都会进行函数的调用，被调用的函数将在硬件设备上执行，而这些函数大多计算速度快，当多次调用函数时，对推理速度起决定性影响的则是启动时间，函数的频繁调用启动会导致更多的内存传输，而TR将网络和张量进行了融合操作，融合包括垂直方向和水平方向2种方式，其中垂直融合将conv层、BN层或偏移层、激活层整合为一个CBR；
水平融合则是使得几个输入相同(例如1×1CBR)融合在融合后网络层数减少，降低了网络的深度和宽度，并且可以减少对函数的调用，提升推理速度。除此之外，TR还可以将连接层移除，将结果输出到预先分配的缓冲区直接输入到下一网络层中。部署阶段则是设备通过长时间运行程序或服务，等待待处理数据的提交，当收到待处理数据时，执行构建阶段产生的优化决策并进行推理，将结果通过设备输出。

4.1 实验数据来源

本文实验数据来源为东北大学采集并发布的NEU-DET钢表面6种缺陷的数据集，这6种缺陷分别为氧化皮、斑块、开裂、麻点状表面、杂质和划痕，共计1 800张图像，其中每一类缺陷包含300个样本。将1 800张图像按照8∶2的比例随机分为训练集和验证集，即训练集1 440张，验证集360张，样例如图3所示。

图3 数据集样例图Figure 3 Sample dataset diagram

4.2 实验平台搭建

本实验采用高性能实验台和低性能实验台来进行实验，高性能实验台配置为ubuntu20.04LTS操作系统，硬件方面采用的处理器为Intel Xeon Gold 5218R，GPU为NVIDIA Quadro RTX5000。低性能实验台配置操作系统为windows10，处理器为Intel®CoreTMi5-8265U，GPU为NVIDIA GeForce MX250。实验使用框架Pytorch(1.7版本)，GPU加速计算方面采用CUDA10.1、CUDNN7.6.0和若干第三方Python库支持模型训练及运算。

4.3 实验流程设计

首先，需要对YOLOv5s网络中添加L1正则化约束；
其次，在搭建好的训练平台进行模型稀疏训练，减小BN层的γ分布，针对接近0分布的γ值，对相应的BN层及其相邻卷积核和卷积核通道进行裁剪，得到裁剪后的新模型；
再次，对该模型再进行训练，得到最终的训练模型并进行TR加速；
最后，在上述搭建好的实验平台上进行模型推理测试，实验流程如图4所示。

图4 本文模型实验流程图Figure 4 Flow chart of the experiment of this model

5.1 评估指标及目标识别过程

目标检测作为机器视觉的一个任务分支，评估指标除了常规有监督机器学习中的精确率P和召回率R之外，还使用各类别平均精确度均值(AP)作为网络模型训练结果好坏的评估标准。类似于其他机器学习算法，目标检测需要正确预测出图像物体所属类别，但同时还需要对其进行位置预测，因此需要引入模型预测边框与真实标签边框的交并比指标IOU，如图5所示，IOU即为左图阴影面积和右图阴影面积比例大小[18]。

图5 锚框的交并集示意图Figure 5 Intersection and union of the anchor box

在图5中，实线框为标注框，虚线框为模型预测框。当对预测情况进行判断时，常用IOU阈值为0.5，即当预测结果IOU大于阈值时，即认为该预测框瑕疵类别为标注框类别。本文数据集缺陷预测的精确率和召回率计算公式如下：

(6)

(7)

式中：TP为模型预测为正样本而验证集中标签为正样本的数量;FN为模型预测为负样本而验证集中标签为正样本的数量;FP为模型预测为正样本而验证集中标签为负样本的数量;TN为模型预测为负样本而验证集中标签为负样本的数量。

在模型对单个类别进行预测时，对IOU置信度取不同的阈值时，P和R的值也会不断在0～1内变化，将改变置信度阈值得到的上述2个值在二维坐标系的变化曲线称为PR曲线，而PR曲线与坐标轴围成的面积即为单个类的AP值，对需要预测的n个类的AP值取平均值即为mAP，公式如下：

(8)

5.2 稀疏训练以及裁剪训练γ分布结果

对原始YOLOv5s模型进行训练，得到γ值分布如图6所示，由图6可以看出,原始网络训练得到的结果接近正态分布。

图6 YOLOv5s训练γ值分布图Figure 6 γ value distribution of YOLOv5s training

对模型进行稀疏训练，得到γ值分布如图7所示，在稀疏训练后绝大多数γ值降低，有约40%的γ值降低到0附近，可以对其进行裁剪。

图7 稀疏训练γ值分布图Figure 7 γ value distribution of sparse training

将稀疏训练后得到的模型进行裁剪，在反复裁剪微调后，最终对33%的接近0的γ值进行裁剪，并对裁剪后的模型进行训练，得到γ值分布如图8所示。在裁剪时为了保证tensor维度可加，对残差结构不进行裁剪，保留原始通道，否则会导致残差结构两端维度不同，前向传播无法进行[19]。同时在裁剪时为了保证网络结构的完整性，裁剪阈值不能超过各个网络层中γ参数最大值的最小值。

图8 裁剪模型训练γ值分布图Figure 8 γ value distribution of pruning model

5.3 对比实验结果

对实验数据集使用几种网络进行训练，其中包括经典two-stage网络Faster R-CNN、经典one-stage网络SSD和YOLOv3-SPP，还有本文使用的YOLOv5s原始网络和裁剪后网络，对各个网络在搭建好的高配置实验台中进行训练，epoch设置为100，batch-size设置为16。网络训练过程中mAP稳定后的15个epoch如图9所示。由图9可以看出，本文模型mAP略高于原始的YOLOv5s模型，与经典two-stage网络Faster-RCNN几乎不相上下，本文网络相比其他网络的优势在于模型的轻量化，训练后权重文件小，可以在低算力硬件设备中展现出不错的推理速度。对每种模型采用验证集中的10张图片进行模型推理测试并计算出单张图片的平均推理时间，对图9中的mAP值取平均，结合模型训练最终的权重文件大小，得出实验结果如表1所示，模型识别效果如图10所示。

图9 各网络训练情况对比Figure 9 Comparison of each network training

表1 各网络模型实验结果Table 1 Experimental result each network model

图10 最终模型检测效果图Figure 10 Final model detection effect

为了体现本文方法的普适性，选用kaggle数据集网站公开的潜水泵叶轮的俯视图数据集进行普适性实验，数据集共有2 400张图片。将数据集图像中缺陷种类分为深槽和毛刺类别并标注数据集，按8∶2将数据集分为训练集和验证集。在高性能实验台上进行训练，最终采用YOLOv5s和本文模型分别在该数据集上得到了72.7%和73.4%的mAP，证明了本文方法在其他数据集上也有优秀的表现。检测效果如图11所示。

图11 普适性实验检测效果图Figure 11 Universal experimental detection effect

由于YOLOv5网络在输入端采用mosaic数据增强，对输入数据集进行随机缩放、裁剪、排布的方式拼接，使其对小目标检测效果较SSD和YOLOv3-SPP网络有所提升，Faster R-CNN模型采用双阶段的目标检测方式，在检测头部先进行候选框提取，再对候选框进行筛选得到预测框，虽然获得了一定的检测精度，但是速度较慢，难以满足流水线产品的快速检测需求，采用本文裁剪方法后的YOLOv5s模型减轻了原模型在本文小样本数据集上的过拟合现象，在保证精度的同时，减少了运算参数，提升了推理速度。由表1可以得出，本文裁剪后的模型在预测精度没有损失的情况下，模型文件大小相比原YOLOv5s模型减少了67.9%，推理时间提升了11.3%，检测效率得到了明显的提升。如表2所示，为展现本文裁剪模型在低算力硬件设备上的推理时间优势，在上述搭建的低性能实验台上对NEU-DET验证数据集中10张图片进行推理时间测试并取平均值，从结果可以看出本文模型在低算力硬件设备上推理时间提升更大。

表2 单图推理时间对比实验结果Table 2 Inference time of single figure comparison experiment result

最后对本文训练得到的模型文件转onnx格式并采用TR推理框架进行加速，在高性能实验台中测试表现可以达到单图5.82 ms的推理时间。本文方法与同样采用本文NEU-DET数据集的文献[9]、文献[20]相比,mAP分别提升了6.56%和0.5%，且本文模型更轻量化，检测速率更快。

本文从企业需要更低的硬件成本的角度出发，通过深度学习网络模型中结构裁剪的思想，提出一种基于YOLOv5s轻量化改进的网络并结合TR加速推理框架。该方法很好地实现了模型权重文件大小缩减，降低运算量，且保证精度没有损失，可以满足产品缺陷实时检测的需求。下一步将基于本文模型继续研究该模型在集成于低算力嵌入式设备和移动端设备的优势，并将本文模型与实际工业场景应用相结合。

猜你喜欢 样本预测检测 无可预测黄河之声(2022年10期)2022-09-27选修2-2期中考试预测卷（A卷）中学生数理化(高中版.高二数学)(2022年4期)2022-05-25选修2-2期中考试预测卷（B卷）中学生数理化(高中版.高二数学)(2022年4期)2022-05-25选修2—2期中考试预测卷(A卷)中学生数理化·高二版(2022年4期)2022-05-09规划·样本领导决策信息(2018年16期)2018-09-27人大专题询问之“方城样本”人大建设(2017年10期)2018-01-23随机微分方程的样本Lyapunov二次型估计数学学习与研究(2017年3期)2017-03-09必修二 Modules 1—6综合检测题时代英语·高一(2016年6期)2017-02-27“整式的加减”检测题中学生数理化·七年级数学人教版(2016年10期)2016-12-22“整式”检测题中学生数理化·七年级数学人教版(2016年10期)2016-12-22