基于YOLOv5s的晶粒检测算法

张子优,韩华超,魏 东,于 霞

1.沈阳工业大学信息科学与工程学院;2.沈阳仪表科学研究院有限公司

0 引言

晶粒是精度要求很高的产品,当前晶粒定位主要通过划片机对准操作实现。张祯[1]在圆投影算法基础上设计了具有旋转不变性的圆投影模板匹配算法。J.Xu等[2]提出一种适用于各种半导体的晶粒对准算法。杨延坤[3]通过分析晶粒通道之间的偏置角度首次对准,采用数理统计重心定位方法对晶粒细对准。魏哲等[4]提出了一种将双线性插值与亚像素边缘检测结合的晶粒边缘检测算法。汪忠伟等[5]通过图像采集器获取图像,采用机器学习算法拟合形成回归算法,实现晶粒图像的对准。鲁沛昕等[6]提出基于分级金字塔算法实现目标区域快速定位,通过将高斯拟合边缘检测进行精确定位。

随着深度学习的发展,目标检测越来越多应用在工业领域,可以将晶粒定位问题转化为晶粒目标检测问题。YOLOv5s深度学习目标检测算法因为具有较高检测精度和较小参数量,广泛应用于各种工业检测中。采用YOLOv5s算法对晶粒进行目标检测面临以下问题:

1)YOLOv5s算法需要大量标注好的图像进行训练,但目前没有标准晶粒检测数据集;

2)YOLOv5s算法相比传统算法特征提取能力更强,但硬件资源要求较高,如何降低模型参数量,是需要解决的问题;

3)在降低模型参数量的同时保证检测精度的问题。

针对没有晶粒检测数据集问题,通过数据增强扩充晶圆图像,使用Labelimg软件标注晶粒,构建晶粒检测数据集。

针对硬件资源要求较高问题,使用轻量级GhostNetv2[7]网络替换YOLOv5s的原有主干特征提取网络,降低模型参数量。

针对保证晶粒检测精度问题,在特征融合网络中引入CA注意力机制[8],采用EIOU[9]替换YOLOv5s算法中原有CIOU[10]定位损失函数,提高算法检测精度。

1 算法思路

首先采集晶圆图像,建立晶粒检测数据集用于模型的训练和测试。

构建YOLOv5s-wafer算法,算法由Backbone特征提取层、Neck特征融合层和Head检测层3部分构成,如图1所示。Backbone特征提取层采用GhostNetv2网络作为主干网络,该网络由GhostV2瓶颈结构堆叠组成。Neck特征融合层将不同层次的语义特征和位置信息融合,提升网络对不同尺度对象的检测效果,在Neck特征融合层添加CA注意力机制,加强特征提取能力。Head部分包含3个检测层,对应Neck部分3个不同尺寸的特征图,检测层对不同尺寸特征图卷积运算输出位置坐标、分类结果及置信度,得到最终检测结果。

图1 YOLOv5s-wafer网络结构

1.1 晶粒检测数据的建立

由于晶粒检测没有标准的数据集,所以本文数据集是自行构建的。使用晶圆划片机在室内不同光照下进行晶粒图像采集。

原始图像共217张,数据量较少,无法让网络训练收敛。将采集的图像通过OpenCV进行水平翻转、叠加噪声等方法随机组合进行数据增强,对原始图像扩充,最终得到晶粒图像2 301张。

为了区别晶粒与背景部分,使用Labelimg数据标注软件对晶粒进行标注,将晶粒图像标注为wafer,标注完成后使用YOLO格式保存,用于后续检测。最后将标记好的图像按照9∶0.9∶0.1的比例随机分为2 083张训练集,200张验证集,18张测试集。训练集和验证集应用于算法的训练过程,测试集用于最后检测算法的识别精度。

1.2 特征提取

为了将目标检测算法部署在实际设备中,将YOLOv5s算法的Backbone主干网络替换为GhostNetv2网络,降低模型参数量。GhostNetv2网络是一种轻量级网络,它在GhostNet[11]网络基础上添加DFC注意力机制以捕捉长距离空间位置的依赖关系,达到更好的特征提取效果。

GhostNetV2网络由GhostV2瓶颈结构堆叠组成,GhostV2瓶颈结构分为2种:步长为1和步长为2,如图2所示。对于步长为1的GhostV2瓶颈结构,DFC注意力分支与第1个Ghost模块并行,增强扩展能力,2个分支乘积作为输出进入下一个Ghost模块,经过批量归一化层(batch normalization,BN),使用Add操作将其与输入特征相加,得到输出特征。对于步长为2的结构,需要使用步长为2的深度可分离卷积(depth wise separable convolutions,DWConv)对特征进行下采样,然后进行与步长为1瓶颈结构相似操作。

图2 GhostV2瓶颈结构图

具体而言,给定具有高度H、宽度W和通道数C的输入特征X∈RH×W×C,它被送到2个分支,一个是Ghost模块分支,另一个是DFC注意力机制分支。Ghost模块卷积运算产生部分输出特征图Y′,将Y′与通过简单线性变换生成的另外一些特征图连接在一起得到最终的输出特征图Y,如式(1)、式(2)所示。

Y′=X*F1×1

(1)

Y=Concat(Y′,Y′*Fdp)

(2)

式中:F1×1为逐点卷积;Fdp为深度卷积。

另一个分支是DFC模块在FC层产生的注意力图A={a11,a12,…,aHW},如式(3)所示。

(3)

式中:⊗表示逐元素乘法;F为全连接层中可学习的权重。

注意力图的2个分支得到的信息聚合后再输出,如图3所示。

图3 信息聚合过程

其输出O∈RH×W×C是2个分支输出的乘积,如式(4)所示,Sigmoid是将注意力图A归一化到范围(0,1)的缩放函数,V( )表示Ghost模块操作。

O=Sigmoid(A)⊗V(x)

(4)

1.3 特征融合

为了解决主干特征提取网络换成轻量化网络后,参数量大幅度下降但晶粒检测精度降低问题,在Neck特征融合网络部分添加CA注意力机制,加强重要特征权重,提高算法检测能力。CA模块结构如图4所示,其中X Avg Pool和Y Avg Pool 分别表示X和Y方向的平均池化。

图4 CA模块结构图

CA注意力机制分为坐标信息嵌入和坐标注意力生成2个阶段,在坐标信息嵌入阶段,对于尺寸为C×H×W的输入特征图X,使用2个尺寸的池化核(H,1)和(1,W)沿水平坐标和垂直坐标对每个通道进行编码,则高度为h的第c个通道和宽度为w的第c个通道的的输出特征图和如式(5)、式(6)所示。生成的2个特征图聚合了2个方向的特征,便于网络对目标检测物体定位。

(5)

(6)

在坐标注意力生成阶段,首先对zh和zw2个特征图进行级联操作,使用卷积变换函数F1生成中间特征映射f,如式(7)所示。

f=δ(F1([zh,zw]))

(7)

式中δ为非线性激活函数。

然后再使用卷积变换Fh和Fw将fh和fw变换为与输入X具有相同通道数的张量,再经过Sigmoid激活函数得到注意力权重gh和gw,如式(8)、式(9)所示。

gh=σ[Fh(fh)]

(8)

gw=σ[Fw(fw)]

(9)

最后,注意力权重gh和gw与输入特征图X相乘,得到注意力权重yc,如式(10)所示。

(10)

1.4 损失函数

YOLOv5s采用的定位损失函数为CIOU Loss,CIOU Loss的原理如图5所示,A为预测框,B为真值框,C为A、B的最小外接矩形。

图5 CIOU Loss原理示意图

IOU和CIOU计算公式如式(11)、式(12)所示。

(11)

(12)

式中:bA和bB分别为预测框和真实框的中心点;ρ为两点之间的欧氏距离;c为预测框和真实框的最小外接矩形的对角线长度;α为权重参数;ν为预测框与真实框宽高比的差异。

但当2个框宽高比成线性比例时,惩罚项就会失效,无法计算损失,因此采用EIOU Loss,EIOU Loss直接计算预测框和真实框的宽高,解决了CIOU Loss存在的问题,EIOU Loss公式如式(13)所示。

(13)

式中cω和ch为预测框与真实框最小外接矩形的宽和高。

2 实验结果与分析

2.1 实验环境

本文实验采用的操作系统为ubuntu 22.04,处理器为Intel Core i7-12700k,16 GB运行内存,显卡为NVIDIA GeForce RTX 3090 24 GB,开发环境Python3.8+pytorch1.10。

2.2 算法评价指标

针对晶粒目标检测实验,选取准确率(P)、召回率(R)、平均精度均值(mAP)、模型体积和参数量作为评估指标。其中准确率、召回率、平均精度均值的计算公式如式(14)、式(15)、式(16)所示:

(14)

(15)

(16)

式中:TP为算法预测为正样本且是正样本的目标;FP为算法预测为正样本但不是正样本的目标;FN为算法预测为负样本但为正样本的目标;k为预测类别数,k=1;AP为平均精度;mAP为平均精度均值。

mAP值越高,即算法检测效果越好,预测精度越高。mAP@.5表示IOU阈值在0.5上的mAP。

2.3 消融实验

为了验证YOLOv5s-wafer各模块的有效性,分别对各模块进行消融实验,实验结果如表1所示。首先引入GhostNetv2轻量网络后,平均精度均值从99.1%降为98.8%,但参数量和模型体积降低,参数量从7.013×106降为4.630×106,模型体积从13.7 MB降为10.0 MB,然后加入CA注意力模块,在参数量和模型体积基本不变情况下,平均精度均值从98.8%提高为99.0%,最后将损失函数CIOU替换为EIOU后,平均精度均值提高为99.3%。

表1 消融实验结果

2.4 对比实验结果

表2为采用单阶段目标检测算法YOLOv5、YOLOv7[12]、YOLOv8、SSD[13]、RetinaNet[14]和YOLOv5s-wafer对晶粒检测的对比实验结果。从表2可知,单阶段目标检测算法对晶粒检测的mAP较高,但模型的参数量和体积较大,无法部署到实际应用中,YOLOv5s-wafer模型参数量和体积较小且检测精度较高,具有实用价值,可以部署到实际应用中。

表2 对比实验结果

2.5 晶粒检测结果

为了体现算法有效性,从晶粒检测数据集中取出一些图片采用YOLOv5s-wafer算法检测,检测效果如图6所示。图中黑色框为晶粒检预测框,框上为检测类别和置信度,置信度均在0.9以上。由检测结果可知,该算法可以在不同光照、旋转角度,不同平整度条件下精确检测出晶粒,表现优异且稳定,满足晶粒加工的实际要求。

(a)普通晶粒

3 结束语

本文针对传统晶粒定位方法容易受外界条件影响、耗费大量资源等问题,基于YOLOv5s算法提出YOLOv5s-wafer晶粒检测算法。该算法首先通过更换轻量级网络主干降低模型参数量,然后在特征融合部分插入CA注意力机制,增强特征提取能力,最后采用EIOU定位损失函数计算预测框回归损失,提高算法检测精度。实验结果表明:该算法平均检测精度可达99.3%,模型参数量降为原来的33.88%,模型体积仅为9.9 MB,符合工业检测上的应用需求。