1. 引言
语义分割作为计算机视觉的基本任务之一,广泛应用于医学图像识别、自动驾驶、地质勘探等领域 [1] [2] 。基于卷积神经网络的语义分割模型需要大量具有像素级别标注的数据,然而获取大量的、高质量的像素级别标注数据费时费力 [3] [4] [5] [6] [7] ,限制了语义分割在很多领域的应用。为了缓解语义分割模型需要大量标注数据的压力,使用少量的带标注样本完成语义分割任务的小样本语义分割问题被提出并受到广泛关注 [8] 。小样本语义分割,也称小样本分割(Few-shot semantic segmentation, FSS)。小样本分割主要采用基于度量的元学习方法进行 [9] [10] 。利用度量的思想,使用度量工具(例如,余弦相似度、卷积等)度量少量带标签样本特征和待分割图片特征在高维空间的距离,对待分割图片中的每个像素点做出预测。
受小样本分类任务的推动,大多数小样本语义分割任务采取了元学习来实现模型的泛化。这些方法在基础数据集中采样一系列学习任务,模拟包含新类别的小样本场景来进行训练。然而这样训练的结果往往并不够充分的。当在基础数据集上进行元训练时,使用丰富的注释样本不可避免地引入了对已知类别的偏见从而妨碍了模型对新类别的学习。所以一些研究人员尝试调整包含足够训练样本的基础数据集的使用 [11] 。他们除了使用小样本分割常见的元学习器外还引入了一个额外分支即基学习器,他们利用基学习器去预测基础类别的目标物体,通过自适应集成这两个学习器并行输出的粗略结果去生成准确的预测结果。这种操作可以帮助模型在传统范式的基础上更好地分割查询集中的易混淆区域。这种方法被称作基学习与元学习结合的小样本分割方法(Base and Meta) [12] ,后文简称BAM。此外在目前的小样本语义分割方法中,文献 [13] [14] [15] 等方法都采用原型学习范式进行元推理,通过平均池化操作生成单个或多个原型,为查询对象推理提供了必要的前景线索。然而,当面对具有模糊对象边界和无关对象干扰等因素的困难小样本任务时,它们的分割能力可能会大打折扣。我们认为这些方法失败的一个可能原因在于未充分利用支持图像–掩膜对中的信息,仅依靠压缩的与对象相关的特征来引导分割是远远不够的。甚至在支持集引导的过程中即使查询集图像与支持集图像一样也无法很好地引导分割 [16] 。
为此,我们在BAM方法中引入一种原型特征增强的方法,使模型对目标类别可以有更强的分辨能力。该方法利用基学习器完成对查询集背景中基类区域有效识别,同时原型特征增强策略进一步加强了对原有原型中学习不充分的部分,这使得模型同时对待分割图像的前景和背景完成了特征增强。具体的,我们给出一种自引导的方式,利用支持集生成辅助原型,再通过辅助原型引导查询集生成激活特征图,最终使用基学习器和元学器得到的结果进行有效集成得到更好的分割结果。该方法在Pascal数据集上取得了较为明显的效果。
2. 相关工作
2.1. 小样本学习
计算机视觉领域多年来一直致力于使网络具备对新类别的泛化能力。目前在小样本学习领域中,大多数方法都遵循 [17] 中提出的元学习框架,其中从基础数据集中采样一组学习任务以模拟小样本场景。基于此,小样本学习方法可以进一步分为以下三个类别:第一是基于度量的方法 [17] [18] 该方法使用特定的距离度量(如欧氏距离和余弦距离)确定支持原型与查询特征之间的关联性。第二是基于优化的方法 [19] [20] ,第三是基于特征增强的方法 [21] [22] 。
2.2. 小样本语义分割
为了解决稀疏注释数据集上的泛化问题,将小样本学习方法应用于语义分割领域。度量学习是小样本学习方法之一,在小样本分割语义任务中得到了广泛应用,例如使用原型向量学习知识。一些小样本语义分割方法 [23] [24] 开始探索建立像素级别的特征匹配与引导。他们有选择地从支持图像中提取有用的信息,通过聚合查询和支持特征之间的关系并在这种关系的引导下进行语义分割。然而这种方法往往需要耗费巨大的计算量。另一些小样本分割方法主要基于原型向量 [13] [14] 。它们对支持特征的前景区域进行全局平均池化,得到全局特征向量,并将查询特征的空间位置与这些特征向量进行比较。一些工作 [16] [24] 尝试添加多个包含目标类别的部分感知特征的原型向量,以丰富支持特征表示。
3. 问题设定
小样本语义分割是一个新兴的研究方向,它结合了语义分割和小样本学习的思想。其主要目标是通过使用少量带标签的未知类别来训练网络,使其能够对该类别的图像进行准确的分割。在这个任务仍然需要训练集和测试集对模型进行训练和评估,但是这个任务的训练集与测试集设定与别的深度学习任务中的设定有一定的区别。
具体而言,我们将训练集称为
,测试集称为
分别用于训练和测试模型。但是训练集中作为分割目标的类别
和测试集中作为分割目标的类别
是完全不同的,并且没有交集,即
。除此以外,在小样本语义分割最初的工作OSLSM [8] 中引入一个“episode”范式,也叫片段范式。每一个片段的构成是同一目标类别的支持集图片和查询集图片。每个片段由同一类别的支持集和查询集图片组成。将支持集称作
,
中包含了通过在包含同一类别的N张图片中随机采样得到的K张图片,将其写作
,其中
代表第i张支持集图片。
代表对应支持集的二进制掩码。同样的,将查询集称作
,
包含带掩码
的图片
,即
,而查询集真实的标签掩码
仅在训练阶段使用。
此外,我们将训练集和测试集的支持集与查询集分别记为
和
以及
和
。在训练阶段,我们使用
来训练模型。一旦模型训练完成,在测试阶段,我们不需要对网络进行微调,而是从测试集中随机采样
,输入模型并获得查询集的预测结果。
4. 方法
在本节中,我们将详细介绍如何在BAM中实现原型特征增强。简而言之,我们利用自引导的方式对支持集进行分割,在对分割正确和分割错误的部分进行特征聚合得到表达信息更完整的辅助原型。最后利用辅助原型结合基学习器与元学习器对查询集进行分割,得到更加精细的分割结果。
4.1. BAM简述
当前小样本分割的泛化性能严重依赖元学习,即使进行了微调过程,其能力也不达标。具体来说,训练的小样本语义模型由于数据分布不平衡会从而对基类产生偏见。在这个领域中,很少有研究明确研究泛化性能下降的问题,而是专注于设计两个分支之间的高容量交互模块。然而Lang [12] 为了减轻当前小样本分割方法中的偏见问题,提出建立一个利用额外学习器的方法,明确预测查询图像中基础类别的区域,从而促进新对象的分割。所提出的方法主要由三个组件构成,包括一对互补学习器(即基学习器和元学习器)和一个集成模块。使用共享主干的两个学习器分别用于识别基础类别和不可见类。然后,集成模块接收它们的粗略预测结果和利用浅层特征计算得到的一个调整因子
,以抑制基础类别错误激活的区域,从而进一步产生准确的分割结果,具体如图1所示。我们将该方法分为第一阶段和第二阶段,第一阶段结合基学习器与元学习生成粗糙结果,第二阶段利用浅层特征得到的调整因子对基学习预测结果和元学习预测结果进行集成调整最后得到最终的结果,我们的改进在第一阶段上进行。
Figure 1. Illustration of the combination of base learning and meta-learning for few-shot semantic segmentation
图1. 基学习与元学习结合的小样本语义分割示意图
4.2. 原型特征增强方法概述
如图2所示,该方法结合基学习与元学习方法,在此之上利用自引导机制生成辅助原型,从而感知原本模型中可能会忽略掉的区域,利用辅助原型激活查询集的不同区域生成激活特征图,最后将激活特征图与拓展维度的原型和原本查询集特征聚合,得到增强后特征,利用增强后的特征引导查询集分割。如此一来该方法不仅增强了前景特征,还通过基类学习器增强了模型对背景的感知。
Figure 2. Diagram of the method we proposed
图2. 提出的方法
4.3. 自引导辅助原型生成
原型学习是小样本分割一种重要的方法,这类方法利用原型网络的思想,首先将支持集图像输入骨干网络中得到支持集图像的高维特征,再结合支持集图像的掩膜进行全局平均池化得到原型。该原型表示通过支持集图像计算得到的该目标类在高维空间中的特征中心。同样的,将查询集图像也送入同样权重的骨干网络得到查询集在高维空间的特征表达。最后利用无参度量工具如余弦函数或欧式距离公式衡量高维特征与原型之间的距离完成对图像像素点的分类,得到最终的预测图。然而这样得到的原型往往会损失不少有用的信息。为此,引入自引导模块可以更加有效利用支持集信息,提供更加全面的信息引导查询集进行分割。
按照在问题设定中的定义,支持集图像被设定为
,我们选择基于ResNet50 [26] 拓展得到的PSPNet [2] 作为模型的骨干网络,支持集图片和查询集图片分别输入到骨干网络得到的特征为
。其中s和q分别代表支持集和查询集,m表骨干网络的层数。对于支持集而言,我们聚合骨干网络中得到的第二层和第三层特征作为骨干网络的输出进行下一步处理。具体的,将第三层和第四层特征在通道维度进行拼接得到待处理的特征
。这里的
表示拼接操作
(1)
之后,为了利用支持集的掩膜对所得特征进行处理生成待分割类别的原型,我们需要将掩膜下采样到和特征
一样的大小。有:
(2)
其中
代表支持集原始的掩膜。
代表对原始掩膜下采样后得到的与特征
高度和宽度大小一致的掩膜。完成上述处理,我们遵循之前工作的操作利用掩膜全局平均池化生成初始原型 [3] [4] [5] 。
(3)
这里的i表示特征图中空间位置的索引。h和w分别是特征图的高度和宽度。
是艾弗森括号。它代表中括号中的内置条件如果成立那么这个表达式的值取1,如果不成立这个表达式的值取0。在这里
表示一个二进制掩膜,
表示在第i个位置上的这个像素点是属于待分割的目标类别。
然而,通过这种方式所求得到的原型往往是对待分割类别信息的粗略表达,在掩膜全局平均池化这一过程中,一些可以有效指导查询集图像分割的信息会被忽略掉,没有在接下来引导分割的任务中保存下来。这样即使使用与查询集图片完全一样的支持集图片引导分割也会损失一定的分割细节。为此我们设置了一种自引导方式,使用之前计算得到的初始原型
首先对支持集进行引导分割。首先参考 [3] [6] 的做法,将初始原型拓展到和支持集特征相同的尺寸,新得到的原型特征称作
。然后将新得到的原型特征
与经过骨干网络得到的支持集特征沿着通道维度进行拼接得到一个新的特征图
。
(4)
通过这样一个融合了初始原型的特征,我们可以顺利对支持集图像进行分割,这样一个代表支持集图像的概率图可以通过解码器被求出来,对此有:
(5)
这里的
是对支持集图片预测的概率图,其中
,
表示解码器,这个解码器与后续对查询集分割的解码器共享权重。共享权重的解码器解码器的具体细节内容会在4.4节介绍。
是一个归一化函数。接下来利用支持集的概率特征图利用一个求最大值的函数完成对支持集预测掩膜的生成。
对此有:
(6)
这里的
依然是一个二进制的掩膜,在这个掩膜当中如果该位置值为0则代表这里被预测为支持集图像的背景,如果值为1则代表这个像素点被预测为目标类。通过对支持集预测的掩膜和支持集掩膜的真实值进一步处理,我们可以生成用来完成原型特征增强的两个辅助原型。
(7)
(8)
通过公式7和公式8可以计算得到用来实现原型特征增强的辅助原型,在公式7中
表示使用初始原型
作为引导分割支持集时正确预测的前景掩码,它代表在这些位置支持集待分割的目标被精确地分割出来。而公式8中
表示初始原型
引导没有被分割出来的部分,即缺失的前景掩膜。从公式7可以看到
可以有效聚合模型预测正确的信息,而
可以聚合之前分割中被遗忘的细节信息。通过这样的操作,之前初始原型中利用支持集图像得到泛化性信息被保留,而在初始原型中被忽略的细节公式8重新搜集。因此得到的辅助原型可以更加充分表达目标类的主体与细节信息,从而在接下来的步骤中更加有效地引导查询集进行分割。
4.4. 原型增强
4.4.1. 伪掩膜生成
遵循大多数小样本分割任务 [12] [14] ,我们也使用了伪掩膜作为最后解码器的输入。利用伪掩膜可以在没有可学习参数的情况下大致定位前景对象。具体而言,它通过测量编码器输出的查询集高层特征和支持集前景特征的余弦相似度实现。在这之后,得到的相似度被归一化,并作为查询集特征成为前景即待分割目标的概率值。我们利用编码器最后一层特征
和
以及变化到对应尺寸的掩膜进行上述操作。
(9)
其中
代表支持集前景特征,
代表变化后尺寸的支持集掩膜。之后对所得前景特征与查询集特征进行余弦相似度计算得到:
(10)
(11)
(12)
之后将
形状变变为三维得到
,其中
被设定为
。
(13)
4.4.2. 激活特征图
在激活注意力图的第一步首先进行特征匹配操作,密集特征匹配是小样本语义分割中的常见方法,它通常将查询集特征图
与全局特征向量在所有空间位置上进行比较实现特征匹配。具体而言我们分别在每个空间位置评估它们与查询特征的余弦相似性。由于在之前的步骤中得到了两个全新的原型,因此要将该操作拓展为多原型设置。值得注意的是,这里的查询集特征与求我们将两个原型按照通道维度进行拼接。具体的有:
(14)
(15)
(16)
是一个双通道的激活特征图,每一个通道与之前求得的辅助原型相对应。如此一来,两个辅助原型感兴趣的不同区域都可以在查询集中被激活。
4.4.3. 增强特征生成
最终增强的特征包含生成的辅助原型、激活特征图、查询集特征和伪掩膜。为了使这些特征融合在一起,我们首先将得到的辅助原型拓展到和查询集特征图相同的尺寸,然后用一种简单的通道拼接的方式将不同特征连接起来,最后通过一层卷积层完成对通道维度的压缩,一方面压缩通道可以降低模型的参数量提高模型的效率和推理速度,另一方面可以使特征更好的聚合在一起,形成一个有效的增强特征。这里的
表示对特征图拓展尺寸大小。
(17)
(18)
(19)
4.5. 解码器
该方法的解码器与BAM [12] 方法中的解码器保持一致,具体结构如图3:编码器得到的特征先经过空洞空间卷积池化金子塔 [27] ,随后经过几个由卷积和激活函数堆叠而成的模块,第一个模块由卷积核大小为1的卷积核和ReLU [28] 激活函数组成,第二个模块由一个卷积核大小为3的卷积和ReLU激活函数,第三个模块由2个连续的卷积核大小为3的卷积和ReLU激活函数构成。值得注意的是第二个模块和第三个模块之间有一个残差连接。最后一个模块依次由卷积核大小为3的卷积、ReLU激活函数、Dropout [29] ,卷积核大小为3的卷积连续堆叠而成。
Figure 3. Diagram of the decoder architecture
图3. 解码器结构示意图
4.6. 损失函数
为了保证
可以从支持特征图中学到大部分有用的信息,我们使用交叉熵损失函数对预测结果进行约束。在公式20到公式23中
与
分别代表模型输出的掩码概率图和对应的真实值。
(20)
类似的有对查询集预测结果进行训练
(21)
上面的公式表示的是利用交叉熵损失函数对查询集分割结果与真实值的差异做评估。其中
表示训练时的批次数量。最后对模型在第二阶段输出的结果在训练阶段利用
进行约束。最终我们将几个损失函数相加得到在元学习阶段的损失函数:
(22)
(23)
同样的,有在基学习阶段的损失函数,其中
表示训练时的批次数量。
(24)
5. 实验
5.1. 实施细节
标准小样本分割是在PASCAL-5i [8] 数据集上进行的。该数据集是基于PASCAL VOC 2012 [30] 中创建的,其中包含来自SDS [31] 的额外掩码注释,由20个语义类别组成,均匀分为4份。模型其中3份上进行训练,并以交叉验证方式在剩余的一份上进行测试。我们从网络未见过的那份数据集中随机采样1,000个片段进行评估。
遵循之前的小样本语义分割方法 [13] ,我们采用平均交并比(mIoU)作为评估指标,其中mIoU指标主要用于反映不同类别的整体性能。
为了与现有的方法进行公平比较,我们采用了相同的骨干网络(ResNet50 [26] )。根据 [13] 的方法,我们对这些骨干网络进行了预训练,并在小样本分割任务训练期间冻结它们以提高泛化能力。预处理技术的设置与 [13] 相同。使用SGD优化器,在PASCAL-5i上使用学习率为0.005进行200个epoch的训练,所有实验都在NVIDIA RTX 3090GPU上使用PyTorch框架进行。我们报告了使用4个不同fold进行训练测试的平均结果,以减少所选的支持–查询图像对对性能的影响
5.2. 方法比较
我们将提出的模型与最近的几个小样本分割方法进行比较。在表1中显示了在ResNet50骨干网络下PASCAL-5i数据集上的结果。从表格中可以看出,在1-shot和5-shot场景相比基线网络BAM [25] ,我们的方法的性能都有一定程度的提升。尤其是在1-shot下,每一次交叉验证中都取得了更好的效果。
Table 1. Performance on PASCAL-5i in mIoU
表1. PASCAL-5i数据集上mIoU指标性能对比
5.3. 消融实验
Table 2. Comparison results of ablation study on PASCAL-5i
表2. PASCAL-5i数据集上消融实验性能对比
我们在PASCAL-5i上进行了消融实验,消融实验结果如表2所示。对最后强化特征中仅保留了原始原型、辅助原型,辅助原型和激活特征图(Activation Map)的三种情况进行了消融实验,可以发现仅仅依靠辅助原型对模型性能提升并不大,但是和激活特征图配合使用会让模型性能有一定的提升。
5.4. 可视化
(a) (b) (c) (d)
Figure 4. Diagram of Visualization results
图4. 可视化结果展示
本小节将对模型在BAM上运行的结果和我们提出的方法上运行的结果进行比较,具体情况如图4所示,第(a)列是支持集图像,第(b)列是BAM方法的分割结果,第(c)列是本文提出的方法分割结果,第(d)列是查询集图像的标签,其中支持集目标对象以蓝色蒙版呈现,其余图像中的目标对象以品红色蒙版展现。可以看出我们的方法在细节上的分割比BAM做的更加出色。
6. 结论
在本文中,我们在BAM方法的基础上提出一个原型特征增强的方法。我们的方法充分利用了支持集的信息,通过自引导的方式构建了辅助原型和激活特征图,完成了原型特征增强。这种方法弥补了BAM对前景特征细节关注不够的问题,我们在PASCAL-5i数据集上进行实验,验证了方法的有效性。在未来,可以尝试更细致的方式完成对前景特征的充分提取,以此提高小样本语义分割的精确性。