改進YOLOv5的鋼材表面缺陷檢測算法

李 鑫，汪 誠，李 彬，郭振平，李秋良，李卓越

(空軍工程大學基礎部，西安，710038)

鋼材作為一種重要的工業產品，是航空航天、機械制造和國防工業等領域不可或缺的原材料。隨著我國工業水平的不斷發展，鋼材生產技術取得了巨大進步[1]。而在鋼材的生產制造過程中，由于受到原料質量、制造設備和生產環境等因素影響，產品表面會形成裂紋、孔洞、劃痕等不同種類的缺陷[2]。這些缺陷會使鋼材的強度、抗疲勞和耐磨性能下降，影響正常使用，甚至可能造成嚴重后果。因此，為保證鋼材的安全性能，必須對其表面進行缺陷檢測。

傳統的鋼材表面缺陷檢測方法有人工檢測方法和頻閃光檢測方法[3]，均為非自動化的檢測方法。檢測人員需要進行大量重復的工作，容易產生視覺疲勞，造成漏檢或誤檢的情況。

近年來，機器視覺技術迅速發展，在圖像分類、人臉識別和目標檢測等領域得到了應用。機器視覺技術具有穩定、高效、準確等優勢[4]，有學者將該技術應用于缺陷檢測，提高了檢測效率和精度。韓強等利用檢測網絡對RPN網絡生成的區域建議框進行分類回歸，提出了一種特征融合和級聯檢測網絡的Faster R-CNN鋼材表面缺陷檢測算法[5]。葉欣提出了一種使用自適應空間特征融合結構和EIoU損失函數的改進YOLOv4算法,提高了對熱軋帶鋼表面缺陷的檢測精度[6]。楊莉等通過引入導向錨點候選區域和多任務FPN結構解決了傳統Faster R-CNN網絡在檢測鋼材表面小目標缺陷時特征融合不充分的問題[7]。

本文基于鋼材表面缺陷檢測任務，提出一種改進的YOLOv5算法，旨在提高算法對鋼材表面缺陷的識別準確率和檢測速度，并在NEU-DET數據集上驗證了算法的有效性。

1 YOLOv5算法介紹

YOLO系列算法[8]是一種基于深度學習和卷積神經網絡的目標檢測算法，具有推理速度快、檢測精度高、可實時檢測等優點。YOLOv5是該系列算法的最新一代，相比前幾代具有更高的精度、更快的速度和更小的體積。該算法共有4種模型：YOLOv5s、YOLOv5m、YOLOv5l和YOLOv5x，網絡的深度和寬度依次增加。本文選擇YOLOv5s作為預訓練模型，在其基礎上進行遷移學習，建立鋼材表面缺陷檢測算法。YOLOv5s的網絡結構如圖1所示。

圖1 YOLOv5s網絡結構圖

YOLOv5s模型主要由4部分組成：Input、Backbone、Neck和Head。Input部分采用了Mosaic數據增強，每次選取4張圖片進行隨機組合，豐富了數據集的多樣性[9]；Backbone部分用于提取圖像特征，主要由Focus、C3、SPP等模塊組成。Focus用來對輸入網絡的圖像進行切片操作,可以增加輸出深度，同時減少計算量；C3主要由若干個Bottleneck結構和卷積層組合而成，有利于加快推理速度；SPP[10]為空間金字塔池化層，分別使用尺寸為5、9、13的3種池化核對輸入圖像進行最大池化操作，有效提升了網絡的感受野；Neck是網絡的融合部分，采用了特征金字塔網絡(feature pyramid networks,FPN)[11]+金字塔注意力網絡(pyramid attention networks,PAN)[12]的結構，二者結合增強了對不同尺度特征的融合效果；Head為網絡的預測部分，輸出3組包含預測框類別、置信度和坐標位置的向量。

2 改進的YOLOv5算法

2.1 使用GhostBottleneck結構

GhostNet[13]是一種新穎的計算模塊，可以通過使用較少的參數來生成更多的特征，原理如圖2所示。使用該模塊可以減少參數量、消除冗余特征、獲得更加輕量化的模型。

圖2 Ghost模塊原理

GhostNet的具體過程為：

Y′=X*f′

(1)

(2)

GhostBottleneck模塊分為2種結構，如圖3所示。步長為1的GhostBottleneck主要由2個Ghost模塊堆疊而成；步長為2的GhostBottleneck的2個Ghost模塊之間通過一個步長為2的DW卷積層連接。

圖3 GhostBottleneck的兩種結構

YOLOv5s網絡中的Bottleneck結構對輸入特征圖先后使用32個1×1卷積核和64個3×3卷積核進行卷積操作,因此本文采用步長為1的GhostBottleneck替換網絡中的C3模塊，使用步長為2的GhostBottleneck替換網絡中的CBS模塊，以減少模型計算量。

2.2 添加注意力機制SE模塊

SENet(squeeze and excitation network)[14]是一種通道注意力機制，獲得了ImageNet 2017競賽分類任務的冠軍。該網絡的原理如圖4所示，通過學習的方式自動獲取每個特征通道的重要程度，并利用得到的重要程度來增強有用的特征，抑制不重要的特征，從而實現對特征權重的調整。

圖4 SE模塊

SE模塊主要包括squeeze、excitation和scale等操作：squeeze將特征圖的每個通道進行壓縮，對H×W的二維向量進行全局平均池化，輸出1×1×C的向量。

(3)

excitation在得到1×1×C向量后，利用全連接層，對每個通道生成一個權重，用來代表特征通道的重要程度。

s=σ(W2δ(W1z))

(4)

最后scale操作是將上一步輸出的權重看作每個特征通道的重要性，再加權到之前的特征上，完成在通道維度上對原始特征的重新標定。

(5)

2.3 增加檢測層

YOLOv5s網絡的head部分共輸出3個尺度的檢測層，輸入尺寸為640×640的圖像，輸出的檢測層尺度為80×80、40×40和20×20，分別用于檢測小、中、大目標。考慮到本文所用數據集中部分缺陷占整張圖片的比例較大，因此在原YOLOv5s網絡的head部分再增加一個尺度為10×10的檢測層，相應地在neck部分添加特征融合結構，以輸出新的特征圖。這樣在改進后的網絡中，head部分共輸出4個尺度的檢測層。

2.4 使用DW卷積

DWConv(depthwise convolution,深度卷積)[16]的原理是將卷積核拆分成單通道形式，在不改變輸入特征圖像深度的情況下，對每個通道進行獨立的卷積運算，得到與輸入特征圖通道數一致的輸出特征圖。相比標準卷積，DW卷積能夠降低運算量，節約計算成本。本文在引入注意力機制并增加檢測層后運算量增加較多，因此在neck部分使用DW卷積替代部分標準卷積，從而達到減少運算量的目的。

3 實驗結果與分析

3.1 實驗環境及數據集

本文的實驗環境基于Windows10操作系統，內存為64 GiB，使用PyTorch 1.9.1作為深度學習框架，Python版本為3.7.11，CUDA版本為10.2。CPU為Intel(R) Xeon(R) Gold 5218 CPU @2.30 GHz。GPU為NVIDIA GeForce RTX 2080Ti。

實驗所用數據集為東北大學宋克臣團隊制作的鋼材表面缺陷數據集(NEU-DET)，包含6種缺陷類型：裂紋(crazing)、內含物(inclusion)、斑塊(patches)、點蝕面(pitted surface)、氧化軋皮(rolled-in scale)、劃痕(scratches)，共計1 800張圖片。圖片的原始尺寸為200×200，類別為灰度圖。采集鋼材表面缺陷的一種方法為：將2個LED光源對稱傾斜安裝于待測鋼材上方，在兩光源的中心軸線上放置一臺工業相機，對鋼材表面的缺陷進行照片采集，如圖5所示。對相機采集到的圖像進行預處理，消除高頻噪聲，進行灰度變換[17]。

圖5 采集缺陷示意圖

數據集圖片均使用LabelImg軟件依次對缺陷標注真實框，得到格式為xml的標簽文件。將數據集的1 800張圖片按照6∶2∶2劃分為訓練集、驗證集和測試集，得到訓練集1 080張，驗證集360張,測試集360張。數據集部分圖片如圖6所示。

圖6 數據集部分圖片

3.2 訓練策略與評估指標

YOLOv5網絡的訓練策略為：使用K-means聚類算法自動根據數據集的標注框計算三組初始錨框的大小；初始學習率lr0為0.01，使用余弦退火算法動態調整學習率；學習率動量momentum為0.937，權重衰減系數weight decay為5×10-4；訓練過程中開啟Mosaic數據增強，提高模型的泛化能力；batch-size設置為16，訓練輪次為300個epoch。依次對原YOLOv5s模型和改進的模型進行訓練和測試，對比分析各模型的參數和性能。

本文從模型參數量、運算量、模型體積、平均推理時間、精確率P、召回率R、mAP[18]@0.5(IoU閾值為0.5時各個類別的平均AP值)等方面對模型進行比較。參數量表示訓練過程中模型的總參數量；運算量FLOPs表示浮點運算數；模型體積表示訓練完成后生成的模型權重大?。黄骄评頃r間是指GPU推理一張圖片所需的時間；精確率指檢測出的正樣本中實際仍為正樣本的概率，召回率指在實際的正樣本中檢測出正樣本的概率；平均精度AP指的是在不同召回率下精確率的均值，即P-R曲線與坐標軸圍成的面積，反映了準確率和召回率的綜合性能，是目前主流目標檢測模型的評價指標。對各個類別的AP(式中簡記為PA)求均值即可得到mAP(式中簡記為PmA)，即所有類別的平均精度。它們的具體計算公式為：

(6)

(7)

(8)

(9)

式中：P、R分別指精確率和召回率；N表示總的類別數；PT表示模型預測為正樣本且實際也為正樣本的數量；PF表示模型預測為正樣本但實際為負樣本的數量；NF表示模型預測為負樣本但實際為正樣本的數量。

3.3 改進的4種模型

本文共設計了4種模型與原YOLOv5s模型進行對比，分別為：

1) YOLOv5s-G：使用GhostBottleneck模塊替換Backbone中的CBS和C3模塊。

2) YOLOv5s-S：在Backbone中增加3個注意力機制SE模塊。

3) YOLOv5s-D：在head部分增加1個尺度為10×10的檢測層，同時在neck部分增加特征融合結構。

4) YOLOv5s-GSD：在原YOLOv5s網絡中同時使用GhostBottleneck替換CBS和C3，添加3個注意力機制SE模塊，在網絡中增加1個10×10的檢測層和對應的特征融合結構，使用DW卷積替換neck中的部分標準卷積。該模型即為本文所提出的改進YOLOv5s模型，其網絡結構如圖7所示，其中GB-1和GB-2分別表示步長為1和步長為2的GhostBottleneck模塊。

圖7 改進的YOLOv5s-GSD模型結構

3.4 實驗對比

3.4.1 模型參數與性能對比

為了驗證單個結構改變對網絡性能的影響，本文進行了消融實驗，對YOLOv5s、 YOLOv5s-G、YOLOv5s-S、YOLOv5s-D、YOLOv5s-GSD進行訓練和測試，得到5種模型的復雜度和性能對比如表1和表2所示。

表1 各模型復雜度對比

表2 各模型性能對比

分析實驗結果可知，YOLOv5s-G模型相比原YOLOv5s的參數量減少了41.1%，運算量減少了40.9%，模型大小減少了39.9%，平均推理時間減少了0.4 ms。這得益于GhostNet使用了線性運算替代部分卷積過程，從而在得到近似于原始特征圖的同時減少計算量和特征圖冗余[19]。YOLOv5s-S相比于原YOLOv5s模型，由于SE注意力機制的引入，雖然增加了模型復雜度，使得模型的參數量、運算量和模型體積略有增大，導致平均推理時間有0.2 ms的增加，但mAP值提升了1.6%，這證明了注意力機制對重要特征的強化作用。YOLOv5s-D由于增加一個檢測層和部分卷積結構，相比原YOLOv5s網絡參數量增加較多，平均推理時間增加0.3 ms，但mAP值提升了0.6%。本文提出的YOLOv5s-GSD模型，相比原YOLOv5s模型，參數量減少了11.5%，運算量減少了44.5%，模型體積減少了10.4%，mAP值提升了3.3%，平均推理時間減少了0.1 ms。

YOLOv5s-GSD和YOLOv5s檢測各類缺陷的平均精度AP值如圖8所示。從圖中可以看出，改進后的YOLOv5s-GSD模型相比原YOLOv5s模型，除scratches類缺陷的AP值降低2.2%外，其他類別均有不同程度的提高，其中crazing、inclusion、patches、pitted surface和rolled-in scale的平均精度分別提高5.9%、2.0%、4.6%、2.6%和7.2%，對于原模型檢測效果較差的crazing和rolled-in scale這兩類缺陷的檢測精度有了明顯提升。

圖8 YOLOv5s-GSD和YOLOv5s各類別平均精度AP值對比

綜合以上分析，可以得出：Ghost模塊和DW卷積對于模型輕量化和減少運算量有顯著作用；SE注意力機制對于精確率有明顯的提升，很好地改善了誤檢的情況；在head部分新增的檢測層對于較大目標的檢測精度有所提高。本文提出的YOLOv5s-GSD模型相比原YOLOv5s模型在參數量、運算量、模型體積、mAP值以及推理速度等方面均有明顯優勢。

3.4.2 訓練過程損失

YOLOv5網絡的損失函數共包含3個部分：Lbox為定位損失函數，用于衡量預測框對坐標定位錯誤的誤差；Lobj為置信度損失函數，反映了預測框的置信度誤差；Lcls為分類損失函數，反映了預測框對目標類別預測錯誤引起的誤差。其中Lbox采用CIoU_Loss[20]函數，Lobj和Lcls采用交叉熵損失函數。本文提出的YOLOv5s-GSD模型在訓練過程中的損失變化如圖9所示。

圖9 YOLOv5s-GSD模型訓練過程損失

訓練設置迭代次數為300次，從圖中可以看出，在0到50次迭代的過程中，Lbox和Lcls急劇下降，之后的訓練過程中緩慢下降，而Lobj在訓練初期出現了局部的過擬合現象；經過300次迭代后，損失值逐漸趨于平穩。

3.5 檢測效果對比

為更加直觀地比較改進后的YOLOv5s-GSD模型與原YOLOv5s模型的實際檢測效果，本文分別用2種模型對測試集的缺陷圖片進行了檢測，部分檢測結果對比如圖10所示。

圖10 改進前和改進后模型的檢測效果對比

圖10(a)為原YOLOv5s模型的檢測效果，(b)為改進的YOLOv5s-GSD模型檢測效果。對比兩者可以看出，原YOLOv5s模型對于一些目標存在漏檢的情況，且置信度普遍較低。改進的YOLOv5s-GSD模型檢測缺陷目標的置信度相比原YOLOv5s模型有了明顯提高，尤其是crazing類缺陷，在使用原YOLOv5s模型檢測時效果較差，且存在較多漏檢目標，改進后的模型明顯提升了檢測效果，而且對于缺陷的定位也更加準確。

3.6 不同算法對比

為了進一步評估YOLOv5s-GSD模型的性能，本文選取了3種主流的目標檢測算法SSD[21]、YOLOv3和YOLOv4，在同樣的數據集上與本文改進的YOLOv5模型進行實驗對比，結果如圖11和表3所示。

圖11 不同算法檢測的各類別平均精度AP值對比

表3 不同算法檢測性能對比

根據圖11和表3的實驗結果可得：本文所改進的YOLOv5s-GSD模型相比幾種主流目標檢測算法，有效提升了對各種類別缺陷的檢測精度，尤其是主流算法檢測效果較差的crazing和rolled-in scale類缺陷，AP值有大幅提高，mAP值達76.8%；參數量和模型體積大幅減小，更加便于移動端的部署；同時具有較高的檢測速度，達到了43.3 FPS，實現了對缺陷目標的實時檢測。因此可以得出，本文改進的YOLOv5模型在模型體積、檢測精度和檢測速度等方面均優于主流目標檢測算法，能更好地完成鋼材表面缺陷檢測的任務。

3.7 實驗結果總結與展望

通過上述的實驗結果可以得出，本文改進的YOLOv5s-GSD模型對比原YOLOv5s模型和幾種主流的目標檢測算法，在模型復雜度和檢測性能等方面都具有明顯優勢。相比原YOLOv5s模型，參數量減少了11.5%，運算量減少了44.5%，模型體積減少了10.4%，實現了模型的輕量化；同時mAP值提升了3.3%，幾乎對各類缺陷的檢測精度均有不同程度的提升。其中crazing類缺陷的檢測精度相對較低，有較大提升空間，本文推測可能是該類別缺陷的標注存在一定的誤差。可使用標注軟件對圖片的標注信息進行調整，提高圖片標注精度，從而進一步提升mAP值。

4 結論

本文提出了一種基于改進YOLOv5算法的鋼材表面缺陷檢測模型，利用GhostBottleneck結構和DW卷積實現了網絡模型的輕量化；引入SE注意力機制，通過強化重要的特征來增強檢測效果；改進了特征融合層結構，增加了一個尺度為10×10的檢測層，提升了對部分類別缺陷的檢測精度。下一步工作將驗證模型在不同數據集上的訓練和檢測效果，同時進一步優化網絡結構，提高網絡的檢測速度和精度。研究如何將模型部署在移動端，在實際的缺陷檢測應用中對模型加以完善和改進 。