谷歌團隊推出新Transformer,優(yōu)化全景分割方案
?最近,谷歌AI團隊受Transformer和DETR的啟發(fā)提出了一種使用Mask Transformer進行全景分割的端到端解決方案。
全稱是end-to-end solution for panoptic segmentation with mask transformers,主要用于生成分割MaskTransformer架構(gòu)的擴展。
該解決方案采用像素路徑(由卷積神經(jīng)網(wǎng)絡(luò)或視覺Transformer組成)提取像素特征,內(nèi)存路徑(由Transformer解碼器模塊組成)提取內(nèi)存特征,以及雙路徑Transformer用于像素特征和內(nèi)存之間的交互特征。
然而,利用交叉注意力的雙路徑Transformer最初是為語言任務(wù)設(shè)計的,它的輸入序列由幾百個單詞構(gòu)成。
而對視覺任務(wù)尤其是分割問題來說,其輸入序列由數(shù)萬個像素組成,這不僅表明輸入規(guī)模的幅度要大得多,而且與語言單詞相比也代表了較低級別的嵌入。
全景分割是一個計算機視覺問題,它是現(xiàn)在許多應(yīng)用程序的核心任務(wù)。
它分為語義分割和實例分割兩部分。
語義分割就比如為圖像中的每個像素分配語義標簽,例如「人」和「天空」。
而實例分割僅識別和分割圖中的可數(shù)對象,如「行人」和「汽車」,并進一步將其劃分為幾個子任務(wù)。
每個子任務(wù)單獨處理,并應(yīng)用額外的模塊來合并每個子任務(wù)階段的結(jié)果。
這個過程不僅復(fù)雜,而且在處理子任務(wù)和整合不同子任務(wù)結(jié)果時還會引入許多人工設(shè)計的先驗。
在 CVPR 2022 上發(fā)表的「CMT-DeepLab: Clustering Mask Transformers for Panoptic Segmentation」中,文章提出從聚類的角度重新解讀并且重新設(shè)計交叉注意力cross attention(也就是將相同語義標簽的像素分在同一組),從而更好地適應(yīng)視覺任務(wù)。
CMT-DeepLab 建立在先前最先進的方法 MaX-DeepLab 之上,并采用像素聚類方法來執(zhí)行交叉注意,從而產(chǎn)生更密集和合理的注意圖。
kMaX-DeepLab 進一步重新設(shè)計了交叉注意力,使其更像一個 k-means 聚類算法,對激活函數(shù)進行了簡單的更改。
結(jié)構(gòu)總覽
研究人員將從聚類的角度進行重新解釋,而不是直接將交叉注意力應(yīng)用于視覺任務(wù)而不進行修改。
具體來說,他們注意到Mask Transformer 對象查詢可以被認為是集群中心(旨在對具有相同語義標簽的像素進行分組)。
交叉注意力的過程類似于 k-means 聚類算法,(1)將像素分配給聚類中心的迭代過程,其中可以將多個像素分配給單個聚類中心,而某些聚類中心可能沒有分配的像素,以及(2)通過平均分配給同一聚類中心的像素來更新聚類中心,如果沒有分配像素,則不會更新聚類中心)。
在CMT-DeepLab和kMaX-DeepLab中,我們從聚類的角度重新制定了交叉注意力,其中包括迭代聚類分配和聚類更新步驟
鑒于 k-means聚類算法的流行,在CMT-DeepLab中,他們重新設(shè)計了交叉注意力,以便空間方面的softmax操作(即沿圖像空間分辨率應(yīng)用的 softmax 操作),實際上將聚類中心分配給相反,像素是沿集群中心應(yīng)用的。
在 kMaX-DeepLab 中,我們進一步將空間方式的 softmax 簡化為集群方式的 argmax(即沿集群中心應(yīng)用 argmax 操作)。
他們注意到 argmax 操作與 k-means 聚類算法中使用的硬分配(即一個像素僅分配給一個簇)相同。
從聚類的角度重新構(gòu)建MaskTransformer的交叉注意力,顯著提高了分割性能,并簡化了復(fù)雜的Masktransformer管道,使其更具可解釋性。
首先,使用編碼器-解碼器結(jié)構(gòu)從輸入圖像中提取像素特征。然后,使用一組聚類中心對像素進行分組,這些像素會根據(jù)聚類分配進一步更新。最后,迭代執(zhí)行聚類分配和更新步驟,而最后一個分配可直接用作分割預(yù)測。
為了將典型的MaskTransformer解碼器(由交叉注意力、多頭自注意力和前饋網(wǎng)絡(luò)組成)轉(zhuǎn)換為上文提出的k-means交叉注意力,只需將空間方式的softmax替換為集群方式最大參數(shù)。
本次提出的 kMaX-DeepLab 的元架構(gòu)由三個組件組成:像素編碼器、增強像素解碼器和 kMaX 解碼器。
像素編碼器是任何網(wǎng)絡(luò)主干,用于提取圖像特征。
增強的像素解碼器包括用于增強像素特征的Transformer編碼器,以及用于生成更高分辨率特征的上采樣層。
一系列 kMaX 解碼器將集群中心轉(zhuǎn)換為 (1) Mask嵌入向量,其與像素特征相乘以生成預(yù)測Mask,以及 (2) 每個Mask的類預(yù)測。
kMaX-DeepLab 的元架構(gòu)
研究結(jié)果
最后,研究小組在兩個最具挑戰(zhàn)性的全景分割數(shù)據(jù)集 COCO 和 Cityscapes 上使用全景質(zhì)量 (PQ) 度量來評估 CMT-DeepLab 和 kMaX-DeepLab,并對比 MaX-DeepLab 和其他最先進的方法。
其中CMT-DeepLab 實現(xiàn)了顯著的性能提升,而 kMaX-DeepLab 不僅簡化了修改,還進一步提升了,COCO val set 上的 PQ 為 58.0%,PQ 為 68.4%,44.0% Mask平均精度(Mask AP),Cityscapes 驗證集上的 83.5% 平均交集比聯(lián)合(mIoU),沒有測試時間增強或使用外部數(shù)據(jù)集。
從聚類的角度設(shè)計,kMaX-DeepLab 不僅具有更高的性能,而且還可以更合理地可視化注意力圖以了解其工作機制。
在下面的示例中,kMaX-DeepLab 迭代地執(zhí)行聚類分配和更新,從而逐漸提高Mask質(zhì)量。
kMaX-DeepLab 的注意力圖可以直接可視化為全景分割,讓模型工作機制更合理
結(jié)論
本次研究展示了一種更好地設(shè)計視覺任務(wù)中的MaskTransformer的方法。
通過簡單的修改,CMT-DeepLab 和 kMaX-DeepLab 重新構(gòu)建了交叉注意力,使其更像一種聚類算法。
因此,所提出的模型在COCO 和 Cityscapes數(shù)據(jù)集上實現(xiàn)了最先進的性能。
研究團隊表示,他們希望 DeepLab2 庫中 kMaX-DeepLab 的開源版本有助于未來對專用于視覺Transformer架構(gòu)設(shè)計的研究。?