K-Net：统一语义分割、实例分割和全景分割

[189****6672]

K-Net：统一语义分割、实例分割和全景分割

论文标题：K-Net: Towards Unified Image Segmentation

论文地址：https://arxiv.org/pdf/2106.14855.pdf

代码地址：https://github.com/ZwwWayne/K-Net/

摘要

1. 引入主题： 尽管语义、实例和全景分割任务之间存在潜在的联系，但目前使用不同的专门框架进行了分别处理。
2. 创新点： 该论文为这些基本相似的任务提供了一个统一、简单、有效的框架。该框架名为K-Net，通过一组可学习的内核对实例和语义类别进行一致性分割。
3. 实验结果： K-Net超过了之前发布的所有最先进的单一模型结果，在MS COCO test dev split上的全景分割和ADE20K val split上的语义分割任务中，分别实现55.2%PQ和54.3%mIoU。在MS COCO任务中，实例分割性能与Cascade MasR-CNN一致，但推理速度提高了60%～90%。

介绍

所有分割任务基本上都会将每个像素分配给预定义的组别group。

由于图像中的组别通常是有限的，所以可以将分割任务的最大组数设置为N。例如，语义分割有N个预定义的语义类，或者实例分割，图像中最多有N个对象。对于全景分割，N是图像中的类和对象的总数。

因此，可以使用N个内核kernels将图像划分为N个组，每个内核负责找到属于其相应组的像素。

具体来说，给定B张图像的神经网络输出特征图$F\in R^{B\times C\times H\times W}$，只需要N个内核$K\in R^{N\times C}$对其进行卷积操作得到分割结果$M\in R^{B\times N\times H\times W}$。

语义分割领域已经广泛使用上述公式进行像素分类。在实例分割中，每个像素组对应一个对象。然而，以前的方法是通过额外的步骤而不是上述内核来分割实例。

该论文首次探讨了语义分割中的内核概念是否同样适用于实例分割，以及更普遍的全景分割。

为了通过内核的思路来进行实例分割，K-Net中的每个内核最多只能分割图像中的一个对象（图1b）。通过这种方式，K-Net可以区分实例并同时执行分割，一次完成实例分割而无需额外步骤。为了简单起见，将这些内核分别称为语义核和实例核，并用于语义和实例分割。实例核和语义核的简单组合可以自然地进行全景分割，即将一个像素分配给实例ID或类别（图1c）。

算法

尽管K-Net很简单，但直接用内核分割实例并不容易。因为实例内核没有语义范畴中明确的特征，需要比静态核有更强的区分能力。

为了克服这个问题，设计出一个迭代式模型（如图3所示），因为更精细的划分通常会减少组别特征中的噪声，这会有利于生成更具辨别力的内核。

模型的大致流程如下：

1. 一组可学习的内核$K_0$首先与特征映射$F$进行卷积，以预测掩码$M_0$。
2. 然后，内核更新头$f_1$将掩码预测$M_0$、学习的内核$K_0$和特征映射$F$作为输入，并生成类预测、组感知（动态）内核$K_1$和掩码预测$M_1$。
3. 生成的掩码预测、动态内核和特征映射被发送到下一个内核更新头。该过程以迭代方式执行，以逐步细化内核和掩码预测。

其中最核心的就是内核更新头模块，如下图所示：

其主要可以分为以下几个模块：

1. Group Feature Assembling
2. Adaptive Feature Update
3. Kernel Interaction

接下来进行详细介绍

Group Feature Assembling

每个内核掩码$M_{i-1}\in R^{N\times H\times W}$表示每一个像素所属于的组别，这一步的主要目的是汇聚相同组别像素点的特征，可以通过以下公式计算：

这样就能得到每个组别的新特征

Adaptive Feature Update

然后，使用获得的$F^K$来更新内核，以提高内核的表示能力。

每个组别的特征可能包含由其他组别像素引入的噪声。为了减少组别特征中噪声的不利影响，设计了一种自适应核更新策略。

具体来说，首先在$F^K$和$K_{i-1}$之间进行元素相乘：

其中，${\varphi }_1$和${\varphi }_2$是线性层。

然后学习两个门控，$G^F$和$G^K$来控制新组别特征$F^K$和旧组别特征$K_{i-1}$的权重来生成新内核$\tilde{K}$：

其中，${\psi }_n,n=1,\dots ,4$是全卷积层+LN，$\sigma$表示sigmoid函数。

Kernel Interaction

内核之间的交互对于向每个内核提供来自其他组的上下文信息非常重要。

这样的信息允许内核隐式地建模和利用图像组之间的关系。为此，添加了一个内核交互过程，以在给定更新的内核$\tilde{K}$的情况下获得新的内核$K_i$。论文简单地采用了多头注意力机制，然后采用前馈神经网络。

最后，内核交互的输出$K_i$通过$M_i=g_i\left(K_i\right)\ast F$生成一个新的掩码预测，其中$g_i$是一个FC-LN-ReLU层。

$K_i$还将用于预测实例和全景分割中的分类分数。

训练实例分割

虽然每个语义内核都可以分配给一个固定的语义类，但缺少一个明确的规则来为实例内核分配不同数量的目标。

论文采用二部匹配策略（bipartite matching），并设置预测损失，以端到端的方式训练实例内核。与之前依赖于bbox的论文不同，实例内核的学习纯粹是掩码驱动的。

损失函数

实例内核的损失函数为$L_K={\lambda }_{cls}{L}_{cls}+{\lambda }_{ce}{L}_{ce}+$${\lambda }_{\text{dice }}{L}_{\text{dice }}$。

其中$L_{cls}$是用于分类的Focal loss ，$L_{ce}$和$L_{dice}$是用于分割的CrossEntropy loss和Dice loss

基于mask的匈牙利匹配算法

采用匈牙利匹配算法进行目标匹配，以端到端的方式训练K-Net。

它基于匹配损失在预测实例掩码和真值（GT）实例之间建立一对一映射。

实验

Benchmark Results

全景分割实验

实例分割实验

语义分割实验

消融实验

可视化