Adaptive Prototype Learning and Allocation for Few-Shot Segmentation

自适应原型学习和分配用于小样本分割
作者来自爱丁堡大学、谷歌研究院、成均馆大学

Abstract

原型学习被广泛用于小样本分割。通常，通过平均全局对象信息，从支持特征获得单个原型。然而，使用一个原型来表示所有的信息可能会导致歧义。在本文中，我们提出了两个新的模块，超像素引导聚类(SGC)和引导原型分配(GPA)，用于多原型提取和分配。具体来说，SGC是一种无参数(自由)和无训练（自由）的方法，它通过聚集相似的特征向量来提取更具代表性的原型，而GPA能够选择匹配的原型来提供更精确的指导。通过将SGC和GPA结合在一起，我们提出了自适应超像素引导网络(ASGNet)，这是一个轻量级模型，并适应对象的规模和形状变化。此外，我们的网络可以很容易地推广到k-shot分割，并有实质性的改进，没有额外的计算成本。特别是，我们对COCO的评估表明，ASGNet在5-shot分割中比最先进的方法高出5%。

Introduction

人类有一种非凡的能力，在只看到少量样本后，就能学会如何识别新奇的物体。另一方面，基于深度学习的计算机视觉系统已经取得了巨大的进步，但是在很大程度上依赖于大规模的训练集。此外，深度网络主要使用预定义的类，不能推广到新的类。小样本学习领域研究机器如何学习人的这种能力，即只有几个新类别的例子可用的情况下，也能识别新类。
在这项工作中，我们解决了小样本分割问题，对少数有标注信息支持图像进行学习，来分割给定的查询图像中的目标。这个任务具有挑战性，因为测试数据是训练集中不存在的新类别，并且在支持图像和查询图像之间通常存在外观和形状较大的不同(同一类也存在差异)。目前的小样本分割网络通常从查询图像和支持图像中提取特征，然后提出不同的特征匹配方法和从支持图像到查询图像的对象掩码转换方法。这种特征匹配和掩模转移通常以两种方式之一执行:原型特征学习或相似性学习。原型学习技术将支持图像中被掩蔽的对象特征浓缩成单个或几个原型特征向量。然后，这些技术在查询图像中找到相似特征的像素位置，以分割期望的对象。原型学习的一个关键优势是原型特征比像素特征更能抵抗噪声。然而，当支持图像和查询图像之间的对象外观存在较大差异时，原型特征不可避免地会丢失空间信息。另外，大多数的原型学习仅通过图一a展示的masked average pooling生成一个原型，丢失一些信息和差异性（可辨别性）。
另一方面，相似性学习技术直接尝试支持图像中的对象像素来匹配查询图像像素，从而转移对象掩模。这些技术使用学习的特征来预测跨图像像素的相似性(也称为连接强度)，与原型学习方法相比，该方法在执行特征匹配的同时更好地保留空间信息。然而，相似性学习技术在试图解决under-constrained的密集相似矩阵像素匹配问题时，所以训练数据容易过拟合。
在这项工作中，我们提出了一种新的原型学习技术，解决了现有技术的一些主要缺点。特别是，我们希望基于图像内容自适应地改变原型的数量及其空间范围，使原型具有内容适应性和空间感知性。这种自适应的多原型策略对于处理不同图像中对象比例和形状的巨大变化非常重要。直观地说，当一个物体占据了图像的很大一部分时，它携带了更多的信息，因此需要更多的原型来表示所有必要的信息。相反，如果物体相当小，背景的比例很大，那么一个或几个原型就足够了。此外，我们希望每个原型的支持区域(空间范围)能够适应支持图像中存在的对象信息。具体来说，我们旨在根据特征相似性将支持特征划分为几个有代表性的区域。我们还希望在定义查询图像中的相似特征时，自适应地选择更重要的prototype。由于不同的对象部分在不同的图像区域和不同的查询图像中是可见的，我们希望跨查询图像动态地分配不同的原型用于特征匹配。例如，对象的某些部分可能在查询图像中被遮挡，我们希望动态选择与查询图像中的可见部分相对应的原型。
我们通过自适应超像素引导网络(ASGNet)实现了这种自适应的多原型学习和分配，该网络利用超像素来调整原型的数量和支持区域。在图1(b)中进行展示。特别地，我们提出了两个模块，超像素引导聚类(SGC)和引导原型分配(GPA)，它们构成了ASGNet的核心。SGC模块对支持图像进行快速的基于特征的超像素提取，得到的超像素质心被认为是原型特征。由于超像素的形状和数量与图像内容相适应，因此生成的原型也变得自适应。GPA模块使用类似注意力的机制来分配最相关的support原型feature给query 图像中的每个像素。总之，SGC模块在原型数量及其空间范围方面提供自适应原型学习，而GPA模块在处理查询特征时提供所学原型的自适应分配。这两个模块使ASGNet高度灵活地适应变化的对象形状和大小，允许它更好地推广到看不见的对象类别。我们做出以下贡献:
 我们提出了自适应超像素引导网络(ASGNet)，这是一种灵活的原型学习方法，适用于不同的目标尺度、形状和遮挡。
 我们引入了两个新颖的模块，即超像素引导聚类(SGC)和引导原型分配(GPA)，分别用于自适应原型提取和分配。它们可以作为特征匹配的组件来使用。
 ASGNet以更少的参数和更少的计算获得了最高性能的结果。具体而言，该方法在Pascal-5i/COCO-20i 5-shot分别中获得了64.36%/42.48%的mIoUs，超过了现有技术的2.40%/5.08%。

Related Work

Semantic Segmentation（语义分割）

现有的语义分割方法大多基于全卷积网络(FCNs)，用全卷积网络代替全连通层进行像素级预测。最近语义分割的突破主要来自多尺度特征聚合或注意机制。这些方法通常使用扩张卷积核，并建立一个编码器-解码器结构，以获得大的感受野，同时保持特征分辨率。虽然这些方法取得了巨大的成功，但它们需要很长的训练时间和大量像素级标记的地面真实来充分监督网络。此外，在推理阶段，经过训练的模型无法识别训练集中不存在的新类。

Few-shot Learning（小样本学习）

小样本学习侧重于模型的泛化能力，这样他们只需给出几个带注释的例子，就能学会预测新的类。现有方法主要集中在度量学习和元学习。度量学习的核心思想是距离度量，它通常被表述为图像或区域之间的距离/相似度的优化。在元学习方法中，主要思想是定义特定的优化或损失函数，以实现快速学习能力。在这些方法中，原型网络的概念被广泛用于小样本分割，在保持高性能的同时大大降低了计算预算。大多数方法集中在图像分类上，而近年来小样本分割受到越来越多的关注。

Few-shot Segmentation（小样本分割）

小样本分割是小样本分类的扩展，它解决了一个更具挑战性的任务，即预测每个像素的标签，而不是预测整个图像的单个标签。这个研究问题是由沙班等人提出的，他提出了一个经典的两分支网络。后来，PL引入了使用原型的想法。在这项工作中，每个预测都是通过测量查询图像中原型和像素之间的相似度来生成的。SGOne提出了掩蔽平均池来获得与对象相关的原型，并在后续工作中得到了广泛应用。PANet 引入了一种新的原型对准正则化，以充分利用支持知识，仅通过余弦距离的测量进行最终预测。基于屏蔽平均池，CANet将原型扩展到相同大小的query feature，并将它们连接在一起。这项工作还使用了迭代优化模块来细化分割结果。PGNet、BriNet和DAN 在支持和查询特征之间引入了密集的像素到像素连接，以维护空间信息。最近，PMMs 利用期望最大化算法生成了多个原型。然而，所有的原型在这个模型中都有相同的相关性，这可能对不匹配的原型很敏感。相反，在我们的工作中，我们利用每个原型和查询特征之间的相似性，在每个像素位置选择最相关的原型。

Superpixel Segmentation（超像素分割）

超像素被定义为具有相似特征(颜色、纹理、类别)的一组像素。超像素在许多计算机视觉任务中是有效的，并且最近被用作小样本分割的基本单元。超像素比像素承载更多的信息，可以为下游视觉任务提供更紧凑、更方便的图像表示。
我们的研究受到maskSLIC和超像素采样网络(SSN) 的启发。MaskSLIC使SLIC适应定义的感兴趣区域(RoI)，主要贡献在于RoI内种子点的放置。SSN通过使SLIC算法可微，提出了第一个端到端可训练超像素算法。受这两种技术的启发，我们提出了特征空间中的掩蔽超像素聚类，它可以将相似的特征聚集在一起，生成超像素质心作为原型。超像素不是代表整个物体的信息，而是代表物体具有相似特征的部分。

Problem Definition

小样本分割和一般语义分割的关键区别在于训练集和测试集中的类别不相交。这意味着，在推理阶段，测试集有训练中完全看不到的类。具体的说，给一个训练集和测试集，训练集和测试集类别的交集为空。I为RGB图像，M为分割掩码。S为support，Q为query。随着第一个提出的one-shot分割工作，我们的训练和测试都采用episode，每一个episode，输入模型的是一个query 图像，和support set中k个样本。Support 和query有相同的类别。我们对预测的query mask和真实的mask进行估计。

Proposed Method

在本节中，我们首先介绍了原型生成和匹配的两个模块，即超像素引导聚类(SGC)和引导原型分配(GPA)。然后，我们讨论了这两个模块的自适应能力。之后，我们介绍了整体网络架构，称为自适应超像素引导网络(ASGNet)，它将SGC和GPA模块集成在一个模型中。整体结构如图2所示。最后，我们详细介绍了ASGNet中的k-shot设置。

Superpixel-guided Clustering（超像素引导聚类）

SGC的核心思想受到超像素采样网络(SSN) 和MaskSLIC 的启发。SSN是第一个用于超像素细分的端到端可训练深度网络。SSN的主要贡献是将SLIC 中的最近邻运算转化为可微运算。传统的SLIC超像素算法采用k均值聚类迭代，分像素-超像素关联和超像素质心更新两个步骤。基于颜色相似性和接近性，像素被分配给不同的超像素质心。具体来说，输入图像I∈Rn×5通常在具有n个像素的五维空间(labxy)中，其中lab表示CIELAB颜色空间中的像素向量，xy表示像素位置。在迭代聚类之后，算法输出关联图，其中每个像素n被分配给m个超像素之一。
这个简单的方法启发了我们一个很有见地的想法，就是将特征图聚类聚合成多个超像素质心，这里超像素质心可以作为原型。因此，我们不是计算图像空间中的超像素质心，而是在特征空间中通过聚类相似的特征向量对它们进行分类。完整的SGC见算法1.
给定支持特征Fs∈ Rc×h×w，支持掩码Ms∈ Rh×w初始超像素种子S0∈ Rc×Nsp (Nsp为超像素个数)，我们的目标是获得最终超像素质心，它构成了多个超像素原型。首先，我们将每个像素的坐标与支持特征图连接起来，以引入位置信息。然后，我们根据SLIC 定义距离函数D，
它由特征和空间距离组成: df、ds是特征和坐标值的欧几里德距离，r是加权因子。我们用支持掩码过滤掉背景信息，只保留支持掩码的特征，从Fs∈ Rc×h×w到F‘s∈ Rc×Nm压缩特征图，其中Nm是支持掩码内的像素数。
然后我们以迭代的方式计算基于超像素的原型。对于每次迭代t，我们首先根据距离函数D计算每个像素p和所有超像素之间的关联图Qt

然后，新的超像素质心被更新为masked特征的加权和:

， 是归一化常数。上述过程图3所示。
在这里，我们详细说明初始种子的选择。一般在超像素算法中，一个H×W图像被均匀划分为大小为H×W的规则网格单元，每个网格单元被认为是一个初始种子(即超像素)。然而，这种初始化并不适合我们的目的，在我们的目的中，为支持图像提供了一个前景遮罩，我们只需要初始化这个前景区域内的种子。为了统一初始化屏蔽区域中的种子，我们参考MaskSLIC 来迭代放置每个初始种子，通道如图5所示。这种种子初始化只需几次迭代就能使超像素引导的聚类更快收敛。

Guided Prototype Allocation

在提取原型之后，以前的方法大多遵循CANet 的设计，将单个原型扩展到相同大小的查询特征，并将它们连接在一起。但是，此操作会使query feature中每个位置都有相同的原型进行匹配。为了使原型匹配更加适应查询图像内容，我们提出了引导原型分配(GPA)，如图4所示。我们首先计算余弦距离来度量每个原型和每个查询特征元素之间的相似性:

是第i个超像素质心（原型）
， 是query feature中（x，y）位置上的特征向量。我们使用这个相似性信息作为两个分支结构的输入。第一个分支计算哪个原型在每个像素位置最相似，如下所示:
，其中采用argmax运算符来获得Gx，y这是一个表示特定原型的单个索引值。将所有索引值放在一起，我们得到一个引导图G ∈ Rh×w，然后，通过在引导图的每个位置放置相应的原型，我们得到引导特征FG∈ Rc×h×w，实现像素级引导。而在另一个分支中，所有的相似性信息C在所有的超像素上相加，得到概率图P。
最后，我们将概率图和带有原始查询特征FQ的引导特征连接起来，提供引导信息，从而得到细化的查询特征F‘ Q:
⊕表示沿通道维度的串联操作，f(.)是1 × 1卷积。

Adaptability（自适应性）

如前所述，我们认为所提出的网络的关键属性之一是它对小样本语义分割的自适应能力。在图6中，我们提供了一些例子来说明SGC和GPA的适应能力。在SGC，为了使其适应物体的尺度，我们定义了一个标准来调节超像素质心的数量为
，其中Nm是支持掩码中的像素数；Ssp是分配给每个初始超像素种子的平均面积，我们根据经验设置为100。当前景相当小时，Nsp= 0或1，这种方法退化为图6(a)所示的一般masked平均池（一个原型）。此外，为了减少计算负担，我们设置了一个超参数Nmax来约束原型的最大数量。在GPA中，我们可以观察到它对物体形状的适应性。换句话说，它对遮挡物体适应性。当查询图像中存在严重遮挡时，例如图6(b)，GPA可以为每个查询特征位置选择最佳匹配的原型。

Adaptive Superpixel-guided Network（ASGNet）

基于上述SGC和GPA模块，我们提出了自适应超像素引导网络(ASGNet)用于小样本语义分割，如图2所示。首先，支持和查询图像被送入一个共享的CNN(在ImageNet 上预处理)以提取特征。然后，通过用support mask和support feature通过SGC，我们获得超像素质心，将其视为原型。之后，为了更精确的像素引导，我们采用GPA模块将原型与查询特征进行匹配。最后，我们使用特征丰富模块并建立一个类似FPN自顶向下的结构来引入多尺度信息。从最后，将所有不同的尺度连接起来，每个尺度产生一个用于计算损失的分割结果。特征从精细到粗糙的转换促进了特征交互，因此我们遵循它们的设计进行快速多尺度聚合。

K-shot Setting

在以前的工作中，k-shot设置通常通过特征平均或基于注意力的融合来解决。然而，事实证明，这种策略的改进很小，同时需要大量的计算。相比之下，基于所提出的SGC和GPA，我们的ASGNet可以轻松地采用高效的k-shot策略，而不会破坏support features。具体来说，在每个support image和mask对中，我们实现SGC来获得超像素质心。通过将所有超像素质心收集在一起，我们可以从k shots中获得整个超像素质心S = (S1，S2，…，Sk)，其中超像素质心的数量是Nsp =
通过这样做，GPA可以从多个shots中接收更大范围的选择，从而产生更准确的指导来分割查询图像中的对象。

Experiments

Datasets and Evaluation Metric

我们选择Pascal-5i和 COCO-20i这两个广泛用于小样本语义分割的数据集来分析我们的模型性能。Pascal-5i包括来自PASCAL VOC 2012 的图像和来自SBD的额外注释。总共20个类别被平均划分为四个部分，模型训练以交叉验证的方式进行。具体来说，在训练过程中选择三个spit，剩下的一个用于测试。在推理过程中，随机抽样1000个support-query对进行评估。与Pascal5i不同，MSCOCO是一个大型数据集，训练集中有82，081幅图像。在FWB 之后，来自MSCOCO 的全部80个类被平均分成四个部分，采用相同的交叉验证策略。为了获得更稳定的结果，我们在评估期间随机抽取了20，000对样本。
我们使用平均IoU(mIoU)作为消融研究的主要评估指标，这在图像分割中是公认的。此外，为了更一致的比较，还展示了前景背景的IoU (FB-IoU)的结果。

Implementation Details（实施细节）

我们采用ResNet作为骨干网络，并且我们连接区块2和区块3以生成特征图。我们用SGD优化器模型在Pascal-5i 200个epochs和COCO-20i 50 epochs的训练。我们将Pascal-5i上batch-size为4的初始学习率设置为0.0025，COCO-20i上batch-size为8的初始学习率设置为0.005。SGC模块中的迭代次数在训练期间设置为10，在推理期间设置为5。我们在训练过程中使用数据增强——输入图像从[-10，10]开始进行随机缩放、水平翻转和旋转变换，然后将所有图像裁剪为473 × 473(Pascal)或641 × 641(COCO)作为训练样本。我们使用Pytorch实现我们的模型，并在带有英伟达特斯拉V100 GPU的工作站上运行我们的实验。为了增加余弦度量的方差，我们在支持和查询特征之前移除了ReLU层，因此相似性度量被限制在[-1，1]而不是[0，1]。

Ablation Study（消融实验）

为了验证所提出模块的有效性，我们在Pascal-5i上使用ResNet-50主干进行了广泛的消融研究。我们使用浮点运算（(FLOPs)来表示计算量，并考虑加法和乘法。
Number of Superpixel Centroids（超像素质心-原型数量）：为了探索超像素质心数量的影响，我们进行了不同Nmaxi在 1-shot上分割的实验。如第4.3节所述，Nmax有一个超参数集来调节原型的最大数量。特别是，当Nmax= 1时，原型生成过程退化为masked平均池（只有一个原型）。如表1所示，当Nmax等于5时，ASGNet在split0和2中获得最佳结果，以及最佳平均性能，这证明了超像素引导聚类的有效性。当数值从1增加到5时，结果逐渐得到改善，然后稍有下降，表示原型过多带来负面影响，容易过度拟合。最后，我们选择Nmax=5进行1-shot和5-shot分割。此外，我们分析了所提出的自适应方案对超像素质心数量Nsp的影响，结果如表2所示。与使用固定数量(5)的超像素相比，结果显示自适应设置(公式7)

可以在获得更好性能的同时减少冗余计算。
SGC和GPA。为了证明所提出的SGC和GPA模块的有效性，我们进行了原型生成和匹配的诊断实验。我们首先用PFENet中提出的单原型学习实现一个基线模型。然后，我们引入SGC生成多个原型，并参考PMMs以密集的方式融合它们。最后，我们用建议的分配方案替换原型扩展。对于SGC模块上的消融，如表3所示，在1-shot设置中，我们观察到用SGC替换平均池(MAP)导致性能下降，但模型在5-shot场景中受益。我们认为主要原因是过多的原型在单个支持样本上变得高度相似，余弦距离不能区分它们。最后，采用GPA模块进行原型匹配时，性能比原型扩展提高了2.70%，计算开销也大大降低。
K-shot融合设置。如第4.5节所述，ASGNet能够在不破坏支持特性的情况下解决k-shot学习问题。为了探索不同融合方法的效果，我们将我们的k-shot设置与其他两种常用解决方案进行了比较:1)特征平均融合和2)注意力加权求和。为简单起见，本消融研究的实验仅用Pascal-50进行。如表4所示，我们的简单策略在没有额外计算的情况下获得了最佳性能和超过1-shot基线(4.82%)的最大增量。相反，基于注意力的融合需要大量的计算，并且性能提升有限。这表明，当给出大量选择时，GPA模块是非常有效的。

Comparison to State-of-the-art（与最先进的技术进行比较）

Pascal-5i. 表5和表6显示了与最先进方法的比较，其中采用了两种不同的指标。在表5中，以ResNet-101为主干，ASGNet在5-shot分割中以2.40%大的差距优于最近的方法，同时与1-shot设置下的最先进水平不相上下。在表6中，我们进一步比较了FB-IoU和可训练参数的数量。再次，所提出的ASGNet在5-shot设置中实现了对最先进技术的显著改进(75.2%对73.9%)，并且它还比1-shot结果有5.0%的最大性能增加。此外，我们的可训练参数比其他方法少得多。在图7中，我们展示了一些代表性的细分示例。我们观察到，即使在支持图像和查询图像之间存在较大的外观和姿态变化时，所提出的ASGNet也能生成准确的分割结果。
COCO20i 在表7中，我们给出了COCO-20i上平均IoU和FB-IoU的性能比较。可以看出，就平均IoU而言，ASGNet在1-shot和5-shot设置中都获得了最先进的结果，并且它大大优于以前的方法。特别是，ASGNet实现了比RPMMs 1-shot和5-shot分别高3.98%和6.96%。此外，我们的ASGNet在FB-IoU方面获得了极具竞争力的1-shot结果和表现最佳的5-shot结果。这些结果表明，所提出的方法能够处理更复杂的情况，因为MSCOCO具有不同的样本和类别，是一个更具挑战性的数据集。

Conclusion

本文提出了一种用于小样本图像分割的ASGNet算法。针对现有基于单个原型的模型的局限性，我们引入了两个新的模块，称为超像素引导聚类(SGC)和引导原型分配(GPA)，用于自适应原型学习和分配。具体来说，SGC通过基于特征的超像素聚类来聚合相似特征向量，而GPA旨在通过余弦距离来度量相似度，从而为每个查询特征元素分配最相关的原型。广泛的实验和消融研究证明了ASGNet的优越性，我们在Pascal-5i和COCO-20i上实现了最先进的性能，没有任何额外的后处理步骤。