文章目錄
- 1 概述
- 2 模型說明
- 2.1 總體架構
- 2.2 輕量pixel decoder
- 2.3 實例激活引導的Query
- 2.4 雙路徑更新策略
- 2.5 GT掩碼引導學習
- 2.6 損失函數
- 3 效果
1 概述
FastInst是一種基于query的實時實例分割方法,它能以32.5FPS的實時速度在COCO測試集上達到40.5的AP。在實例分割領域,基于query的方法源自Mask2former,FastInst也不例外,它也沿用了MaskFormer的大架構。其核心設計包括:
(1)實例激活引導的Query
(2)雙路徑更新策略
(3)GT掩碼引導學習
這些設計使FastInst能夠使用更輕量的像素解碼器和更少的 Transformer 解碼器層,同時獲得更優的性能。
2 模型說明
2.1 總體架構
如圖2-1-1所示的左半部分所示,FastInst和MaskFormer一樣,網絡結構上主要分為三大部分,分別是Backbone、Pixel decoder和Transformer decoder。
假設網絡的輸入為I∈RH×W×3I \in R^{H \times W \times 3}I∈RH×W×3,經過Backbone之后,可以得到下采樣8倍、16倍和32倍的特征圖C3C_3C3?、C4C_4C4?和C5C_5C5?,此時這三個特征的shape可以表示為[BatchSize,2048,H8,W8][BatchSize, 2048, \frac{H}{8}, \frac{W}{8}][BatchSize,2048,8H?,8W?]、[BatchSize,1024,H16,W16][BatchSize, 1024, \frac{H}{16}, \frac{W}{16}][BatchSize,1024,16H?,16W?]和[BatchSize,512,H32,W32][BatchSize, 512, \frac{H}{32}, \frac{W}{32}][BatchSize,512,32H?,32W?]。
這三個特征就是pixel decoder的輸入,輸出對應的三個特征圖E3E_3E3?、E4E_4E4?和E5E_5E5?。C3C_3C3?、C4C_4C4?和C5C_5C5?經過1×11 \times 11×1的卷積,統一將通道數變為256,此時的shape仍舊是[BatchSize,256,H8,W8][BatchSize, 256, \frac{H}{8}, \frac{W}{8}][BatchSize,256,8H?,8W?]、[BatchSize,256,H16,W16][BatchSize, 256, \frac{H}{16}, \frac{W}{16}][BatchSize,256,16H?,16W?]和[BatchSize,256,H32,W32][BatchSize, 256, \frac{H}{32}, \frac{W}{32}][BatchSize,256,32H?,32W?]。然后,從E4E_4E4?特征圖中選取NaN_aNa?個實例激活引導查詢,并于NbN_bNb?個輔助可學習查詢連接,作為最終的總查詢Q∈RN×256Q \in R^{N \times 256}Q∈RN×256,其中N=Na+NbN=N_a+N_bN=Na?+Nb?。
Transformer decoder將QQQ和高分辨率像素特征圖E3E_3E3?作為輸入,以雙路徑的方式更新像素特征QQQ和E3E_3E3?,并在每個解碼器層預測對象類別和分割掩碼。
2.2 輕量pixel decoder
多尺度上下文特征圖對于圖像分割至關重要。然而,使用復雜的多尺度特征金字塔(FPN)會增加計算負擔。與直接采用pixel decoder輸出的底層特征圖的先前方法不同,FastInst使用 Transformer decoder中經過精煉的像素特征來生成分割掩碼。這種設置降低了pixel decoder對大量上下文聚合的需求,因此使用輕量級的pixel decoder模塊就夠了。換句話說,也就是主力在Transformer decoder了,pixel decoder搞得簡單點就行了。
為了在精度和速度之間取得更好的平衡,FastInst使用一種名為 PPM-FPN 的變體來替代普通的 FPN。
pixel decoder的示意圖如下圖2-2-1所示(假設整個網絡的的輸入為I∈R1024×1024×3I \in R^{1024 \times 1024 \times 3}I∈R1024×1024×3),PPM只作用于C5C_5C5?的特征,這里的CBR指的就是[Conv, BN, Relu]。
2.3 實例激活引導的Query
基于Query的這種方式最初來自于DETR,但是DETR的Query是全零初始化的,所以其收斂速度很慢。收到Deformable DETR的啟發,作者提出了一種實例激活引導的Query,直接從多尺度的特征圖上選取包含高級語義的Queries。具體來說,給定pixel decoder的輸出特征圖,在特征圖E4E_{4}E4?上添加一個輔助分類頭,隨后進行softmax激活,以生成每個像素的類別概率預測pi∈ΔK+1p_{i} \in \Delta^{K+1}pi?∈ΔK+1,其中ΔK+1\Delta^{K+1}ΔK+1是(K+1)(K+1)(K+1)維概率單純形,KKK是類別數,為“無目標”(?)增加了一個類別,i是像素索引,輔助分類頭分別由兩個內核大小為3×33×33×3和1×11×11×1的卷積層組成。通過pip_{i}pi?,獲得每個像素的前景概率pi,kip_{i, k_{i}}pi,ki??、ki=argmaxk{pi,k∣pi,k∈pi,k∈1,?,K}k_{i}=argmax_{k}\{{p_{i, k} | p_{i, k} \in p_{i}, k \in}{1, \cdots, K}\}ki?=argmaxk?{pi,k?∣pi,k?∈pi?,k∈1,?,K}。這里有點對特征圖進行語義分割的味道。
然后,從特征圖E4E_{4}E4?中選擇具有高前景概率的NaN_{a}Na?個像素嵌入作為目標查詢。在這里,首先選擇在相應類別平面中具有八鄰域內局部最大值的那些點,然后在pi,kii{p_{i, k_{i}}}_{i}pi,ki??i?中選擇具有最高前景概率的那些點。這里默認了特征圖上相鄰像素點為同一個類別時,他們表示的是同一個實例,這在大部分情況下是成立的,但是對于緊密貼合的同類別小面積實例是有風險的。這樣做還有一點好處就是防止最終NaN_aNa?個Queries都集中在某一塊置信度較高的前景區域。
在訓練過程中,應用基于匹配的匈牙利損失來監督輔助分類頭,記作LIA?qL_{IA-q}LIA?q?。與DETR不同,FastInst僅使用類別預測和位置成本LlocL_{loc}Lloc?來計算分配成本。位置成本LlocL_{loc}Lloc?被定義為一個指示函數,當像素位于該物體的區域內時為0;否則為1。這種成本背后的直覺是,只有落在物體內部的像素才有理由推斷該物體的類別和掩碼嵌入。此外,位置成本減少了二分匹配空間并加速了訓練收斂。
使用這種匹配方式,就是為了讓每個GT掩碼只對應于一個特征點。另一種比較直覺的方式是直接對掩碼所對應的區域進行交叉熵損失計算,這樣每個GT掩碼就對應于一片特征點了,不利于后續topk的篩選,可能會被一個概率高的區域主導。所以還是需要二分匹配。
將通過上述策略生成的Queries稱為實例激活引導(IA引導)的查詢。與零初始化Queries相比,IA引導的查詢在初始階段就包含了關于潛在對象的豐富信息,并提高了Transformer decoder中查詢迭代的效率。
其實,這個策略也可以在E3E_{3}E3?或者E5E_{5}E5?上進行,更大的特征圖包含更豐富的實例線索,但計算負擔更重。為了權衡,作者使用中等大小的特征圖E4E_{4}E4?。
2.4 雙路徑更新策略
在從特征圖中選擇NaN_{a}Na?個由實例激活引導的Queries后,將它們與NbN_{b}Nb?個輔助可學習Queries連接起來,以獲得總查詢Q,其中輔助可學習Queries用于在后續的雙重更新過程中促進背景像素特征的分組,并提供通用的與圖像無關的信息。然后,將總查詢Q與展平的1/8高分辨率像素特征E3E_3E3?輸入到Transformer decoder中。在Transformer decoder中,為查詢Q和像素特征E3E_3E3?添加位置嵌入,隨后經過連續的Transformer decoder層對它們進行更新。一個Transformer decoder層包含一次像素特征更新和一次查詢更新。整個過程類似于EM(期望最大化)聚類算法。E步:根據像素特征所屬的中心(Query)更新像素特征;M步:更新聚類中心(Query)。與單路徑更新策略相比,雙路徑更新策略共同優化像素特征和Query,減少了對重型pixel decoder的依賴,并獲得了更細粒度的特征嵌入。最后,使用細化后的像素特征和查詢來預測每一層的目標類別和分割掩碼。
(1)位置編碼
位置編碼是非常重要的信息,特別是對于不同位置但類別相同的物體而言。本文使用一個固定尺寸的可學習位置嵌入向量P∈RS×S×256P \in R^{S \times S \times 256}P∈RS×S×256,其中SSS作者根據經驗設定為Na\sqrt{N_a}Na??,比如Na=100N_a=100Na?=100,則P∈R10×10×256P \in R^{10 \times 10 \times 256}P∈R10×10×256。如果某個特征圖需要使用位置編碼,則將PPP上采樣到對應的尺寸即可。
官方對應的代碼為
pixel_pos_embeds = F.interpolate(self.meta_pos_embed, size=pixel_feature_size,mode="bilinear", align_corners=False)
proposal_pos_embeds = F.interpolate(self.meta_pos_embed, size=proposal_size,mode="bilinear", align_corners=False)
其中,self.meta_pos_embed
就是PPP,也就是所有的特征圖都共用一個可學習的位置編碼。作者認為,這種方式可以在不影響模型效果的前提下,提高模型的推理速度。實際情況下,PPP作用于兩種特征,一種是從E3E_3E3?得到的像素特征,也即是圖2-1-1中的pixel features;另一種是從E4E_4E4?得到的NaN_aNa?個實例激活引導的Queries,也即是圖2-1-1中的NaN_aNa? IA-guided features。
對于NbN_{b}Nb?個輔助可學習Queries而言,使用了額外的可學習位置編碼,shape為Nb×256N_b \times 256Nb?×256。這是因為NbN_{b}Nb?個輔助可學習Queries學習到的Queries對應的位置,我們并不知道。
(2)像素特征更新
首先更新像素特征。給定展平的像素特征XXX(最開始就是E3E_3E3?)和查詢QQQ,像素特征更新的流程包括一個交叉注意力層和一個前饋層,如圖2-1-1右側所示。在每個交叉注意力層,都會將位置嵌入添加到query和key中。在更新像素特征時,不使用自注意力,因為像素特征的序列較長,自注意力會帶來巨大的計算和內存成本。全局特征可以通過對query的交叉注意力進行聚合。
假設輸入的XXX的shape為[N,256][N, 256][N,256],QQQ的shape為[100+Nb,256][100+N_b, 256][100+Nb?,256],NNN要比100+Nb100+N_b100+Nb?大得多。像素特征更新是,將XXX作為query,QQQ作為key和value。經過cross attention之后,XXX會被替換更新,輸出的shape仍為[N,256][N, 256][N,256]。也就是對于每一個query,都可以通過key對value進行加權求和,換句話說,XXX會被QQQ中與其相似的特征加權求和更新。
看到這里,我們其實可以進一步理解為什么作者要增加NbN_{b}Nb?個輔助可學習Queries,因為NaN_aNa?里的都是前景特征,而XXX中既有前景,又有背景,在更新XXX時,不能只使用前景特征的加權求和得到,需要一些反應背景的特征的加權求和,這個反應背景的特征就是NbN_{b}Nb?個輔助可學習Queries。
(3)Query更新
作者像Mask2Former中那樣,使用掩碼注意力,隨后再使用自注意力和前饋網絡來進行Query更新。掩碼注意力將每個Query的注意力限制在僅關注前一層預測掩碼的前景區域內,而上下文信息則被假設通過后續的自注意力來收集。這種設計顯著提升了基于Query的模型在圖像分割任務中的性能。在此,位置嵌入同樣會被添加到每個掩碼注意力層和自注意力層的query與key中。
這里的masked attention中使用的mask會在2.5中進一步說明。
(4)預測
作者在每個decoder layer中,使用兩個獨立的三層MLPs作用于細化后的實例激活引導queries來預測物體的類別和mask embeddings。每個實例激活引導查詢(IA-guided query)都需要預測所有目標類別的概率,包括 “無目標”(?)類別。作者對經過精煉的mask features進行線性投影以獲得掩碼特征。然后,將mask embeddings與mask features相乘,得到每個query對應的分割掩碼。此處,每個 Transformer decoder層中多層感知機(MLPs)和線性投影的參數是不共享的,因為query和pixel features會交替更新,且在不同的解碼器層中,它們的特征可能處于不同的表征空間。此外,實例分割需要為每個預測結果提供一個置信度分數用于評估。作者沿用先前的工作的方法,將類別概率分數與掩碼分數(即前景區域內掩碼概率的平均值)相乘,作為最終的置信度分數。
2.5 GT掩碼引導學習
mask attention的方法是以模型之前預測的mask作為先驗來對attention進行掩碼的,如果這個mask學習的不好,會影響之后的學習效果。為緩解這一問題,作者引入了GT掩碼引導學習。首先,使用最后一層的二分匹配GT真值掩碼替換第lll層掩碼注意力中使用的預測掩碼。對于在最后一層中未匹配到任何實例的Query(包括輔助可學習查詢),采用標準的交叉注意力,即
Mil={Mjgtif(i,j)∈σ?otherwise(2-5-1)M^l_i = \begin{cases} M^{gt}_j & if (i, j) \in \sigma \\ \empty & otherwise \end{cases} \tag{2-5-1} Mil?={Mjgt???if(i,j)∈σotherwise?(2-5-1)
其中,MilM^l_iMil?是第iii個query在第lll層的注意力掩碼,σ={(i,j)∣i∈{1,...,Na},j∈{1,...,Nobj}}\sigma = \{ (i, j) | i \in \{ 1, ..., N_a \}, j \in \{ 1, ..., N_{obj} \} \}σ={(i,j)∣i∈{1,...,Na?},j∈{1,...,Nobj?}}是decoder最后一層的匹配關系,并且MjgtM^{gt}_jMjgt?是第iii個query在最后一層匹配到的GT掩碼。NobjN_{obj}Nobj?表示GT目標的數量。
然后,使用該注意力掩碼MlM^lMl來進行masked attention。圖2-1-1右側所示的結構會重復多次,每一次都輸入是上一次的輸出。新的一輪根據固定匹配σσσ對新輸出進行監督,也就是σσσ的匹配過程只在第一次進行。這種固定匹配確保了每個Transformer decoder層的預測的一致性,并在訓練過程中節省了匹配的計算成本。通過這樣的指導學習,允許每個query在訓練期間看到其目標的整個區域預測對象,這有助于掩碼注意力在更合適的前景區域內進行。
2.6 損失函數
FastInst的整個損失函數可以寫作
L=LIA?q+Lpred+Lpred′(2-6-1)L = L_{IA-q} + L_{pred} + L'_{pred} \tag{2-6-1} L=LIA?q?+Lpred?+Lpred′?(2-6-1)
其中,LIA?qL_{IA-q}LIA?q?就是2.3中提到的輔助分類頭損失,也就是先使用匈牙利匹配,后進行交叉熵損失。
LpredL_{pred}Lpred?定義為
Lpred=∑i=0D(λceLcei+λdiceLdicei)+λclsLclsi(2-6-2)L_{pred} = \sum^D_{i=0} (\lambda_{ce} L^i_{ce} + \lambda_{dice} L^i_{dice}) + \lambda_{cls} L^i_{cls} \tag{2-6-2} Lpred?=i=0∑D?(λce?Lcei?+λdice?Ldicei?)+λcls?Lclsi?(2-6-2)
其中,DDD表示Transformer decoder的層數,并且i=0i=0i=0表示在進入Transformer decoder之前的prediction損失。LceiL^i_{ce}Lcei?和LdiceiL^i_{dice}Ldicei?分別表示分割任務的二元交叉熵損失和dice損失。LclsL_{cls}Lcls?表示類別交叉熵損失。同樣,利用匈牙利算法來尋找目標分配的最佳匹配。對于匹配成本,為每個query添加了額外的位置成本λlocLlocλ_{loc}L_{loc}λloc?Lloc?。
Lpred′L'_{pred}Lpred′?表示GT掩碼引導的損失,與LpredL_{pred}Lpred?相似。區別是它不計算第0層的損失,并使用固定的二分匹配結果。LpredL_{pred}Lpred?是使用模型預測的掩碼引導的,Lpred′L'_{pred}Lpred′?是GT掩碼引導的。
3 效果
FastInst和其他模型的效果對比可見表3-1。其中,FastInst-D1表示一層transformer decoder層,FastInst-D3表示三層transformer decoder層,層數速度越慢,但是精度越高。
圖3-1是實例激活引導的query的可視化結果。第一行和第二行顯示了實例激活引導的query的分布,對輔助分類頭部的監督損失不同。
第一列:密集的像素語義分類損失
第二列:基于二分匹配的匈牙利損失,沒有位置成本
第三列:基于二分的匹配匈牙利損失,有位置成本
第四列:GT
本文設計的損失(第三列)下的查詢點更集中于每個前景對象的區域。最后一行顯示四個預測的掩碼(藍色),具有相應的實例激活引導query位置(紅色)。