[小样本图像分割]PANet: Few-Shot Image Semantic Segmentation with Prototype Alignment

PANet:基于原型对齐的Few Shot图像语义分割（ICCV19）

论文地址
开源代码

摘要

本文从度量学习的角度来解决Few Shot Segmentation问题，提出一种新的原型对齐网络来更好地利用支持集信息。PANet从嵌入空间内的一些支持图像中学习特定类的原型表示，然后通过将每个像素与学习到的原型进行匹配，对查询图像进行分割。通过非参数度量学习，PANet提供了代表每个语义类的高质量原型，这些原型同时还存在对不同类的判别信息。此外，PANet还引入了一种支持和查询之间的原型对齐正则化。在此基础上，PANet充分利用了来自支持的知识，在少镜头分割中提供了更好的泛化。

目前存在的问题和方案

现有的少镜头分割方法通常从少量的支持图像中学习，然后将学习到的知识输入到参数化模块中进行查询分割。然而，这类方案存在两个缺点，推广效果不理想。首先，它们没有区分知识提取和分割过程，这可能会有问题，因为分割模型的表示与支持集的语义特征混合。因此，我们建议将这部分分为原型提取和非参数度量学习两部分。原型被优化为每个语义类的紧凑和鲁棒表示，非参数度量学习通过嵌入空间中的像素级匹配进行分割。此外，我们不像以前的方法一样使只使用支持图像的标注做masking，我们将这些标注加入了小样本学习的监督过程。

为此，我们引入了一种新的原型对齐正则化方法，这是一种反向分割方法。即将查询图像及其预测掩码作为一种新的支持集，用于对之前的支持图像进行分割。通过这种方式，鼓励模型在支持和查询之间生成更一致的原型，从而提供更好的泛化性能

PANet首先通过共享的特征提取器将不同的前景对象和背景嵌入到不同的原型中。这样，每一个学习到的原型都是对应类的代表，同时也有足够区别于其它类。然后，将查询图像中的每个像素标记为这个像素的嵌入表示离它最近的类特定原型。

所提出的PANet结构设计有几个优点。首先，它不引入额外的可学习参数，因此不容易过度拟合。其次，在PANet中，原型嵌入和预测在计算的特征图上执行，因此分割不需要额外通过网络。此外，由于正则化仅用于训练，因此不会增加推理的计算成本。

方法

现有的少镜头分割方法将提取的支持特征与查询特征融合，以参数化的方式生成分割结果，我们提出的模型旨在在嵌入空间中学习并对齐每个语义类的紧凑和健壮的原型表示，然后通过非参数度量学习在嵌入空间内进行分割。

原型学习（Prototype learning）

我们的模型基于度量网络，能够有效地学习每个语义类易于分离的原型表示。PANet不是对整个输入图像进行平均，而是利用支持图像上的掩码注释来分别学习前景和背景原型。有两种方法来利用分割掩码，前期融合和后期融合，本文采取的是后期融合，也就是将mask盖在特征图上以产生前景和背景特征（而不是在输入特征提取器前进行融合），这样做有利于保持共享特征提取器的输入一致性。

具体来说，给定一个支持集 S i = { ( I c , k , M c , k ) } S_i=\left\{ \left( I_{c,k},M_{c,k} \right) \right\} Si={(Ic,k,Mc,k)}， F c , k F_{c,k} Fc,k表示图像 I c , k I_{c,k} Ic,k经过特征提取后得到的特征图。 c c c表示类别， k k k则表示支持图像序号。类 c c c的原型通过掩码平均池化得到：

背景类的原型则可表示为：

( x , y ) (x,y) (x,y)表示位置， 1 \mathbb{1} 1表示框内式子为真则值为1，假为0

非参数度量学习（Non-parametric metric learning）

我们采用非参数度量学习方法来学习最优的原型并进行相应的分割。由于分割可以看作是在每个空间位置的分类，我们计算查询图像在每个空间位置上的特征向量与支持图像原型之间的距离。然后在距离上使用softmax得到语义类（包括背景）的概率映射 M q ^ \hat{M_q} Mq^

具体来说，给定一个距离函数 d d d，让 P = { p c ∣ c ∈ C i } ∪ { p b g } \mathcal{P}=\left\{ p_c|c\in \mathcal{C}_i \right\} \cup \left\{ p_{bg} \right\} P={pc∣c∈Ci}∪{pbg}， F q F_q Fq查询图像特征映射。对于每个 p j ∈ P p_j\in\mathcal{P} pj∈P，我们有：

然后给出了预测的分割掩码：

距离函数 d d d一般采用余弦距离或平方欧几里得距离，根据经验，我们发现使用余弦距离更稳定，性能更好，可能是因为它是有界的，因此更容易优化。

分割损失的表达如下：

M q M_q Mq表示查询图像的ground-truth分割掩码，N表示所有空间位置数量，对上述损失进行优化将为每个类派生出合适的原型。

原型对齐正则化（Prototype alignment regularization）

直觉上来说，如果模型能够使用从支持集中提取的原型来预测出一个好的查询分割掩码，那么这个预测出的查询图像分割掩码应该能够很好地分割支持图像。因此PAR鼓励最终的分割模型反向进行few-shot learning，即以查询和预测的掩码作为新的支持样本来学习分割原始支持图像（也就是这里的查询图像和预测得到的查询图像掩码作为新的支持图像–掩码对）。这将在支持原型和查询图像之间形成对齐，并从支持集中学习到更丰富的知识。

使用查询图像和其掩码预测支持图像掩码的损失函数如下：

在没有PAR的情况下，信息只能单向地从支持集流向查询集。通过将信息流回支持集，我们迫使模型学习一个一致的内嵌空间，以对齐查询和支持原型。

那么PANet的总损失可以表示为：

实验结果

结论

提出了一种新的基于度量学习的多镜头网络分割算法。PANet能够从支持集中提取鲁棒原型，并使用非参数距离计算进行分割。在此基础上，我们的模型可以进一步利用支持信息来辅助训练。在没有任何解码器结构或后处理步骤的情况下，我们的PANet大大优于之前的工作。