基于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,符合工業檢測上的應用需求。