基于YOLOv5的表面缺陷檢測優化算法

渠 逸,汪 誠,余嘉博,孔亞康,陳賢聰

(空軍工程大學基礎部,西安,710051)

隨著現代工業的發展,各行業對金屬材料性能的把控能力愈加重要,鋼材作為目前產量最高、應用最廣的金屬材料,在交通運輸、機械制造、航空航天和國防工業等領域其他金屬材料是無法替代的。隨著我國工業水平的迅速發展,市場需求不斷變化,各領域對鋼材的質量要求也在不斷提升。鋼材在生產制造的過程中,受到生產設備、工藝水平等因素制約,鋼材表面會產生如裂紋、劃痕、凹坑、斑塊等無規律的缺陷[1],導致在后續應用中,材料的力學性能、抗腐蝕性、抗磨損性能都會有不同程度地下降,嚴重的缺陷甚至會引發安全事故[2]。因此,探索一種準確、高效率的表面缺陷檢測方法成為了當前工業發展的迫切需求。

近年來,隨著深度學習的發展,計算機視覺廣泛應用在各個領域,對材料表面的目標檢測是計算機視覺中重要的一環,該項技術大幅降低了人工成本,相較于傳統人工檢測手段具有穩定、高效、準確度高等優勢。

當前基于深度學習的目標檢測算法主要分為one-stage和two-stage 2種。one-stage目標檢測算法的檢測速度快,但精度相對較低,代表算法有SSD[3]、YOLO[4]等。two-stage目標檢測算法檢測精度較高,但檢測速度慢,以Fast R-CNN[5]、Faster R-CNN[6]、Mask R-CNN[7]等為代表。對于一般的物體表面缺陷檢測,one-stage目標檢測算法的精度已經可以滿足檢測需求,又因為其具有更高的速度,所以被廣泛應用[8]。YOLO網絡作為one-stage檢測算法里最受歡迎的系列,近年來被大量學者使用和優化。Fang等[9]通過改進的YOLOv3算法實現了對金屬表面的缺陷檢測,通過對輸入端的圖像進行處理和損失函數優化提升了檢測效果。李維剛等[10]使用加權K-means聚類算法,并通過融合淺層特征與深層特征的方式改進了YOLOv3網絡,在帶鋼表面缺陷數據集上的檢測精度比原始網絡提高了11%。Li等[11]通過構建全卷積YOLO檢測網絡檢測帶鋼表面缺陷,為此類檢測提供了端到端的檢測方案,并取得了97.55的mAP。Zhang等[12]提出了改進的YOLOv5算法,該算法在原有算法的基礎上增加了微尺度檢測層,并加入了CBAM關注機制,以控制小目標缺陷等缺陷的特征信息丟失。

大量研究已經初步使目標檢測的精度和檢測速度得到提升,但各類材料的應用場景變化較大,為滿足各類環境下個人檢測任務需求,檢測設備需盡可能輕量化。在此基礎上,本文針對工業應用場景,以提升檢測精度與速度、減小模型體積為目標,提出了改進的YOLOv5算法,并在NEU-DET數據集上進行測試。

1 實驗基礎

1.1 數據集

本文所涉及的實驗數據集均采用東北大學宋克臣團隊制作的鋼材表面缺陷數據集(NEU-DET),數據集內包含6個缺陷類別：crazing,inclusion,patches,pitted surface,rolled in scale,scratches,每類缺陷300張識別圖像,各類別的缺陷標注信息通過xml文件保存,共計1 800張灰度圖像,4 189個檢測缺陷邊界框。

將數據集內1 800張圖像按照8∶1∶1的比例劃分為訓練集、驗證集和測試集,得到訓練圖片1 440張,驗證圖片和測試圖片各180張,數據集內各類缺陷圖片如圖1所示,其中缺陷目標所在位置位于紅框標注處。

(a)crazing (b)inclusion (c)patches

1.2 YOLOv5

YOLO是一種通用的one-stage目標檢測算法,側重于推理速度,被廣泛應用在圖像或視頻的目標分類、定位工作中。迄今為止,學者通過不斷改進,發布了各種版本的YOLO模型,其中YOLOv5便是在YOLOv4的基礎上改進而來,其在輸入端的模型訓練階段加入了Mosaic數據增強、自適應圖片縮放和錨框計算功能;融合了Focus結構與CSP結構改進了基準網絡;在Neck部分添加了FPN+PAN結構;改進了預測篩選框的DIoU NMS和訓練時的GIOU Loss,對比前代在精度和速度上都得到了很大提升。YOLOv5目前已經更新到7.0版本,根據不同的使用場景,有n/s/m/l/x 5個模型可以使用,由n至l版本,模型大小和運算量逐步增加,精度也同時提高。根據本文的實驗對象,經過對比試驗后,采用體積相對較小、速度較快的s版本模型。

2 算法改進

2.1 添加目標檢測頭

DyHead是由Dai等[13]提出的一種基于注意力機制的目標檢測頭,通過將多頭self-attention注意力機制結合在尺度感知的特征層、空間感知的空間位置、任務感知的輸出通道處,并將其統一在一個框架內,以此在沒有多余算力開銷的情況下有效提升目標檢測頭的檢測性能。其結構如圖2所示。

圖2 DyHead結構

從圖2可以看出,DyHead將尺度注意力、空間注意力和任務注意力疊加在了一個模塊內,將Head部分的輸入,也就是General View視為一個三維的tensor,其中L為不同尺度(即不同層次和階段)的特征圖;S為空間位置信息,也就是特征圖的寬高乘積;C為通道信息。DyHead使用分離式的注意力機制,使以上每個維度都能獨立通過注意力機制進行特征感知,3個注意力機制如下：

1)Scale-aware Attention為添加在level維度的尺度感知注意力,不同level的特征圖所對應的目標尺度也不同,該機制的引入可以增強模型的尺度感知能力。

2)Spatial-aware Attention為添加在spatial維度的空間感知注意力,當空間位置發生變化,所對應的便是檢測目標的幾何變換,該機制的引入可以增強模型的空間位置感知能力。

3)Task-aware Attention為添加在channel維度的任務感知注意力,不同的channel所對應的檢測任務也不同,該機制的引入可以增強模型對不同任務的感知能力。

2.2 改用aLRP Loss

一般的目標檢測任務需要通過結合平衡超參數,來聯合優化分類目標與定位目標,但這種方法在定位任務和分類任務之間缺乏關聯,無法同時兼顧定位精度和分類精度;在訓練過程中,對超參數的調節較為困難,導致在需要較長時間訓練的網絡內收斂變慢;容易受到正負樣本不均勻或者正負樣本異常值變動的影響,在參數調節上需花費更多精力[14]。為解決以上問題,Oksuz等[15]提出了aLRP Loss。這是一種基于ranking的損失函數,具有統一、有界、平衡等特點,通過對AP Loss的擴展,建立一個能夠優化不可導排序損失的通用框架,能有效用于目標檢測中的分類和定位任務。與精確度和AP損失之間的關系類似,aLRP Loss為正樣本LRP值(lLRP(i))的平均值,如式(1)所示：

(1)

式中：P為精確率。LRP值定義為：

rank(i)=NFP+NTP

(3)

式中：NFP為錯誤預測的數量;NTP為正確預測的數量;當k≥0時,H(k)=1,否則H(k)=0;τ為TP標記閾值,一般取0.5。

相對于傳統損失函數,aLRP Loss具有以下優點：

1)在分類和定位任務中首次采用排名策略,將獨立工作的2個分支合為統一整體,兩者同時調用檢測器的全部輸出,從而同時兼顧檢測精度和定位質量。

2)相對于目前SOTA方法中的6個超參數,aLRP Loss僅有一個不需要調試的超參數,能夠減少參數調節的工作量,并使模型更容易訓練。

3)aLRP Loss由AP Loss拓展而來,能夠替換AP Loss、SmoothL1 Loss等損失函數,并以單個損失函數參與訓練,具有出色的正負樣本策略,以提高網絡性能。

2.3 基于FasterNet輕量化模型提出C3-Faster

為提高算法的檢測速度,Chen等人[16]提出了一種新的部分卷積,其基本思想是在卷積操作中只對輸入中非零的部分進行卷積,對輸入中全零的部分不進行操作,可以同時減少計算冗余和內存的訪問,更有效地提取空間特征,并在其基礎上提出了FasterNet,在不影響檢測精度的同時,大幅提升了檢測速度。

在此基礎上,本文提出了C3-Faster,替換了YOLOv5中用于提取特征的C3模塊。在C3-Faster中調用了C3模塊的構造函數,初始化部分基本參數,并引入了Partial Convolution部分卷積層和Faster_Block模塊,二者都使用了類似PConv[17]的思想,其通過使用一個額外的掩蔽矩陣實現,該矩陣與輸入圖像大小相同,其中缺失或遮擋的的像素位置用0填充,而正常的像素位置用1填充。隨后卷積核僅對被預設為1的像素進行卷積運算,其中在掩蔽矩陣中對應為0的像素位置,卷積核中相應的權重也被設為0。這意味著在卷積計算過程中,這些被掩蔽的權重不會對輸出特征圖產生影響,不但降低了運算量,在處理不完整的輸入數據時也能提供更好的結果,從而提高了模型的魯棒性和泛化能力,可以增強模型的非線性表達能力,提高檢測精度;Faster_Block通過部分卷積來處理輸入特征圖,從而減少計算量和模型參數。

2.4 添加輕量級上采樣算子CARAFE

上采樣在目標識別算法中可以將低分辨率的物體特征圖擴展到更高分辨率,在卷積神經網絡中,高層特征圖中的目標特征比較抽象,分辨率較低;低層特征圖中物體特征更為精細,分辨率更高。通過將低層特征圖上采樣到高層尺寸,可以將低層精細的特征和高層抽象的特征結合起來,從而更好地捕獲物體的特征,提高識別的準確性。

為了深度神經網絡中實現特征圖的高效上采樣,Wang等[18]提出了輕量級通用上采樣算子(content-aware reassembly of features,CARAFE),這是一種通用性強、輕量化、高效率的模塊,相比于傳統上采樣模塊,CARAFE在更輕、更快的條件下,感受野進一步擴大,不依靠亞像素鄰域運行,而是在更大的接受域內融合信息,并且可以針對特定內容動態生成自適應內核,能夠更優秀地進行內容感知。

CARAFE運行的具體步驟是在每個位置為中心的預定義區域內,使用加權組合來重新組裝特征,并通過內容感知的方式生成權重,并且上述的每個位置都會存在多組生成的上采樣權重,最后將生成的權重重新排列為一個完整的空間快,以此進行特征上采樣。具體來說,給定一個大小為C×H×W的特征圖X和一個上采樣率σ(假設σ為一個整數),將生成一個大小為C×σH×σW的新特征圖X′,其原理如式(5)、(6)所示：

Wl′=ψ(N(Xl,Kencoder))

(5)

對于輸出X′的任一個目標位置l′=(i′,j′),在輸入X中都能找到一個相應的源位置l=(i,j),其中N(Xl,k)為以l為中心的特征圖X的k×k子區域,Kup為重組核的大小,同時考慮到性能和效率,使用核大小的卷積層Kencoder一般取值為Kup-2。

式(5)中,內核預測模塊ψ根據Xl的子區域,為每個l′預測該位置上的內核Wl′,即輸出大小為Cup×H×W的重組核,Cup為：

式(6)中φ是內容感知重組模塊,可以將Xl的子區域與內核Wl重新組合,如式(8)所示：

(8)

式中：r=Kup/2

2.5 改進的5種模型

基于以上4點改進方向,本文共提出5種改進的YOLOv5模型如下：

1)YOLOv5-D：為YOLOv5網絡添加基于分離式注意力機制的DyHead目標檢測頭;

2)YOLOv5-A：改用由AP_Loss拓展而來的aLPRLoss損失函數;

3)YOLOv5-F：將YOLOv5算法中的C3模塊替換為基于FasterNet網絡提出的C3-Faster;

4)YOLOv5-C：為YOLOv5網絡添加CARAFE輕量級上采樣算子;

5)YOLOv5-DAFC：將Dyhead目標檢測頭、aLPRLoss損失函數、C3-Faster模塊、CARAFE上采樣算子同時部署在YOLOv5網絡中。

改進后的YOLOv5-DAFC網絡結構如圖3所示：

圖3 YOLOv5-DAFC結構

3 實驗

本文實驗環境基于Windows操作系統,運行內存128 GB,使用Pytorch 1.11.0作為深度學習框架,Python版本為3.8,CUDA版本為11.2。CPU為Intel(R) Xeon Gold 6330 CPU @2.00 GHz。CPU為NVIDIA GeForce RTX 2080Ti 11G。

3.1 實驗評價標準

本文從模型自身參數、模型檢測參數兩方面進行效果對比。

其中模型自身參數包括以下幾類：模型參數量,即網絡結構中所有可訓練參數的數量,通常是衡量模型復雜度和計算所需資源的指標;FLOPs,指模型在推理期間進行的浮點運算次數,一般采用GFLOPs,通常被用來評估模型的計算復雜度;權重體積,指模型訓練后所生成權重文件的大小。

模型的檢測參數包括以下幾類：精確率,指模型正確預測的目標數量與總預測目標數量之比;召回率,即在所有真實目標數量中,模型準確預測的目標數量占比。AP[19]@0.5(IoU閾值為0.5時各個類別的平均AP值)而mAP則是在不同召回率下的精確率的均值,這一指標考慮了準確率和召回率的綜合表現,是當前主流的目標檢測評價指標。相關計算公式為：

(11)

(12)

式中：P、R分別為精確率和召回率;N表示目標類別數;Tp為IoU>0.5時的檢測框數量(同一個Ground Truth值計算一次);Fp為IoU≤0.5時的檢測框數量(或者是檢測到同一個Ground Truth的多余檢測框的數量);FN為指沒有檢測到的Ground Truth數量;AMP為均值平均精度,即mAP。

3.2 訓練參數設置

YOLOv5使用了K-means聚類算法來生成錨框(anchor box),通過統計訓練集中不同大小的目標,自動生成合適的K值(錨框的數量),以便網絡能夠檢測不同大小和比例的目標。本文針對目標數據集與服務器性能,初始學習率設為0.01,使用隨機梯度下降法(stochastic gradient descent,SGD),warmup_epoch設為5,學習率動量為0.98,權重衰減系數weight_decay設為5e-5,work-number設為16,原始YOLOv5模型batch-size設為32,后續訓練根據模型大小與顯存冗余調整至最優。依次對原YOLOv5模型的s版本與5個改進模型進行訓練與測試,對比分析各模型的參數和性能。

3.3 訓練過程損失

在訓練過程中,本文提取了YOLOv5原模型與YOLOv5-DAFC模型的損失函數值,如圖4所示,其中綠色線條為原版YOLOv5模型訓練損失,橙色線條為改進后的YOLOv5-DAFC模型訓練損失。3個不同的損失變化圖可以客觀地描述本文中檢測算法的訓練變化過程。

圖4中Box loss是定位損失,通過計算預測邊界框與實際邊界框差異所得,用于衡量模型對物體位置的預測能力。Obj loss是置信損失,通過Focal Loss計算所得,用于衡量模型對預測物體與背景的置信度誤差。Cls loss是分類損失,通過計算分類預測結果與實際標簽的交叉熵損失函數所得,用于衡量模型對每個類別的分類能力。

可以看出,在訓練epoch達到300時,2個模型均已收斂,損失值趨于平穩。添加了aLRP Loss的改進模型損失曲線與原版區別較大,在Box loss曲線中,DAFC模型在訓練初期出現了過擬合[20]現象,相比原版的平滑下降,呈階段性下降趨勢,并且收斂速度明顯快于原版,僅在50 epoch左右就已經進入收斂狀態,且收斂后曲線平穩度高于原版;在Obj loss曲線中,DAFC模型訓練初期仍有過擬合現象存在,但模型收斂后曲線平穩度仍高于原版。可以看出,引入了aLRP Loss的DAFC模型在收斂速度和收斂后曲線穩定性方面均高于原版模型。

4 實驗結果對比

4.1 模型性能對比

為確定每個優化結構對網絡性能的影響,本文進行消融實驗,分別對YOLOv5、YOLOv5-D、YOLOv5-A、YOLOv5-F、YOLOv5-C、YOLOv5-DAFC 6個模型進行訓練,得到不同模型自身參數與測試性能如表1所示。

表1 不同模型的參數與檢測效果對比

通過分析表1實驗結果可知,在網絡中增加DyHead目標檢測頭、用C3-Faster替換C3模塊后,網絡整體的訓練參數量和計算復雜度會有較為明顯的變化,前者增加了8%的參數量和7.5%的浮點運算量,將AP@0.5提升了1.101%,同時權重體積增加了4.3%;后者則減少了17.62%的參數量和20.63%的浮點運算量,同時將AP@0.5提升了2.825%,提升效果顯著。

而在改用ALRPLoss損失函數與增加CARAFE上采樣算子的A、C模型中,相應的改進主要作用在檢測精度的提升上,A、C兩模型相比于原版模型,AP@0.5分別提升了0.49%和1.071%。

DAFC模型則是4個模型優點的整合,保留4個改進點對檢測精度提升的同時,綜合減少模型參數量與計算復雜度。相比于原模型,改進的DAFC模型在參數量減少了11.25%,計算復雜度減少了13.75%,權重體積減少了10.72%的同時,AP@0.5增加了4.174%。可以看出,改進的DAFC模型同時發揮了4個改進點的優勢,在盡可能節約內存占用與計算資源的同時,做到了檢測效果的最佳優化。

4.2 不同缺陷檢測性能對比

為研究本文所提出的模型對于不同類別缺陷的檢測效果,摘取實驗數據如圖5所示。

圖5 各類缺陷檢測效果

從圖5的分析結果可以看出,對于本數據集內的6種缺陷類別,最優檢測效果均分布在5個改進模型上,其中C模型對pitted surface和scratches類別缺陷檢測效果最佳,檢測精準度分別提升了0.5%、1.6%;對于crazing、inclusion、patches、rolled in scale 4種缺陷,均為DAFC模型檢測效果最佳,相比于原版模型,改進后的DAFC模型對crazing的識別準確率提升了1.5%,inclusion提升了13.7%,patches提升了16.6%,rolled in scale提升了8.6%,scratches提升了0.9%。改進后的模型對本數據集整體檢測效果大幅提升,并對部分難檢目標的檢測效果提升明顯,對實際應用場景的可用性有實質性的加強。

4.3 檢測效果對比

為了直觀了解改進后的模型對目標缺陷的檢測效果提升,使用原模型與改進后的5種模型對測試集180張圖片進行了檢測,YOLOv5s與改進后的YOLOv5-DAFC模型對部分測試圖片識別效果如圖6所示。

(a)YOLOv5部分檢測效果圖

通過對比可以看出,原版模型對較小目標容易發生漏檢,對較大目標存在定位錯誤,并且存在重復描框、同一目標檢測為多個等問題,改進后的DAFC模型相對于原版YOLOv5,感受野增大后對較大目標識別效果提升,對較難檢測的目標更為敏感,可以有效避免目標漏檢的情況發生,并且對于緊鄰的多個目標、疊加目標的檢測策略更加智能,顯著減少了對同一個目標重復描框、分段描框的情況和將多個目標識別為同一對象的情況,有效避免原版模型的檢測錯誤問題。

4.4 主流檢測算法性能對比

為了進一步評估本文提出的模型的性能,在相同實驗條件下,在同一數據集上選擇了SSD、RetinaNet、FCOS、YOLOv3、YOLOv4和YOLOv7 6種主流的一階目標檢測算法,與本文提出的YOLOv5-DAFC模型進行實驗對比分析,通過比較實驗結果,可以更加客觀地評估本文所提出改進模型的有效性和優越性。具體實驗結果如表2所示。

表2 與主流目標檢測算法的性能對比

通過以上對比可以明顯看出,本文所提出的YOLOv5-DAFC模型相比于幾種主流的one-stage目標檢測算法,在整體檢測精度、檢測速度、模型體積3個方面都有較大提升,并且在各方面均超過了更新的YOLOv7模型,基本可以實現高精度、高速度、輕量級需求的工業檢測。

5 結語

為實現精準、快速的材料表面缺陷檢測,本文提出了改進的YOLOv5優化算法,添加DyHead目標檢測頭,通過分離式注意力機制將多種注意力同時應用在同一個模塊中,有效提升目標檢測頭的檢測性能;更換ALRPLoss損失函數,簡化訓練參數、優化訓練效果的同時兼顧分類和定位精度;基于FasterNet網絡提出C3-Faster模塊,用于替換原版C3模塊,提升模型魯棒性同時減少權重體積;添加CARAFE上采樣算子,擴大感受野,提升內容感知性能。實驗證明,改進后的YOLOv5-DAFC模型相比原版YOLOv5模型,檢測精度提升的同時,模型參數量、計算復雜度、權重體積都有明顯優化,對于輕量化工業檢測領域有較高應用價值。