用于PCB缺陷檢測系統(tǒng)的YOLOv5算法設計

張振亞 趙成林

（北京郵電大學信息與通信工程學院 北京市 100876）

缺陷目標檢測是工業(yè)領域很有價值的研究方向。其中印刷電路板（Printed Circuit Board，PCB）的缺陷檢測對工業(yè)生產十分重要。傳統(tǒng)的PCB 缺陷檢測方法存在檢測速度慢、漏檢率高、魯棒性不佳的種種問題?，F(xiàn)代基于CNN 的深度學習目標檢測算法對于大目標有很強的檢測能力，但對于PCB 中這種小目標缺陷，檢測能較弱。小目標不易于檢測的原因是，缺陷目標像素過少，在不斷卷積的過程中丟失了大部分的有效特征和細節(jié)特征。因此，無論是對于傳統(tǒng)檢測技術還是現(xiàn)代檢測技術，小目標檢測都是叩待解決的問題。

1 算法設計思路與貢獻介紹

為了解決上述問題，對YOLOv5[1]框架進行了研究，提出了ATCSP-YOLOv5 用于PCB 的缺陷檢測。實驗所使用的PCB 數(shù)據(jù)集，具有分辨率高，但缺陷目標與圖像的相對比例非常小的特點。其中最小的一類缺陷占整個圖像的比例僅為0.00893%。結合數(shù)據(jù)集的特點，ATCSP-YOLOv5 算法采用注意力機制來解決問題。注意力機制在不增加復雜度的前提下，通過給模型分配合適的權重，對小目標的特征信息進行增強。不僅為模型獲取了特征圖的全局信息，又使目標上下文語義信息得到了豐富。由于圖像分辨率較高，所以ATCSP-YOLOv5 選擇使用空間注意力機制Self-Attention[2]來進行學習。同時，受到ViT（Vision Transformer）[3]的啟發(fā)。ATCSP-YOLOv5 借鑒ViT 的思想將Transformer[2]模塊嵌入到主干網(wǎng)絡中，對主干網(wǎng)絡特定深度的特征圖做多頭自注意力機制的處理。實驗證明，此方法有效改善了小目標的檢測效果。同時，ATCSP-YOLOv5 還受到CSPNet（跨階段局部網(wǎng)絡）[4]的啟發(fā)，將SCPNet 的思想與Transformer 模塊結合，通過將梯度的變化從頭到尾地集成到特征圖中，增強了CNN 的學習能力、消除了計算瓶頸、減少了內存成本。其中Transformer 中多頭自注意力機制（Multi-Head Attention）和SCPNet 結合的模塊為ATCSP（多頭自注意力跨階段局部網(wǎng)絡）。

研究所做出的貢獻如下：

（1）將多頭自注意力機制和SCPNet 的思想結合，生成ATCSP 模塊。提高了模型的小目標檢測性能。

（2）在特征金字塔部分融合生成感受野更小的檢測層，添加了一個小目標探測頭；同時去掉了原本感受野最大的探測頭，減少了計算量。

（3）通過使用特定的數(shù)據(jù)增強方法和SRGAN（SRGAN-VGG54）[5]來預處理和擴充數(shù)據(jù)集，為PCB 數(shù)據(jù)集的擴充和預處理提供了有效策略。

（4）將訓練好的ATCSP-YOLOv5 模型中的PANet[6]換成FPN[7]，對模型進行微調訓練，降低了模型的計算量和體積，提升了模型精度。最終將算法封裝成檢測系統(tǒng)如圖1所示。

2 ATCSP-YOLOv5模型

2.1 YOLOv5簡介

YOLOv5 有五種不同的模型?？紤]到PCB 在流水線檢測上實時性和精度的要求，選定YOLOv5s 作為基礎框架來進行改進。YOLOv5s 的結構如圖2所示。

圖2：YOLOv5s 結構圖

2.2 ATCSP-YOLOv5模型結構

相比于YOLOv5s，ATCSP-YOLOv5 將主干網(wǎng)絡中80×80 特征圖大小的C3 結構替換為ATCSP 模塊。同時在特征金字塔部分融合生成特征圖大小為320×320 的小目標檢測層，添加了新的YOLOHead3，刪除了YOLOv5s 中的YOLOHead1。ATCSP-YOLOv5 的結構如圖3所示。

圖3：ATCSP-YOLOv5 結構圖

將訓練好的ATCSP-YOLOv5 模型中的PANet 替換為FPN 再進行模型參數(shù)微調得到ATCSP-YOLOv5_Pruned。

2.3 ATCSP模塊

Base Layer 是上一個卷積層的輸出。Part1、Part2 是Base Layer 經過特定卷積后分成的兩個部分。Part1、Part2一部分不處理，一部分經過原來的模塊，然后在最后將兩部分進行Concat。ATCSP 的結構如圖4所示。

圖4：ATCSP 結構圖

2.4 Transformer Layer

Transformer Layer 的輸入是ATCSP 中的Part2 特征圖。將輸入的特征圖分割成多個16×16 像素的小塊。然后將各個小塊通過卷積展平成特定維度的一維矢量序列，再給每個矢量加上位置矢量后，便可以送入Transformer Encoder 進行注意力機制的學習了。Transformer Encoder 輸出的結果是一個一維矢量需要通過卷積和上采樣將其恢復成與輸入尺寸相同的特征圖。Transformer Layer 的內部結構示意圖如圖5所示。其中，Transformer Encoder 主要包括兩個模塊，分別為多頭自注意力模塊和全連接層（MLP）模塊，兩模塊均使用殘差連接。模塊中有Norm 層負責歸一化處理，有助于網(wǎng)絡收斂，防止過擬合。將多頭自注意力模塊和MLP 作為一個子模塊，進行S 次堆疊最終形成Transformer Encoder。

圖5：Transformer Layer 示意圖

2.5 Multi-Head Attention（多頭自注意力機制）

多頭自注意力模塊是ATCSP 的核心部分。他有效增強了特征圖中小目標周圍的上下文信息，加強了模型的全局信息關注能力。多頭自注意力機制可解釋性示意圖如圖6所示。

圖6：多頭自注意力可解釋性示意圖

Multi-Head Attention（多頭自注意力）公式如下：

（1）（2）為自注意力機制公式，（3）（4）為多頭自注意力機制公式。其中，是特征圖生成的輸入矩陣，C是特征圖的通道數(shù)，Din是輸入矢量的維度；TQ,TK,TV是可訓練的矩陣（本質上是全連接層的權重矩陣）對Xin做線性變換生成V,Q,K作為AT的輸入。ATHeadi是多頭自注意力編碼器中的第i個注意力頭部；Dmodel是Transformer Encoder 輸入矢量的維度；N是Encoder中的ATHead的數(shù)量；

3 實驗

3.1 數(shù)據(jù)集和預處理

訓練使用的數(shù)據(jù)集是來自北京大學智能機器人開放實驗室所提供的公共數(shù)據(jù)集。數(shù)據(jù)集中包含693 張圖像，包括六種缺陷分別為，漏焊、鼠咬、斷路、短路、毛刺、余銅。由于數(shù)據(jù)集數(shù)量相對較少，為了增加訓練出模型的魯棒性，實驗對數(shù)據(jù)集進行了預處理。首先使用SRGAN（超分辨率對抗生成網(wǎng)絡）對數(shù)據(jù)集進行超分辨率數(shù)據(jù)增強，生成細節(jié)紋理更加清晰的數(shù)據(jù)集，增加了小目標的細節(jié)特征信息。然后使用復制增強[8]的方法，在圖像中復制多份小目標，增加了小目標在圖像中出現(xiàn)的次數(shù)。最后在此基礎上，再使用圖像旋轉、鏡像對稱、增加缺陷和背景的對比度、隨機亮度調節(jié)、HSV 調節(jié)幾種方式,將數(shù)據(jù)集擴充到了8316 張。

3.2 實驗細節(jié)

在實驗方面，訓練使用的是GPU 型號是NVIDIARTX3090。實驗環(huán)境使用的是官方給出的YOLOv5 Pytorch版環(huán)境。在訓練細節(jié)上，按訓練集和測試集8:2 的比例進行訓練。為了防止過擬合，實驗采用提前停止訓練的策略。經過實驗，發(fā)現(xiàn)改進后的模型收斂速度很快，模型在56 輪epoch 左右趨近于收斂。訓練結束后，將ATCSP-YOLOv5中的PANet 替換為FPN 再進行微調訓練。微調訓練時，模型在第22 輪epoch 左右收斂，得到了模型體積更小，檢測精度更高的ATCSP-YOLOv5_Pruned 模型。

3.3 算法比較

實驗選擇了兩種最為經典的算法與ATCSP-YOLOv5 進行比較，分別是兩階段算法中的Faster R-CNN 和一階段算法中的SSD。訓練曲線對比和檢測精度、檢測速度、參數(shù)量的對比結果，如圖7 和表1所示。

表1：各算法mAP@0.5 和檢測速度的比較

很明顯可以看出ATCSP-YOLOv5 收斂速度明顯快于原始YOLOv5，同時mAP@0.5 增加了3.06%。進行FPN 替換后，ATCSP-YOLOv5_Pruned 相較于原始YOLOv5，mAP@0.5增加了3.24%。由圖表易見，ATCSP-YOLOv5 的性能遠高于Faster R-CNN 和SSDLite_MobileNetV2。

3.4 檢測效果展示

如圖8所示。

圖8：檢測效果展示圖

4 結束語

ATCSP-YOLOv5 是用于檢測與PCB 數(shù)據(jù)集特點類似的小目標檢測算法。首先，該算法借鑒了一些前沿技術的思想，例如Self-Attention、ViT 以及CSPNet 的思想，并將他們巧妙融合起來，作為算法的核心。然后，又提供了許多用于小目標數(shù)據(jù)集預處理的有效策略。在此基礎上，還通過網(wǎng)絡替換微調的嘗試，得到了一個檢測PCB 數(shù)據(jù)集精度更高的模型，并將此模型做成了一款檢測系統(tǒng)。最后，實驗表明ATCSP-YOLOv5 算法明顯優(yōu)于已有算法，并保證了檢測的實時性。