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 这类步骤引入的一些新的问题，宰割工作也由此有了新的挑战
。