基于特征融合注意力的小样本语义分割算法*

李呓瑾,李少龙**,贺 彦,刘 炜

(1.云南电网有限责任公司信息中心,云南昆明 650200;2.北京国科恒通科技股份有限公司,北京 100085;3.清华大学电机工程与应用电子技术系,北京 100084)

语义分割的目的是为图像中的每个像素点分配一个类标签,在医疗诊断、无人驾驶、图片编辑等领域具有广泛的应用前景[1]。近年来,基于深度学习的语义分割模型取得了突破性的进展,如全卷积神经网络FCN、DeepLab、UNet和PSPNet等[2],并利用膨胀卷积来增大感受野[3],从而增强了模型的分割性能。然而,基于深度学习的主流语义分割模型需要大量逐像素的标注数据,这类标注数据的获取费时费力成本高。虽然弱监督学习方法可以缓解模型对像素级标注的依赖,但仍然需要大量的弱标注数据[4]。此外,基于弱监督学习方法的模型对于新类或标注不充分的目标类的泛化性能较差。

受小样本学习的启发,Shaban等[5]提出了一种基于双分支结构的小样本语义分割模型,其中支持分支将支持图片及对应的标注掩码作为输入,学习指导信息;查询分支以查询图片为输入,预测对应的分割掩码。之后,众多研究者基于该双分支结构去构造各种变体分割模型,提高模型对目标物体的分割性能[6,7]。

现有的小样本语义分割模型主要包括度量学习和元学习两大类[8]。度量学习方法首先利用支持分支获得每一像素的特征表示,然后通过计算查询图片中每一像素与特征表示之间的相似度,给出查询图片中目标物体的分割。Wang等[9]提出了一种特征对齐的小样本语义分割模型,该模型利用全局平均池化策略获取支持图片的全局特征并将其作为目标物体的特征表示。Zhang等[10]针对全局平均池化策略极易造成目标背景对前景干扰的问题,提出了一种掩码平均池化策略,该策略利用支持图片的掩码来分离目标的前景和背景,充分挖掘前景信息从而提高分割的性能。然而,仅利用全局平均特征构造单一原型不足以充分表示目标的不同部分。为此,Liu等[11]提出了一种多特征表示的小样本语义分割模型,其通过均等划分支持分支中目标物体的多个区域并借助掩码平均池化获取每个区域的平均特征,从而构造目标物体的多个特征表示。类似地,Li等[12]提出了一种自适应多特征表示的小样本语义分割模型,其借助支持分支的掩码信息获取多个不同大小的目标区域,并利用掩码池化策略提取每个区域的特征表示。贾熹滨等[13]提出了一种金字塔原型对齐的小样本语义分割模型,它通过提取不同尺度的特征构造目标物体的特征表达,并通过计算特征表达与查询分支目标物体之间的相似度分割目标区域。Liu等[14]设计了一种动态原型卷积网络来构造类的特定多原型表示,该网络能够充分捕获目标的细节特征从而提高分割性能,并且可以被应用在小目标、多目标等复杂场景中。

元学习又称“学习如何学习”,指的是通过在多个任务中学习分割先验知识(一组参数)来指导新任务的过程,旨在强化模型的泛化性能。Liu等[15]提出了一种基于Transformer的小样本语义分割模型,该模型利用Transformer编码块提供的注意力机制,构造动态权重的分类器,并采用预训练模型固化的策略,解决了小样本数据样本不足的问题。刘宇轩等[16]针对支持图像和查询图像共性信息利用不足的问题,提出了一种结合全局和局部特征的小样本语义分割模型,并在PASCAL数据集上验证了其有效性。Pambala等[17]提出了一种基于元学习的小样本语义分割模型,在视觉信息的基础上通过引入文本语义信息来构造多尺度的融合特征,并利用解码器分割目标物体。虽然上述模型取得了一定的成功,但高质量的特征提取更能够进一步提升下游的分割性能。Tian等[18]在特征编码阶段提出了一种强区分性的元学习模块,通过挖掘局部和全局特征图来进一步提高特征的表达能力。Wu等[19]在支持分支和查询分支特征提取模块之后引入了一种元记忆学习模块,通过学习记忆支持图片和查询图片的相似性语义信息,来强化特征的表达能力。

虽然基于度量学习的小样本语义分割模型结构简单,参数较少,但其分割性能过度依赖于原型的质量。此外,仅利用无参数的相似度度量计算极易导致信息丢失或歧义。基于元学习的小样本语义分割模型虽然可以解决上述性能不佳的问题,但现有的该类分割模型主要聚焦于支持分支中支持图片的前景和背景的信息挖掘,而无法有效利用查询图片的信息。因此,本文提出的多特征融合的小样本语义分割模型仍然采用元学习的架构并在充分利用支持图片的前景、背景信息的同时,进一步挖掘查询分支中查询图片的信息,从而解决了现有的基于元学习模型对于查询图片信息利用不充分的问题。具体来说,所提出模型首先通过引入注意力模块从支持图片和查询图片中学习语义关联特征图;然后,借助掩码平均池化方法提取支持图片中前景和背景信息的全局特征;最后,利用特征融合模块将前景和背景特征进行尺度融合,并将其融合特征作为解码块的输入实现查询图片的掩码预测。

1 任务定义

θ的训练和测试采用episodic机制,每个episodic可以简化为一个C-way K-shot的学习任务。支持集Si由K张图片及其对应的掩码组成,即S=(image,mask);类似地,查询集Qi由与支持集语义类相同的L张图片组成Q=(image,mask)。然而,测试阶段中的查询图片没有对应的mask,即image∈3×h×w,mask∈h×w。

2 小样本语义分割模型

2.1 模型结构

图1给出了所提出模型的结构,包括特征提取、语义关联注意力模块、多特征融合模块和解码器等4部分。首先,使用同一个特征提取器编码查询图片和支持图片;其次,利用注意力机制来映射支持特征和查询特征之间的强语义相关;最后,通过特征融合注意力模块融合多个不同来源的特征,并利用解码块分割查询图片。

2.2 特征提取

特征提取的质量直接影响分割的效果。虽然浅层特征携带颜色、边缘等低层次线索,但其在语义层面上的区分性不强;而高层特征尽管具有较强的语义类区分能力,但难迁移到对未知类的分割任务中。此外,小目标物体由于下采样操作,其信息极易丢失,从而影响分割效果。为此,本文通过构造多尺度特征提取网络来强化深度特征空间的语义表达能力。特征提取网络结构如图2所示。

图1 模型架构

具体来说,首先利用在ImageNet上预训练的VGG-16、ResNet-50和ResNet-101分别作为主干网络进行多尺度特征提取。多尺度特征表示如公式(1)所示。

Fs=fsl⊕fsm⊕fsh,

(1)

其中,Fs表示融合后的特征;fsl表示低层特征,fsm表示中间层特征,fsh表示高层特征;⊕表示向量的连接操作。

考虑到深度卷积网络对于硬件要求较高,此处采用深度可分离卷积代替主干网络中的部分卷积块,通过减少模型参数计算量来降低运行时间开销,深度可分离卷积的结构如图3所示。具体来说,首先利用主干网络的第一个block将输入的支持图片Is和查询图片Iq映射到浅层特征空间,获得特征图Fo;然后,作为深度卷积(Depth-Wise Convolution,DWC)的输入沿着通道维度进行分解。将标准卷积过程分解成K个等效的DWC和N个逐点卷积(Pointwise Convolution,PC),值得一提的是PC能够在保持特征提取质量不变的基础上降低计算复杂度。此处,DWC的输出特征图表示为Fdw,支持图片Is和查询图片Iq的深度融合特征表示为Fpc。最后,将主干网络中每个block替换成图3所示的卷积结构,从而获得最终的支持图片融合特征Fs和查询图片融合特征Fq。值得注意的是,此处特征包括前景区域特征和背景区域特征。

2.3 语义关联注意力

现有的小样本语义分割模型大多使用从支持分支中提取类的特征表达来指导查询图像的分割[9,12]。然而当带标注的支持图片数量有限时,极易导致从中提取的特征不足以表达类的强语义相关性,从而使得模型对于查询图片的预测不精确。为了解决该问题,受注意力机制在视觉和文本领域成功应用的启发,本文提出一种语义关联注意力模块(图4)。

图3 深度可分离卷积

图4 语义关联注意力模块

考虑到支持分支和查询分支中的图片具有相同的语义信息,因此利用图4所示的语义关联注意力模块获取新的强语义特征。首先,将{Fs,Fq}∈512×56×56作为注意力模块的输入;然后,利用余弦相似度计算支持特征和查询特征的Top-H得分,从而选取最相似的H个特征构造新的语义特征。利用以上语义关联注意力模块获取的语义特征能够在特征融合阶段帮助定位查询图片的目标区域,从而有效解决现有模型对于目标区域定位不精确的问题。

图4中,模块首先计算了支持特征Fs和查询特征Fq在每一位置(x,y)处的余弦相似度[11],计算公式如下:

(2)

然后,利用公式(3)选择得分最高的H个相似性特征构造语义关联特征Rs∈K×56×56。

Rs= argmaxTop-H(Ds)。

(3)

最后,利用全局平均池化(Global Average Pooling,GAP)[20]获取每个特征区域的平均特征Rs′∈K×1,并利用softmax函数沿通道方向计算每个位置的语义注意力Asi,计算公式如下:

(4)

2.4 多特征注意力融合模块

由于拍摄角度、颜色变化和遮挡等外界因素的影响,来自同一语义类的支持图片和查询图片会存在很大差异,这极大地增加了分割的难度。因此,本文设计了一种多特征注意力融合模块,在支持图片和查询图片语义注意力生成图的基础上,该模块通过融合支持图片掩码过滤后的前景信息和查询图片的深度语义信息来进一步强化支持图片和查询图片之间的语义关联。具体来说,特征注意力融合模块的输入端包括语义关联注意力特征Rs∈K×56×56、掩码后的支持特征Fs′∈512×448×448和查询特征Fq∈512×56×56这3个分支。此处借助注意力机制的思想,融合(R,Fs′)、(R,Fq)和(Fs′,Fq)3个不同尺度的特征,并将融合后的特征进行归一化操作,得到最终的融合特征FF∈512×56×56。特征融合流程如图5所示。此处,以{Fs′,Fq}为例说明特征融合的流程。

图5 特征融合注意力

首先,将{Fs′,Fq}∈512×56×56利用reshape函数转换为{Fs′,Fq}∈512×N,N=56×56,并利用公式(5)计算单特征之间的相似性矩阵M1∈512×512。

M1=Fs′FqT。

(5)

然后利用softmax函数按行计算注意力权重W,并与单特征矩阵相乘获得融合特征F1。计算如公式(6)所示。

F1=softmax(M1)×Fs′+softmax(M1)×Fq。

(6)

接着利用点乘运算将所有特征两两融合,得到最终的强语义特征FF=F1(Fs′,Fq)·F2(R,Fq)·F3(R,Fs)。最后,将强语义特征FF∈512×56×56作为解码器的输入,得到查询图片的分割掩码,并利用交叉熵损失函数实现模型端到端的优化。

3 实验与结果分析

3.1 数据集

采用经典的PASCAL-5i[21]和COCO-20i[22]数据集进行模型训练与测试。其中PASCAL-5i包含5 953张训练图片和1 449张测试图片,共包含20个类,划分为4折,其中3折(15类)用于训练,剩余用于测试;COCO-20i每张图片中包含的类别更多,包括82 081张训练图片和40 137张测试图片,共包含80类,其中3折(60类)用于训练,剩余用于测试。数据集详细信息如表1所示,部分样本如图6所示。

表1 数据集描述

图6 样本图片

3.2 实验环境与评价指标

硬件:Nvidia A100 8×40 GB GPU。软件:Python 3.7,pytorch深度学习框架。利用文献[7]中的VGG-16和文献[15]中的ResNet-50和ResNet-101等预训练模型作为特征提取器。输入图片的大小为448×448;优化器采用Adam;初始学习率为0.000 1,权重衰减因子为0.000 5,batch_size为8。

采用平均交并比mIoU (mean Intersection over Union)和前景背景二分类交并比FB-IoU(Foreground and Background IoU)[5]作为评价指标。其中,mIoU为所有类别真实值和预测值的交集和并集之比,FB-IoU表示类别为2的二分类任务。计算公式如(7)所示。

(7)

其中,k+1表示识别的总类别数,i表示真实标记,j表示预测标注,pij表示将i预测为j。

3.3 对比实验

3.3.1 PASCAL-5i

为验证所提出模型的优越性,在PASCAL-5i数据集上与当前主流模型进行对比实验,详细结果见表2。当VGG-16作为主干网络时,所提出模型在1-shot任务中可以获得50.6%的mIoU和69.2%的FB-IoU。在5-shot任务中,虽然在mIoU评价指标下所提出模型略逊于PANet,但在FB-IoU指标下所提出模型比PANet提高了0.84%;当ResNet-50作为主干网络时,所提出模型在1-shot和5-shot任务上可以达到59.4%和60.1%的mIoU以及72.8%和73.4%的FB-IoU,整体优势明显;当ResNet-101作为主干网络时,在1-shot和5-shot任务中,尽管在mIoU评价指标下所提出模型略逊于GL模型,但在FB-IoU指标下,所提出模型分别在1-shot和5-shot任务中比GL模型提高了1.22%(73.8%→74.7%)和0.94%(74.8%→ 75.5%)。

表2 1-way 1-shot和1-way 5-shot在PASCAL-5i上的分割结果

3.3.2 COCO-20i

为进一步验证所提出模型的有效性,在更具有挑战性的COCO-20i数据集上与当前主流模型进行对比实验,详细结果见表3。当VGG-16作为主干网络时,所提出模型在1-shot任务中可以达到35.8%的mIoU和61.6%的FB-IoU;在5-shot任务中,所提出模型略逊于FFNet;当ResNet-50作为主干网络时,所提出模型在1-shot和5-shot任务上可以达到38.6%和41.3%的mIoU以及64.2%和66.8%的FB-IoU;相比1-shot的分割任务,5-shot的mIoU增加7.0%(38.6%→41.3%),这表明所提出模型可以充分利用支持图片的信息来提高模型的分割性能。当ResNet-101作为主干网络时,所提出模型在mIoU和FB-IoU评价指标下,相比所有对比模型,整体优势明显。

3.3.3 2-way 1-shot和2-way 5-shot

为了验证所提出模型的鲁棒性,选择ResNet-50作为主干网络,PANet[9]为对比模型,在PASCAL-5i和COCO-20i数据集上分别进行2-way 1-shot和2-way 5-shot的对比实验,结果详见表4。所提出模型在PASCAL-5i和COCO-20i数据集上整体2-way的分割性能明显强于PANet。具体地,在PASCAL-5i数据集上,所提出模型在1-shot任务上可以达到48.3%、55.1%和56.8%的mIoU,相比PANet至少提升10.9%;在5-shot任务上可以达到50.6%、57.1%和58.3%的mIoU。在更具有挑战性的COCO-20i数据集上,所提出模型在5-shot任务上可以达到36.8%、37.9%和32.9%的mIoU,相比PANet至少提升9.3%。上述结果进一步验证了所提出模型更加鲁棒。

表3 1-way 1-shot和1-way 5-shot在COCO-20i上的分割结果

表4 2-way 1-shot和2-way 5-shot在PASCAL-5i和COCO-20i上的分割结果

3.3.4 可视化结果

图7是所提出模型和PANet模型的分割可视化结果。从图7可以看出,虽然PANet可以准确地找到待分割目标的位置,但仍会出现丢失部分细节或误分割的情况。相比而言,所提出模型可以准确地将轮船、飞机、猫和自行车等类别从查询图片中分割出来,整体分割效果更好,可视化结果进一步验证了所提出的多特征融合的小样本语义分割模型的优越性。

图7 分割可视化结果

3.4 消融实验

首先,为了降低模型的参数量,所提出模型采用深度可分离卷积代替主干网络中的卷积块。为了探究深度可分离卷积对整体性能的影响,在1-way 1-shot 任务上利用PASCAL-5i和COCO-20i数据集,采用ResNet-50作为主干网络设计消融实验。实验结果如表5所示,利用深度可分离卷积块代替标准卷积块的设计虽然对mIoU和FB-IoU评价指标的提升效果不明显,但是能够极大地减少模型参数量,从而有效地解决模型对硬件过于依赖的问题。

表5 深度可分离卷积消融实验

其次,为了探究语义关联注意力模块对模型性能的影响,在1-way 1-shot 任务上利用PASCAL-5i和COCO-20i数据集,采用ResNet-50作为主干网络设计消融实验,实验结果如表6所示。结果表明引入语义关联注意力模块能够显著提升模型的分割性能,究其原因是语义关联注意力模块引入了支持图片和查询图片之间的强语义关联,这有助于捕获到目标物体的位置及其更多细节信息。

表6 语义关联注意力模块消融实验

最后,为了验证所提出设计的多特征注意力融合模块的优越性,在1-way 1-shot任务上利用PASCAL-5i和COCO-20i数据集,采用ResNet-50作为主干网络设计消融实验,基线模型采用简单的向量拼接方式将多个输入特征图进行融合,实验结果如表7所示。相比基线模型,将语义关联注意力特征图、掩码后的支持特征和查询特征利用多特征注意力融合模块进行特征融合,在两个数据集上分别提升了10.4% mIoU和13.4% FB-IoU(1-shot),18.0% mIoU和8.3% FB-IoU(5-shot),分割性能明显提升。

表7 多特征注意力融合模块消融实验

4 结论

针对现有模型对于支持分支查询图片信息利用不充分的问题,本文在元学习方法的基础上,提出了一种新的小样本语义分割方法。该方法利用深度可分离卷积代替原始卷积块,降低了其对硬件资源的要求。此外,受注意力机制的启发设计一种语义关联注意力模块,该模块通过计算支持特征和查询特征之间的相似性来引入上下文语义信息,从而提高了模型对未知目标物体分割时定位的准确性。在PASCAL-5i和COCO-20i数据集上进行测试,实验结果表明所提出模型比起对比模型,整体优势较为显著。在未来工作中,可以利用交叉注意力机制建立查询编码特征与支持前景特征之间的关联,从而强化目标前景原型的表达能力。