NeurIPS 2021丨K-Net- 邁向統一的圖像分割

K-Net: Towards Unified Image Segmentation

Wenwei Zhang¹  Jiangmiao Pang^2,4  Kai Chen^3,4  Chen Change Loy^1?

¹S-Lab, Nanyang Technological University  ²CUHK-SenseTime Joint Lab, the Chinese University of Hong Kong ³SenseTime Research ⁴Shanghai AI Laboratory

{wenwei001, ccloy}@ntu.edu.sg      panjiangmiao@gmail.com      chenkai@sensetime.com

Part 1 TL；DR

實例分割（instance segmentation）已經被以 Mask R-CNN 為代表的“先檢測后分割”的框架主導了多年，之前涌現的單階段實例分割算法也依然需要遍歷圖片中的每一個位置（feature grids）來預測實例分割掩碼（instance masks），因此它們都需要一些額外的組件（如檢測框或/和NMS）來區分不同實例或清理相同實例的掩碼。

我們希望實例分割的推理過程可以像語義分割那樣簡單：由一組卷積核（convolutional kernels） 生成一組 mask，每一個 mask 至多只分割圖片中的一個物體，且不同的 kernel 負責不同物體的 mask 生成。這樣就可以不借助任何額外的組件來完成實例分割任務（box-free and NMS-free)，并且在提高推理效率的同時實現端到端的訓練優化，同時也自然地統一了語義、實例乃至全景分割任務的解決范式。

我們提出 K-Net 來作為這個思路的一種探索，在全景分割（COCO-panoptic，test-dev set， 55.2 PQ）和語義分割（ADE20K val set，54.3 mIoU）上取得了新的 state-of-the-art 結果，在相同的實例分割精度下模型的推理速度比 Cascade Mask R-CNN 快 60-90% 。

Part 2 實例分割的N種姿勢

自 Mask R-CNN 被提出以來，“先檢測后分割”的思路統治了實例分割任務很長一段時間，直到最近兩年才涌現出一些單階段實例分割的探索。近些年來，實例分割方法的思路總結如下圖：

上圖中的4類方法都引入了額外的組件來區分不同的實例或者消除重復的實例，我們一一來看：

Top-down 方法：例如 Mask R-CNN，Cascade Mask R-CNN，以及 HTC 等，都是采用的先檢測，后分割的思路，通過框來區分不同物體并獲得不同物體的特征圖，再進行實例分割。這類算法都會依賴檢測框和 NMS。

Bottom-up 方法：例如 Associate Embedding 和 Instance Cut 采用的是先語義分割，然后通過一些 grouping 的過程來區分不同實例。這類算法依賴聚類的操作（grouping process）。

Dense Mask Prediction：從2019年TensorMask開始，有一些方法嘗試不依賴檢測框，直接從每個feature grid 來預測 instance mask，比如 TensorMask 使用 sliding window，SOLO 把圖像拆成不重疊的 grids。因為他們都是遍歷 CNN 的 feature grids 來生成密集的 instance mask，所以都需要 NMS 來消除重復的 mask。

Dense Kernel Prediction：還有一些新的探索通過預測 kernel 來生成 mask，但是 kernel 的生成來自于 dense feature grids，是一個位置一個 kernel，因此需要檢測框或者NMS來消除重復的實例（例如 SOLO v2 和 CondInst） 。

Part 3 讓實例分割像語義分割那樣簡單

此時反觀語義分割，我們會發現，自全卷積網絡 （FCNs）提出以來，語義分割任務的解決思路就沒怎么變過，后來的工作主要是在表征上下功夫，如PSPNet，DeepLab系列，以及各類 attention 網絡等。

其預測語義分割 mask 的核心結構如下圖所示，就是由一組 kernel 來負責語義 mask 的生成，同時由于語義分割任務的特點，我們可以讓kernel數量和語義類的數量保持一致，并讓每一個 kernel 負責一個固定語義類別 mask 的生成。

理論上來說，圖像分割的本質就是把圖片中的 pixel 分成具有不同特性的 group，對語義分割來說，一個 group 代表一個語義類型，對實例分割來說，一個 group 代表一個實例。既然語義分割可以使用 kernel 來解決，做到一個 kernel 負責分割 一個 group，實例分割能不能也用這樣簡單的框架來解決呢？

看起來，我們也可以引入一組卷積核來負責 mask 的生成，只要我們限定一個 kernel 只分割一個物體，同時讓每個kernel負責分割不同的物體，實例分割任務的推理好像就完成了，順帶著全景分割也被統一到一個框架內了？

順著這個思路，實例分割和全景分割的模型框架就變成了下面這樣簡單：

如上圖所示，在通過 backbone 和 neck 得到 2D 特征圖以后，有一組可學習的 kernel 和特征圖卷積得到初始的 mask 預測。

在拿到 mask 預測和 mask 類別預測以后我們可以像 DETR 那樣將 instance mask 和 ground truth masks 以 mask loss 作為 cost 進行 Bipartite matching ，匹配得到模型學習的target，然后整個模型就可以 end-to-end 進行訓練推理了，具體使用的網絡結構、 loss 和超參等細節可以見 paper。

Part 4 Why now?

這個時候你其實會想：這么簡單的思路，為什么之前大家沒有想到？

原因其實挺簡單的，就是 DETR 之前大家都很難想到，原來目標檢測可以直接學一組數量有限的 query，然后基于 transformer + 足夠久的模型訓練（之前大家也沒那么富足，一般不會訓 300 個 epoch），就可以使每一個 query 學到只負責一個物體的檢測框預測。

DETR 除了使用 transformer，其實還有使用了一個之前不受大家重視的關鍵技術，那就是 Bipartite matching。

目標檢測/實例分割和其他任務不一樣的點就在于，這個任務的 target 是一個檢測框或者實例分割掩碼的集合，而 CNN 學出來的表征是 dense 的，所以大家之前基于 dense 的 feature 都是做 dense 的 prediction，然后通過 NMS 等手段清理重復的預測。

而 Bipartite matching 實質上解決了由一組 query 預測得到的 instance set 如何去匹配 ground truth instances 的問題，也使得這樣的一個框架不需要 NMS。然后 DETR 用 transformer + 300 epoch 訓練就把這樣一套框架做 work 了。

當然，在 DETR 的 follow-up 工作中，Sparse R-CNN 也進一步證明了可以直接學一組 Bounding box proposal，然后通過逐階段 refine 來得到很高的目標檢測精度。Sparse R-CNN 在檢測任務上的成功進一步驗證了我們的想法，最終促使我們在這樣一個時間節點去勇敢嘗試 K-Net。

Part 5 Group-aware Kernels

雖然理論上一組 instance kernel 就可以得到實例分割的預測結果，但實際上我們會發現這樣得到的結果效果不盡如人意。

我們認為這實際上是因為實例分割任務對它所需要的 kernel （instance kernel）相比語義分割需要的 kernel （semantic kernel）有更高的要求，主要原因有如下兩點：

1.instance kernel 其實不像 semantic kernel 那樣可以具備一些顯式的特性來便利學習。例如，由于每一個單獨的 semantic kernel 都可以和一個唯一的語義類別（ semantic class ）綁定，因此在學習的時候它在每張圖上都可以學習著去分割同一個語義類別，而 instance kernel不具備這樣的特性，所以我們是通過 Bipartite matching 來做的 target assignment，這就導致了每個 kernel 在每張圖上學習的目標是根據他們當前的預測情況動態分配的。

2.特點1就導致了，instance kernel 實際上要去區分外觀（appearance）和尺度（scale）高度變化的物體，需要具備更強的判別特性（discriminative capability）。

此時一個直觀的想法就是，直接用圖片里的內容來增強 kernel，讓它能獲取當前圖片的一些信息（content-aware）。

那么，圖片里的哪些內容是這個 kernel 所需要的呢？我們認為，應該就是 kernel 和特征圖響應產生 mask 的部分，因為 mask 本質上是 kernel 對每個 pixel 是否屬于它對應的 group的一種 prediction 或者 assignment，如果讓 kernel 通過 mask 獲取到 kernel 所對應的 pixel group 的信息，理論上新的 kernel 再去做分割的時候，得到的結果不應該比當前的分割結果要差。

因此，我們設計了一個 Kernel Update Head 基于 mask 和特征圖來將 kernel 動態化。如下圖所示，Kernel Update Head 首先獲得每個 kernel 對應 pixel group 的 feature，然后以某種方式動態地更新當前的 kernel。

為了讓 kernel 還能夠 modeling 全局的信息，我們增加了一個 kernel interaction 模塊，最終得到的特征可以用于分類并產生的 dynamic kernel 來和特征圖卷積得到更加精確的 mask prediction。Adaptive Kernel Update 和 Kernel Interaction 的形式都可以用很多種，我們效仿 LSTM 設計了一種 Adaptive Kernel Update，然后為了方便再 Kernel Interaction 里用了 MultiHeadAttention。

具體每個 component 的設計細節和 ablation study 歡迎大家參考我們的 paper 和 code。

我們可以添加多個 Kernel Update Head 來對 mask 和 kernel 進行 iterative 的 refine。最終，完整的 K-Net pipeline 如下圖所示。論文中僅使用 3 個 Kernel Update Head 和 100 個instance kernel 就可以得到各個benchmark 上 state-of-the-art 的結果。

Part 6 實驗結果

我們把 K-Net 和最近的一些全景分割算法做了一些比較，在 COCO-panoptic 上以最樸素的訓練方式（多尺度訓練 36 epoch，訓一個 K-Net 只需要 16張 V100 訓兩天半）取得了高于其他方法的結果。

PS：這個版本只用了 100 個 instance kernel，用的是 window size=7 （MaskFormer 用的 window size=12）的 Swin-Large backbone ，所以理論上結果還能更高，給大家留足充分的刷點空間。

我們在實例分割上也做了一些對比，K-Net 在比之前SOLOv2/CondInst/Mask R-CNN 等算法都要快的情況下取得了更優的精度，在和 Cascade Mask R-CNN 精度持平的情況下推理速度快 60-90%。

同時，K-Net 也可以直接和當前基于 semantic kernel 的語義分割算法組合，進一步提升他們的性能，將 FCN，PSPNet，DeepLab v3，UperNet 等算法提高了 1.1-6.6 mIoU。即使在 UperNet + Swin-L 的情況下，仍然能將模型提高 1.2 mIoU ，超越了 UperNet + Swin-L。

Part 7 討論

7.1 和 MaskFormer，MaX-DeepLab 的區別

其實看到 MaX-DeepLab 的時候 K-Net 的實例分割部分已經完成，全景分割也調地差不多了。看到 MaskFormer 也已經是7月份了。

K-Net 和 Max-DeepLab 差別還是挺明顯的，MaX-DeepLab 主要針對end-to-end 全景分割從 backbone 到 head 都用 transformer 增強了一遍，相應的 loss 設計等也都是針對全景分割的 metric，最后的推理方式也和 K-Net 不太一樣，K-Net 沒有拘泥于全景分割，提供了一個統一不同圖像分割任務的視角，里面其實也沒有 transformer 這個概念（就是不太想用）。

而我們看到 MaskFormer 的時候有一種殊途同歸的感覺，K-Net 最開始是希望把實例分割做地像語義分割那樣簡單，抓住的是語義分割里由一組 kernel 出一組 mask 的范式，通過這個范式統一了各個分割任務；而MaskFormer 的本意是 rethinking 語義分割，抓住 instance-level segmentation 里 mask classification 的核心 design 來解決語義分割任務，最終也統一了不同分割任務的框架。

兩個方法最終輸出 mask prediction 的本質是一樣的，都是一組 kernel 分割得到一組 mask，進而對 mask 做分類，只是兩個方法生成 kernel 的方式不同，一個是 Transformer 大法好，一個是 iterative refine 的思想。

7.2 What is next？

我們認為如何最有效地生成具有高 discriminative ability 的 kernel 仍然是一個值得探索的問題， K-Net 只是對這一范式的一種簡單探索，很多結構其實我們都還沒來得及嘗試。

另外，我們觀察了 K-Net 的failure case，發現了兩個常見問題，一個是 mask classification 容易出錯，尤其是具有相同紋理的類別，另一個是 mask 的 boundary 有時候還不夠好，也是在有相同紋理的地方容易分割出奇怪的內容，具體可以參考 paper 的 appendix。我們覺得這可能是 Mask Classification 這類方法引入的一些新的問題，分割任務也由此有了新的挑戰。