一種使用Mask Transformer進行全景分割的端到端解決方案

【導讀】近日，谷歌團隊推出了一項新Transformer，可用于優(yōu)化全景分割方案，還登上了CVPR 2022。

最近，谷歌AI團隊受Transformer和DETR的啟發(fā)提出了一種使用Mask Transformer進行全景分割的端到端解決方案。

全稱是end-to-end solution for panoptic segmentation with mask transformers，主要用于生成分割MaskTransformer架構的擴展。

該解決方案采用像素路徑（由卷積神經(jīng)網(wǎng)絡或視覺Transformer組成）提取像素特征，內存路徑（由Transformer解碼器模塊組成）提取內存特征，以及雙路徑Transformer用于像素特征和內存之間的交互特征。

然而，利用交叉注意力的雙路徑Transformer最初是為語言任務設計的，它的輸入序列由幾百個單詞構成。

而對視覺任務尤其是分割問題來說，其輸入序列由數(shù)萬個像素組成，這不僅表明輸入規(guī)模的幅度要大得多，而且與語言單詞相比也代表了較低級別的嵌入。

全景分割是一個計算機視覺問題，它是現(xiàn)在許多應用程序的核心任務。

它分為語義分割和實例分割兩部分。

語義分割就比如為圖像中的每個像素分配語義標簽，例如「人」和「天空」。

而實例分割僅識別和分割圖中的可數(shù)對象，如「行人」和「汽車」，并進一步將其劃分為幾個子任務。

每個子任務單獨處理，并應用額外的模塊來合并每個子任務階段的結果。

這個過程不僅復雜，而且在處理子任務和整合不同子任務結果時還會引入許多人工設計的先驗。

在 CVPR 2022 上發(fā)表的「CMT-DeepLab: Clustering Mask Transformers for Panoptic Segmentation」中，文章提出從聚類的角度重新解讀并且重新設計交叉注意力cross attention（也就是將相同語義標簽的像素分在同一組），從而更好地適應視覺任務。

CMT-DeepLab 建立在先前最先進的方法 MaX-DeepLab 之上，并采用像素聚類方法來執(zhí)行交叉注意，從而產(chǎn)生更密集和合理的注意圖。

kMaX-DeepLab 進一步重新設計了交叉注意力，使其更像一個 k-means 聚類算法，對激活函數(shù)進行了簡單的更改。

結構總覽

研究人員將從聚類的角度進行重新解釋，而不是直接將交叉注意力應用于視覺任務而不進行修改。

具體來說，他們注意到Mask Transformer 對象查詢可以被認為是集群中心（旨在對具有相同語義標簽的像素進行分組）。

交叉注意力的過程類似于 k-means 聚類算法，（1）將像素分配給聚類中心的迭代過程，其中可以將多個像素分配給單個聚類中心，而某些聚類中心可能沒有分配的像素，以及（2）通過平均分配給同一聚類中心的像素來更新聚類中心，如果沒有分配像素，則不會更新聚類中心）。

在CMT-DeepLab和kMaX-DeepLab中，我們從聚類的角度重新制定了交叉注意力，其中包括迭代聚類分配和聚類更新步驟

鑒于 k-means聚類算法的流行，在CMT-DeepLab中，他們重新設計了交叉注意力，以便空間方面的softmax操作（即沿圖像空間分辨率應用的 softmax 操作），實際上將聚類中心分配給相反，像素是沿集群中心應用的。

在 kMaX-DeepLab 中，我們進一步將空間方式的 softmax 簡化為集群方式的 argmax（即沿集群中心應用 argmax 操作）。

他們注意到 argmax 操作與 k-means 聚類算法中使用的硬分配（即一個像素僅分配給一個簇）相同。

從聚類的角度重新構建MaskTransformer的交叉注意力，顯著提高了分割性能，并簡化了復雜的Masktransformer管道，使其更具可解釋性。

首先，使用編碼器-解碼器結構從輸入圖像中提取像素特征。然后，使用一組聚類中心對像素進行分組，這些像素會根據(jù)聚類分配進一步更新。最后，迭代執(zhí)行聚類分配和更新步驟，而最后一個分配可直接用作分割預測。

為了將典型的MaskTransformer解碼器（由交叉注意力、多頭自注意力和前饋網(wǎng)絡組成）轉換為上文提出的k-means交叉注意力，只需將空間方式的softmax替換為集群方式最大參數(shù)。

本次提出的 kMaX-DeepLab 的元架構由三個組件組成：像素編碼器、增強像素解碼器和 kMaX 解碼器。

像素編碼器是任何網(wǎng)絡主干，用于提取圖像特征。

增強的像素解碼器包括用于增強像素特征的Transformer編碼器，以及用于生成更高分辨率特征的上采樣層。

一系列 kMaX 解碼器將集群中心轉換為 (1) Mask嵌入向量，其與像素特征相乘以生成預測Mask，以及 (2) 每個Mask的類預測。

kMaX-DeepLab 的元架構

研究結果

最后，研究小組在兩個最具挑戰(zhàn)性的全景分割數(shù)據(jù)集 COCO 和 Cityscapes 上使用全景質量 (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ù)集。

從聚類的角度設計，kMaX-DeepLab 不僅具有更高的性能，而且還可以更合理地可視化注意力圖以了解其工作機制。

在下面的示例中，kMaX-DeepLab 迭代地執(zhí)行聚類分配和更新，從而逐漸提高Mask質量。

kMaX-DeepLab 的注意力圖可以直接可視化為全景分割，讓模型工作機制更合理

結論

本次研究展示了一種更好地設計視覺任務中的MaskTransformer的方法。

通過簡單的修改，CMT-DeepLab 和 kMaX-DeepLab 重新構建了交叉注意力，使其更像一種聚類算法。

因此，所提出的模型在COCO 和 Cityscapes數(shù)據(jù)集上實現(xiàn)了最先進的性能。

研究團隊表示，他們希望 DeepLab2 庫中 kMaX-DeepLab 的開源版本有助于未來對專用于視覺Transformer架構設計的研究。

審核編輯 ：李倩