ISTR：實例分割Transformer，首次驗證了Transformer在端到端實例分割中的潛力，表現(xiàn)SOTA！性能優(yōu)于Sparse R-CNN、SOLOv2等網絡，代碼已經開源！
作者單位：廈門大學(紀榮嶸團隊), Pinterest, 騰訊優(yōu)圖, IIAI

1 簡介

在本文中提出了一種稱為ISTR的實例分割Transformer，它是首個基于Transformer的端到端框架。ISTR通過預測低維Mask嵌入，并將其與Ground-Truth Mask嵌入進行匹配以得到Set Loss。此外，ISTR同時使用循環(huán)細化策略進行檢測和分割，與現(xiàn)有的自上而下和自下而上的框架相比，它提供了一種實現(xiàn)實例分割的新方法。

得益于所提出的端到端機制，即使使用基于近似的次優(yōu)嵌入，ISTR仍能展示出最先進的性能。具體而言，ISTR在MS COCO數(shù)據(jù)集上使用ResNet50-FPN獲得46.8/38.6box/maskAP，并使用ResNet101-FPN獲得48.1/39.9 box/mask AP。

2 所提方法

2.1 Mask Embeddings

為了提供一個有效提取Mask Embeddings的公式，作者對原始和重構Mask之間的互信息進行了約束:

其中， 為2個隨機變量之間的互信息，M為Mask集合， 為用于提取嵌入的Mask編碼器， 為用于重建Mask的解碼器，上式保證了編碼和解碼階段的信息損失最小，這隱含用嵌入表征Mask。推導后可以得到Mask Embeddings的廣義目標函數(shù):

式中， 為Mask Embeddings， 為L2-norm，通過使用矩陣( )對編碼器和解碼器進行簡單的線性變換，最終目標函數(shù)可以表達為：

式中，l為Mask Embeddings的維數(shù)， 為 單位矩陣。上式與PCA的目標函數(shù)有相同的公式，它提供了一個closed-form solution來學習transformation。

2.2 Matching Cost and Prediction Loss

在得到Mask Embeddings的編碼器和解碼器后，本文還定義了端到端實例分割的bipartite matching cost和prediction loss。將Ground truth bounding boxes, classes和masks表示為 。predicted bounding boxes, classes和mask embeddings表示為，其中

2.2.1 Bipartite Matching Cost

對于Bipartite Matching，以最小代價尋找n個非重復整數(shù) 的排列，如：

將bounding boxes的matching cost定義為:

class的matching cost為:

其中 表示平衡cost的超參數(shù)， 表示L1成本， 表示generalized IoU cost， 是類 的概率。

這里不再直接匹配高維mask，而是使用mask embeddings之間的相似性度量來匹配它們，定義為：

其中mask embeddings L2歸一化，2個歸一化向量之間的點積用于計算余弦相似度。結果加上1，然后除以2保證取值范圍在[0,1]之間。

2.2.2 Set Prediction Loss

對于set prediction loss使用匹配預測回歸mask ground truth。set prediction loss定義為:

定義與 相同, 為用于分類的Focal loss。對于masks，加入dice loss來改進學習到的嵌入以重構mask。mask loss定義為:

為L2 loss， 表示dice loss。

2.3 Instance Segmentation Transformer

ISTR的結構如圖2所示。它包含4個主要組件:

• CNN backbone with FPN 用于提取每個實例的特征；
• a Transformer encoder with dynamic attention，用于學習物體之間的關系；
• a set of prediction heads，用于同時進行檢測和分割；
• N-step recurrent update用于細化預測集。

2.3.1 Backbone

對于ROI特征： 首先，這里使用一個帶有FPN的CNN Backbone來提取特征金字塔的P2到P5級的特征。

然后，利用初始化k個可學習query boxes  覆蓋整個圖像的，通過RoIAlign提取k個RoI特征 。

對于Image特征：對P2到P5的特征進行平均求和，提取圖像特征 ，

2.3.2 Transformer Encoder and Dynamic Attention

首先對圖像特征和位置嵌入的和進行3個可學習權矩陣的變換，得到輸入 ; ; ,self-attention模塊定義為:

multi-head attention包括多個self-attention塊，例如最初的Transformer中的8個用來封裝不同特性之間的多個復雜關系。為了更好地融合RoI和圖像特征增加了dynamic attention模塊，定義為第 步中對RoI特征 的關注:

其中 為全連接層，用來生成dynamic parameters。然后在預測頭中使用獲得的特征 來產生輸出。

2.3.3 Prediction Heads

預測由Prediction Heads計算得到，包括class head, box head, mask head以及fixed mask decoder。

box head：預測歸一化中心坐標、高度和寬度的殘差值，用于在第i步中更新query boxes ，

class head：使用softmax函數(shù)預測類。

mask head：輸出掩模嵌入，然后使用預先學習的mask decoder重建以預測mask。

2.3.4 Recurrent Refinement Strategy

query boxes  被預測框周期性地更新，這改進了預測并使同時處理檢測和分割成為可能。整個過程如下：

3 Experiments

上表作者分別對于所設計的Mask embedding、Mask Cost Function、LossFunction以及Attention type進行了消融實驗，各位老鐵可以看出本文設計的有效性。

通過SOTA對比可以看出，在使用相同的Backbone的基礎上，ISTR比Mask R-CNN、BlendMask、SOLO V2以及Sparse R-CNN都要好，同時推理時間并沒有增加多少，確實是一個不錯的工作。

4 參考

[1].ISTR: End-to-End Instance Segmentation with Transformers
[2].https://github.com/hujiecpp/ISTR

5 推薦閱讀

CVPR2021 | 重新思考BiSeNet讓語義分割模型速度起飛（文末獲取論文）

YOLO在升級 | PP-YOLO v2開源致敬YOLOV4攜帶Tricks又準又快地歸來（附論文與源碼）

簡單有效 | Transformer通過剪枝降低FLOPs以走向部署（文末獲取論文）

又改ResNet | 重新思考ResNet：采用高階方案的改進堆疊策略（附論文下載）

掃描下方二維碼關注【AI人工智能初學者】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行?。?！

點“在看”給我一朵小黃花唄