K-Net：Towards Unified Image Segmentation，基于动态内核的通用分割网络，（NMS-free and Box-free），从语义/实例分割到全景分割。

K-Net：《K-Net: Towards Unified Image Segmentation》
发布于2021 。

论文地址：K-Net: Towards Unified Image Segmentation

本文将包含以下内容：

1. K-Net实现原理。
2. 动态内核如何实现。
3. K-Net如何应用于语义分割。
4. K-Net如何实现NMS-Free与Box-Free的实例分割。
5. K-Net如何实现全景分割。

前言

语义分割、实例分割和全景分割存在着潜在的链接，但现有的分割方法往往将其割裂开来，譬如最经典的FCN用于语义分割，而在实例分割中，Mask-RCNN则是较为经典的网络。而全景分割则是语义和实例两者的结合，要求在分割实例的情况下提供更全面的场景理解。

• 语义（Semantic）分割模型可以直接根据图像像素进行分组，转换为密集的分类问题，如FCN中，根据分类类别来学习一组卷积核来将像素进行分类，即可简单而优雅地完成语义分割任务，而以后的众多语义分割方法均基于FCN，无论是DeepLab系列的扩展卷积和ASPP，还是PSPNet，以及Non-Local（Self-Attention）形式的多数网络。

• 但是，考虑到图像之间的不同实例，将语义分割扩展到实例（Instance）分割是很不错的思路。当然，实例分割需要更加复杂的模型构成和后处理机制，诸如NMS和Box。实例分割一般可分为两个传统框架，“自上而下” 和 “自下而上”。自上而下的框架是先计算实例的检测框，在检测框内进行分割；相反，自下而上的框架则是先进行语义分割，在分割结果上对实例对象进行检测。直觉上而言，自上而下的方法会相对更为准确，事实上结果也如此。如经典的实例分割网络Mask-RCNN即采用自上而下的方法，在Faster R-CNN中加上了一个FCN语义分割模块。当然，现在涌现出了更多的实例分割方法，包括一阶段的YoloAct、BlendMask、EmbedMask等；还有一些更加骚的操作，基于Contour的PolarMask还有发在2022CVPR的E2EC；以及基于Transformer的Query-based的QueryInst、SOLQ、Mask2Former等。

• 全景（Panoptic）分割结合了实例和语义分割，以提供对场景的更丰富的理解。主流框架在实例分割框架上添加语义分割分支，或基于语义分割方法采用不同的像素分组策略。

• Dynamic Kernel：卷积核通常是静态的，但由于输入是不可知的（即输入往往与训练数据有差异），因此模型的效果十分受限。在一些工作中，如Dynamic filter networks、Deformable convolutional networks(DCN)的可变形卷积等等，通过核的动态变形来提高模型的灵活性和性能。一些语义分割方法也应用动态内核来通过扩大的接收场（LS-DFN）或多尺度上下文（DMNet）来改善模型表示。

• K-Net：针对以上三个任务，K-Net提出了一种基于动态内核的分割模型，为每个任务分配不同的核来实现多任务统一（图1）。

原文中有如下表述：Instead of generated from dense grids, the kernels in K-Net are a set of learnable parameters updated by their corresponding contents in the image. K-Net does not need to handle duplicated kernels because its kernels learn to focus on different regions of the image in training, constrained by the bipartite matching strategy that builds a one-to-one mapping between the kernels and instances.
翻译一下就是：K-Net的kernels是一组科学系的参数，根据图像对应的内容动态更新。同时，K-NET不需要处理重复的内核，因为其内核学会专注于图像的不同区域，受到双方匹配的约束策略，该策略在内核和实例之间建立了一对一的映射。

图1 不同的内核用于不同分割任务

K-Net

K-Net的本质

对于所有细分的分割任务，其内核都是对有意义的像素进行分组，譬如语义分割，将不同的类别像素分组。理论上而言，分割任务的分组是有上限的，因此，可以把分组数量认为设置为N ，比如，有N个用于语义分割的预定义语义类或图像中最多有N个实例对象，对于全景分割，N是图像中的stuff类和Instance类的总数。
那么，可以使用N个Kernel将图像将图像划分为N组，每个Kernel都负责找到属于其相应组的像素（也就是先前讲的，Kernel与Content实现一对一映射）。具体而言，给定由深神经网络产生的B图像的输入特征映射$F\in R^{B\times C\times H\times W}$，我们只需要N个kernels $K\in R^{N\times C}$即可用F执行卷积以获得相应的的预测结果$M\in R^{B\times N\times H\times W}$。
$$M=\sigma (K\ast F)$$

使用$\sigma$对结果进行激活，设置对应阈值后即可得到N个二进制掩码Mask（这也是语义分割长期以来的基本思想 (图1-a)。
但是，为了实现实例分割，需要对每一个Kernel进行限制，也就是每个Kernel最多只能处理图像中一个对象(图1-b)，通过这种方式，K-NET可以区分实例并同时执行分割，从而在一个特征映射中实现实例分割，而无需额外的步骤（实现了NMS-Free和Box-Free）。
为简单起见，我们可以将这些内核称为本文中的语义内核和实例内核，分别用于语义和实例分割。
实例内核和语义内核的简单组合即可实现全景分割，该分割将像素分配给实例对象或一类东西 (图1-c)。

Group-Aware Kernels

虽然使用Kernel来区分语义类别是十分简单的，但是要区分实例对象就显得比较困难。因为实例内核需要区分图像内部和跨图像内外部的对象。不像语义类别具有共同和明确的特征，实例内核需要拥有比语义内核更强的判别能力。

因此作者这里提出了一个Kernel Update策略，来使每一个内核对应一个像素组。

Kernel Update Head $f_i$包含三个关键步骤：group feature assembling、adaptive kernel update 和 kernel interaction。

1. group feature assembling：首先，通过Mask $M_{i-1}$来计算聚合出一个组特征映射$F^K$。其中，每一个组group都对应着一个语义类\实例对象。
2. adaptive kernel update： 这个$F^K$内的每一个组都将会被用来更新内核$K_{i-1}$ 。
3. kernel interaction：随后，这些内核$K_{i-1}$进行交互，互相能够提供上下文信息以对全局特征图进行建模，获得新的内核$K_i$ 。
4. 最后，使用$K_i$对特征图F进行卷积，得到预测结果更加精确的Mask $M_i$ 。
5. 用公式表达就是：
   $$K_i,M_i=f_i(M_{(i-1)},K_{(i-1)},F)$$

图2 Kernel Update Head

可以细看一下Kernel Update Head中的三个步骤。

Group Feature Assembling

Kernel Update Head首先聚合每个group的特征，随后使用这些特征来实现核的group-aware，也就是每一个核能够对特征图正确感知。由于$M_{i-1}$中每个内核的掩码本质上定义了像素是否属于该内核的相关组，可以通过将特征映射$F$乘$M_{i-1}$作为新的组装特征$F^K$:
$$F^K=\sum _u^H\sum _v^W{M}_{(i-1)}(u,v)\cdot F(u,v),F^K\in R^{B\times N\times C}$$

Adaptive Feature Update

经过Group Feature Assembling之后，我们可以使用$F^K$计算得到一组新的Kernels。但作者考虑到mask $M_{i-1}$可能不够准确，可能包含了其他组被误分类进来的噪音，因此设计了一个自适应的内核更新策略，首先在$F^K$和$K_{i-1}$之间执行元素乘法(${\varphi }_1$、${\varphi }_2$为线性变换)：
$$F^G={\varphi }_1(F^K)\otimes {\varphi }_2{K}_{i-1},F^G\in R^{B\times N\times C}$$
随后，计算两个门控gates，$G^F$和$G^K$：
$$G^K=\sigma ({\Phi }_1(F^G))、G^F=\sigma ({\Phi }_2(F^G))$$

再由这两个gates计算出一组新的kernels $\tilde{K}$

$$\tilde{K}=G^F\otimes {\Phi }_3(F^K)+G^K\otimes {\Phi }_4(K_{i-1})$$
其中，${\Phi }_n$函数均为Fully connected layers（全连接层）+ Layer Norm（层归一化）。计算结果$\tilde{K}$则将用于Kernel Interaction中。

在这里作者提到，计算出来的两个gate，$G^F$和$G^K$，可以看做是Transformer中的self-attention机制，本质上是计算一个权重对特征$F^K$和核$K_{i-1}$做一个加权求和，当然，公式中所表达的含义也是如此，无非是对一些噪声信息做加权罢了。

Kernel Interaction

Kernel Interaction可以使不同的kernel之间互相信息流通，也就是能够提供上下文信息，这些信息允许kernel隐式利用图像group之间的关系。为了从上诉计算出来的$\tilde{K}$中计算出一组新的kernels $K_i$，作者这里采用了Multi-Head Self-Attention+Feed-Forward Neural Network的形式。也就是用 $MSA-MLP$的形式来输出一组新的$K_i$（其实这里有点类似于Transformer结构）。$K_i$将用来计算新的Mask：$M_i=g_i(K_i)\ast F$， 这里的 $g_i$ 为 $FC-LN-ReLU$操作。

## K-Net用于分割 ### 训练Instance Kernels的一些细节 1. 这一部分的损失函数选择了Dice Loss，原因是局部细节使用交叉熵的话容易导致不平衡问题。 2. 使用匈牙利算法（Mask-based Hungarian Assignment）来处理Mask配对问题。

全景分割

图3 K-Net用于全景分割

介绍完Kernel Update Head之后，K-Net的结构就比较简单了。通过一个BackBone和Neck（作者这里使用了FPN）来生成一组特征图F。由于语义分割和实例分割所要求的特征图有所差异，所以作者通过两个独立的卷积分支对F进行处理生成$F^{ins}$和$F^{seg}$，当然，这里使用的卷积核为初始化的$K_0^{ins}$和$K_0^{seg}$，这样就生成了一组新的Mask：$M_0^{ins}$、$M_0^{seg}$。

由于$M_0^{seg}$中自然而然的包括了全景分割中所需求的"things"和"stuff"（只不过没有区分实例），那么只需要从$M_0^{seg}$中将包含"stuff"的部分提取出来，再和$M_0^{ins}$（区分实例）直接进行通道相加，即可得到全景分割所需要的Mask：$M_0$。同理，对于卷积核Kernel，也只需要提取对应的$K_0^{ins}$和$K_0^{seg}$组合成新的$K_0$。而对于特征图F，将$F^{ins}$和$F^{seg}$简单通道相加即可（毕竟信息越多越好）。对于新得到的$K_0$、$M_0$和F，经过S次Kernel Update Head处理，可以得到最终的Output：$M_S$。

实例分割

与上述全景分割类似，在实例分割中，只需要删除内核和掩模的串联过程即可执行实例分割。在这一步并不需要删除语义分割分支，因为语义信息仍然是互补的（还是遵循信息越多越好的原则，毕竟实例信息可以从语义信息中提取）。不过，需要注意的是，在这种情况下，语义分割分支不使用额外的真实标签。语义结果的标签是通过将实例掩码转换为相应的类标签来构建的（这里其实在一定程度上起了一个监督的作用，也就是使用语义分支来监督实例分支）。

语义分割

对于语义分割的流程而言，结构是最为简单的，只需要简单地将Kernel Update Head附加到任何依赖语义内核的语义分割方法中即可执行语义分割，比如ASPP、UperNet、FCN、PPM均可。

总结

本文介绍了：

1. K-Net的实现原理。
2. K-Net如何根据不同的语义信息更新核Kernels。
3. K-Net如何应用于语义分割、实例分割、全景分割。

除此之外，K-Net的实现代码可以参考mmcv官方的mmsegmentation和mmdetection，分别可实现语义分割和实例/全景分割。当然作者也给出了他的代码仓库，不过还得配合mmcv食用，K-Net代码链接。