基于動態Transformer的輕量化目標檢測算法

方思凱, 孫廣玲, 陸小鋒, 劉學鋒

(上海大學,上海 200000)

0 引言

目標檢測[1]旨在從一幅圖像中同時分類和定位出多個感興趣的物體,近年來因其快速的技術突破和廣闊的應用場景而受到關注。傳統的目標檢測方法[2-3]通常使用卷積神經網絡(CNN)[4]作為骨干來提取視覺特征,其通過堆疊卷積層逐漸增大模型的感受野以檢測整幅圖像。同時,具有不同感受野的特征圖可以實現對不同尺度物體的檢測。然而,LUO等[5]研究發現CNN中特征圖的實際感受野只占其理論值的一部分,這表明模型的最終感知區域難以覆蓋整幅圖像,也嚴格限制了其檢測性能。

隨著深度學習的發展,誕生于自然語言處理領域[6]中的Transformer網絡框架[7]為人們帶來了新的思路。如文獻[8]提出了結合Transformer和卷積神經網絡的檢測模型;文獻[9]引入了位置編碼和編碼器-解碼器框架范式;更進一步地,文獻[10]徹底移除了卷積結構,提出了一種稱為Vision Transformer的純Transformer網絡框架,該框架將圖像分割為均勻的小塊,使用自注意力機制在各塊之間建立相關性,以充分利用圖像的視覺特征。然而,Transformer視覺模型在目標檢測任務中仍然存在問題。首先,全局自注意力機制擅長遠程依賴性建模,適用于分類任務(其實例往往占據整幅圖像)以及目標檢測中的大尺度物體檢測。但對于小尺度和中尺度物體的檢測精度往往比CNN模型低幾個百分點,這表明完全排除局部機制是不合理的。此外,全局自注意力機制的計算復雜度隨著圖像塊的數量增加呈二次指數增長,尤其在目標檢測等高分辨率輸入的任務中,對設備的計算能力要求較高,難以在實際應用中大規模拓展。

針對Transformer在目標檢測任務中檢出效率低、多尺度泛化能力弱的問題,本文在Vision Transformer的框架基礎上,提出了一種輕量化的動態Transformer目標檢測算法。改進算法在降低計算復雜度的同時,增強了對中小尺度目標的檢測性能,保證了模型的實時性和準確性。

1 改進的Transformer目標檢測算法

改進的算法即動態Transformer目標檢測算法的網絡整體框架如圖1所示。

圖1 改進后網絡的整體框架結構Fig.1 Overall framework structure of the improved network

首先通過4×4的卷積層將輸入圖像映射成空間尺度為原圖的1/4、通道數為96的特征圖,然后依次經過4個網絡階段進行特征提取,每個階段由3個編碼層組成,同時,在網絡中插入由2×2的最大池化層和1×1的卷積層組成的下采樣模塊,將特征圖的空間尺度和通道數分別減小至原來的1/2和擴展2倍,最后使用Faster R-CNN[11]目標檢測頭得到最終的檢測結果。編碼層主要由改進的動態稀疏自注意力(DS-Attention)模塊和全連接層組成,對于前者,本文算法使用由卷積組成的動態門從全局范圍內為每個查詢向量Q選擇一定比例(由動態因子μ決定)的鍵向量K和值向量V做自注意力計算。同時,在網絡的淺層階段使用較小的比例閾值,關注圖像的局部區域,隨著網絡的階段加深逐漸增大該比例閾值,以最終實現全局關注的自注意力模式。為了進一步加快模型的計算效率,本文算法還將動態門機制應用于網絡的結構層面,通過在每個編碼層前面插入動態跳層門(DLG),模型可根據不同的輸入自適應地選擇不同的編碼層參與計算,以跳過重要性較弱的編碼層,從而減少模型的冗余。此外,層歸一化(LayerNorm)和批歸一化(BatchNorm)被插入網絡以優化訓練梯度。

1.1 動態稀疏自注意力機制

傳統自注意力機制的關鍵組件有查詢向量Q、鍵向量K和值向量V,由輸入特征x線性變換而來。隨后進行線性映射轉換為多頭狀態,即

(1)

在所提出的動態稀疏自注意力機制中,如圖1右下角所示,對于第l個編碼層,通過動態門Ql生成一組概率向量Gl,并假設生成的Gl中元素的值越大,選擇對應的K和V的概率就越高,其算式為

(2)

式中:C(·)為4×4的卷積;B和δ分別為批歸一化和GeLU激活函數;同時,將當前和前一編碼層生成的概率向量相加,以記憶先前層的選擇結果,并定義了一個記憶因子α,通常被設為0.2。最后,生成的概率向量Gl∈Rh×N轉換為掩碼,即

(3)

式中:j的區間為(1,N),N為圖像塊的數量;μ為動態因子;閾值t為Gl中第μ大的值。Gl中小于閾值t的元素被設置為0,其余的被保留,相當于過濾器。然后,將掩碼MMask擴展到與K∈Rh×N×dk相同的維度,并使用哈達瑪積將MMask與K和V進行逐元素乘積,以對鍵向量和值向量進行掩碼,算式為

(4)

最終的自注意力函數為

(5)

式中,D(Kmasked)和D(Vmasked)都屬于Rh×μ×dk,由最大池化函數D(·)生成,以顯式地減少掩碼后鍵向量和值向量的特征維度。

由式(5)可知,提出的動態稀疏自注意力機制的計算復雜度降低到與N呈線性關系,且可以在全局水平上動態選擇重要性較高的鍵向量和值向量參與自注意力計算。同時,由于自注意力的多頭特性,可以在每個頭中為查詢向量生成不同的概率向量,使得模型即使在μ很小的情況下也有能力關注到圖像的各個區域,而殘差連接能夠將當前層未被選擇的信息向后傳遞。由于其輸出維度與原始自注意力機制相同,且以上操作都是可微的,所以可以直接代替原始自注意力機制進行端到端訓練。

1.2 從局部到全局的自注意力模式

本文進一步將動態稀疏自注意力機制中的動態因子歸一化表示為自注意力選擇的百分比,并為每個網絡階段分配了一個固定的μ值。因此設計了幾種特定的自注意力模式:局部到全局型、全局到局部型、中等型和局部型。隨著網絡階段的加深,其μ值分別為(0.2,0.4,0.6,0.8),(0.8,0.6,0.4,0.2),(0.5,0.5,0.5,0.5)和(0.2,0.2,0.2,0.2)。通過Grad-Cam[12]將上述4種模式以及傳統(全局)模式對小、中和大尺度物體的關注效果進行可視化展示,如圖2所示。

由圖2可以觀察到,局部型和中等型的自注意力模式能夠識別出小尺度和中尺度物體,但其受限的感受野不足以識別大尺度物體。傳統(全局)自注意力模式雖然可以識別出小尺度物體的語義信息,但很難在全局范圍里準確定位出它們的位置,導致小目標的檢測性能較差。

對于局部到全局型,當網絡初始階段的動態因子較小時,這種較小的感受野有利于模型在早期階段識別和定位出小目標。隨著網絡階段的加深,增大的動態因子允許模型關注更大的特征區域,特征也逐漸轉化為高層次的語義信息,使模型能夠識別出中型和大型目標。此外,對于全局到局部型的注意力模式,較淺的網絡階段缺乏高級語義信息而較深的階段缺乏足夠的感受野,使其始終無法識別出大目標。綜上所述,與自然語言處理領域相比,在圖像處理領域,尤其是目標檢測任務,在網絡的早期階段不需要對圖像中的太大區域進行關注,否則會導致模型丟棄低層次的細粒度特征,同時還會增加計算復雜度。類似于卷積神經網絡的逐級提取特征的方法更有利于模型取得更好的性能。因此,本文最終選擇了從局部到全局的自注意力模式。

1.3 動態跳層網絡結構

根據常識,對于一幅簡單圖像,人眼只需較短時間就能分辨清楚;而對于一幅復雜圖像,通常需要較長時間且仔細的判斷才能辨認。因此,本文從人眼的視覺機制出發,嘗試將動態決策方法進一步擴展到模型的結構層面上(如圖1左下角所示)。

首先,對于第l個編碼層,在其前面插入動態條件門來生成重要性權重因子Hl∈R1,其算式為

Hl=max(0,tanh(C1(δ(B(C2(Xl))))))

(6)

式中:Xl為當前編碼層的輸入特征;C1(·)和C2(·)為3×3的卷積;tanh(·)為雙曲正切函數。因此,Hl的結果區間為[0,1],表示當前層的重要性。

在網絡訓練時,不論Hl是否為0,每個編碼層的所有結構都參與計算,防止網絡中某些參數沒有被訓練。在推理時,若Hl為0,則直接跳過當前層。訓練和推理時編碼層的輸入輸出分別有以下關系

(7)

(8)

(9)

重要性權重Hl的值越小,網絡在當前層的期望計算量就越少,反之亦然。則新的損失函數為

(10)

式中:Laccuracy為精度項損失函數;Lefficiency為效率項損失函數;λ為精度和效率之間的權衡因子,λ越大,網絡越輕量,反之網絡精度越高;θ為衰減系數,控制網絡提升效率的上限,通常設為0.2。

最后,為了使網絡更好地訓練,本文算法在訓練的早期階段只采用精度項損失函數,稱為預加熱階段,之后才加入效率項損失函數。通過動態跳層機制,網絡在訓練時能夠利用損失函數,學習在不同的輸入情況下動態地選擇編碼層,以更好地平衡網絡精度和效率。

2 實驗結果與分析

2.1 實驗數據集與評價指標

本文使用MS COCO-2017數據集[13]作為訓練和測試時使用的數據集,共包括80個目標類別,其中訓練集和測試集分別擁有11萬張圖片和5000張圖片。采用浮點運算次數(FLOPS)和檢測幀率(FPS)分別驗證模型的計算量和計算速度。同時,使用平均精度均值(mAP)、小尺度目標精度(APS)、中尺度目標精度(APM)和大尺度目標精度(APL)等作為驗證模型檢測性能的主要指標。其中

(11)

(12)

(13)

式中:P為查準率;R為召回率;PT為真正例;PF為假正例;NF為假反例;PAP為平均精度;P(R)為P-R曲線;C為目標類別數。APS,APM和APL分別為像素面積小于322、像素面積在322與962之間以及像素面積大于962的目標框的AP值。AP50和AP75分別表示交并比為0.5和0.75時的AP值。

2.2 實驗細節與運行環境

本文在訓練時采用的操作系統為Windows 10,GPU為NVIDIA-RTX3090(24 GiB)。為了測試模型的實際檢測能力,使用“大疆妙算”(機載計算機)作為測試設備。機載計算機的操作系統為Ubuntu-16.04,GPU為NVIDIA-Jetson TX2(8 GiB)。其中,模型的初始學習率為8E-5,AdamW優化器和分段式學習率衰減策略被用于訓練模型。網絡訓練的前20個輪次被設定為預加熱階段,此時只有精度項損失函數生效,在預加熱階段結束后效率項損失函數才開始生效。對于輸入的圖像,其短邊被調整為512像素,而長邊保持在800像素以下。模型的具體參數以及當輸入為H×W×3時模型各部分的輸出特征大小如表1所示,其中,k和s分別代表核大小和步長。

表1 改進模型的具體參數Table 1 Detailed parameters of the improved model

2.3 對比與消融實驗

圖3展示了本文算法在MS COCO-2017數據集上的部分檢測結果,模型展現了對各種尺度目標以及復雜圖像的檢測能力。同時,圖4還展示了在不同輸入條件下模型的編碼結果。由圖4可知,模型能夠根據輸入動態選擇部分編碼層參與計算以達到檢測效果,體現了算法的自適應能力。

圖3 本文算法在MS COCO-2017數據集上的部分檢測效果實例Fig.3 Some examples from MS COCO-2017 test dataset detected by the proposed algorithm

圖4 部分編碼結果展示Fig.4 Display of partial encoding results

為了對模型進行更詳細的性能測試,本文對比了改進算法模型與其他幾種基于Transformer的目標檢測基準模型的性能,如表2所示。

表2 模型性能對比Table 2 Comparison of model performance

與其他算法相比,本文算法在檢測精度與效率方面都取得了更好的性能。其中,模型的平均精度均值達到了43.8%,尤其是對小、中尺度目標的檢測性能分別為27.6%和47.4%,相比其他模型有較大的提升。此外,模型的計算量相對較小,其檢測速度達到了37幀/s,實時性大幅提升。本文在圖5中進一步說明了改進算法模型的收斂性,由圖5可知,模型在300個訓練輪次就可取得較為理想的結果。

圖5 模型訓練損失變化Fig.5 Variation of model training loss

為了研究本文算法中的不同組件對模型的影響,將所提出的改進模塊逐一添加到原Vision Transformer框架中,其結果如表3所示。動態稀疏自注意力在提升小尺度和中尺度目標檢測性能的同時減輕了網絡的一部分計算量,動態跳層機制在略微損失精度的同時大幅降低了模型的計算復雜度,對于提升效率的貢獻度較大,通過它們的組合實現了更高效的Transformer檢測模型。

表3 不同組件對模型的影響Table 3 Impact of different modules on the model

為了驗證不同的自注意力模式對模型性能的影響,本文將傳統(全局)模式與本文提出的4種模式的性能進行了測試對比,如表4所示。相比傳統的全局型模式,所提出的幾種模式在小型和中型物體上都取得了明顯的性能提升,得益于受限的關注區域。本文建議的局部到全局型的自注意力模式能夠關注所有尺度的物體,取得了最好的性能。該實驗結果與圖2中觀察到的現象一致。

表4 不同自注意力模式的性能對比Table 4 Comparison of performance of different self-attention patterns

為了驗證動態跳層機制在模型訓練時的工作模式,圖6顯示了在前50個訓練輪次中,損失函數中的不同權衡因子對網絡訓練的影響。在前20個訓練預加熱輪次中,算法為了提高精度增加了模型的計算量。在20個輪次后,效率項損失函數開始生效,模型的計算量逐漸下降,同時,模型精度在略微衰減后繼續上升?？梢钥闯?權衡因子越大,模型越輕量,對精度的影響也越大。

圖6 權衡因子對訓練的影響Fig.6 Impact of the trade-off factor on training

3 結束語

為了提高Transformer檢測框架的實時性和準確性,本文提出了基于動態Transformer的輕量化目標檢測算法。實驗表明,本文算法改善了原Transformer框架對于小尺度和中尺度物體檢測性能不足的問題,同時模型的計算量有效減少,實時性得到有效提升,在機載計算機上的檢測速度達了37幀/s,模型的實際應用空間獲得大幅增加。