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@qidebank.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 这类方法引入的一些新的问题，分割任务也由此有了新的挑战。