基于深度学习的Attention U-Net语义分割模型研究

薛泽民,邹连旭,黄志威,冉杰,余若岩,郑国勋

(1.长春工程学院,长春 130012; 2.长白山历史文化与VR技术重构吉林省重点实验室,长春 130012)

0 引言

近年来,深度学习在计算机视觉领域大放异彩,已逐渐成为主流研究方向,越来越多的研究人员将其深度学习技术应用于图像语义分割。在深度学习进行图像处理的研究中,通常是将图像输入到编码器中,得到1个中间的上下文,再使用解码器对该上下文进行解码,最后还原成1个输出的序列。而目前的做法是使用1个循环神经网络实现编码器,用另外1个循环神经网络去实现解码器。针对编码器和解码器的实现,还有一些人采用不同模型结合的方法,比如使用双向循环神经网络或者采用带有长短期记忆的双向循环神经网络等[1]。这些方法均需要将大量的信息压缩成1个固定长度的中间上下文向量,因此使得大量的信息在被压缩后产生损失,尤其是先输入的信息会被后输入的信息稀释,且随着输入的序列越长,问题越严重,最后会导致解码质量严重下降。

Jay Alammar为解决这种编码器-解码器结构存在的问题,在传统的编码器中引入了注意力机制。加入注意力机制后,编码器需要将源图像编码成一个个向量。与传统的编码器-解码器结构不同的是,传统方式只需要把最后1个向量当作上下文向量传给解码器即可,但在注意力机制中,需要把所有的向量传给解码器,而在解码时,会多一个额外的步骤,即通过注意力机制找出与当前输出时间步长最相关的编码器时间步长的输入向量,再产生当前时间步长的输出,使在产生每一个输出的时候,都能够做到充分利用输入序列携带的信息,CHO等[2]采用这种方法在图像语义分割中取得了非常不错的成果。本文针对通过生成对抗网络GAN扩充的数据集,提出了融入注意力机制的U-Net网络训练的图像语义分割模型。

1 生成对抗网络

生成对抗网络是一种深度学习模型,被认为是近年来最有前景的无监督学习方法之一,适用于处理复杂分布的数据。与传统方法需要标记大量的数据相比,生成对抗网络可以在无需标注数据的情况下进行生成任务的学习。该网络结构由生成器和判别器两部分组成。首先生成器从潜在空间中进行随机取样并将其作为输入,生成高度逼真的样本。再把判别器接收的真实样本和生成器生成的虚假样本作为输入,并努力将它们区分开。生成器和判别器相互对抗、不断学习,模型通过这两个模块的博弈和优化逐步达到平衡点,从而实现扩充数据的目的[3]。为了保持生成器与判别器之间平衡的稳定性,即纳什均衡,必须满足式(1),结构如图1所示。

图1 生成对抗网络结构图

(1)

式中V(D,G)表示生成对抗网络整体待优化的损失函数,z表示随机变量,x表示真实数据,Pz表示生成数据的概率分布,P表示x的概率分布,D(x)表示判别x是否为真实的概率,D(G(z))表示判别重建数据是否真实的概率。

2 注意力机制

注意力机制是一种通过模拟人类视觉和认知系统的方法,通过注意力机制可以让模型知道图像中不同局部信息的重要性,从而提高模型的性能和泛化能力。主流的注意力机制采用编码器-解码器结构,如图2所示,该结构将一个变长的输入X=(x1,x2,…,xn)映射到一个变长输出Y=(y1,y2,…,ym)。其中编码器把一个变长的输入序列X,通过非线性变换转化为一个中间的语义表示C(C=f(x1,x2,…,xn))。解码器则根据输入序列X的中间语义表示C和先前已经生成的y1,y2,…,yi-1,预测并生成i时刻的输出yi=g(y1,y2,…,yi-1,C)[4],在前述过程中,f()和g()为非线性转换函数。

图2 注意力机制的编码器-解码器结构

注意力机制包含多种,本文使用的是空间注意力。空间注意力是通过空间转换模块将原始图片中的空间信息转换到另一个空间,并保留关键信息,使关键区域的特征表达得到提升。它为每个位置生成权重掩膜并对特征进行加权输出,从而增强感兴趣的特定目标区域并减弱不相关的背景区域。空间注意力模块的实现如图3所示,首先对给定的H×W×C特征F′进行全局平均池化和全局最大池化操作,得到2个H×W×1的特征图,分别代表不同信息;然后将这2个特征图进行拼接,经过一个具有较大感受野的7×7卷积层进行特征融合;再通过Sigmoid操作生成一个权重图Ms,并将其叠加到原始输入特征F′上;最后根据权重图对输入特征进行缩放,从而完成增强目标区域的特征表达,计算过程如式(2)所示。

图3 空间注意力实现图

Ms(F)=σ(f7*7([AvgPool(F),MaxPool(F)])),

(2)

式中σ表示Sigmoid函数,f7*7表示7×7大小的卷积核,AvgPool表示平均池化,MaxPool表示最大池化。

3 融入空间注意力机制的U-Net

在U-Net网络中,融入空间注意力模块,可以使网络更加关注感兴趣的区域,并减少对不相关区域的依赖,我们称其为Attention U-Net。即在U-Net的编码器和解码器之间插入空间注意力模块,通过该模块计算输入特征图的空间相关权重来调控信息的传递。具体来讲,我们使用空间映射的方法对输入特征图进行卷积操作来获取感兴趣区域的权重分布,通过卷积层和激活函数来学习感兴趣区域的空间特征,然后使用空间调整法将计算得到的注意力权重与输入特征图逐元素相乘得到经过调整的特征,最后将调整后的特征送入下一层网络进行处理,从而增强感兴趣区域的特征响应并有效抑制不相关区域的干扰。为了实现Attention U-Net,需要完成3部分内容:U-Net的实现、注意力门模块的实现和U-Net与注意力门模块的融合。

3.1 U-Net的实现

U-Net是一种基于FCN架构构建的U型网络,其结构如图4所示。左侧部分是收缩路径,经过一系列卷积和池化进行下采样,提取图像特征。右侧部分是扩展路径,进行上采样,将特征图与左侧提取的特征相融合,在卷积层进行特征维度的调整,然后再进行上采样和融合。重复前述过程直到上采样到与原图相同的尺寸时,添加1个卷积层用于对每个像素点进行分类来实现图像的语义分割。这种结构充分利用了上采样和下采样过程中提取的多尺度特征信息以及跳跃连接的方式,能够有效地获取全局和局部的上下文语义信息,并将其应用于像素级的分割任务[5]。

图4 U-Net结构图

3.2 Attention U-Net的实现

3.2.1 实现注意力门模块

图5 注意力门模块结构图

(3)

式中σ1表示激活函数,*表示卷积操作,k表示卷积核参数,Fl表示在第l层的感受野内的所有位置的集合。

注意力系数α用来识别图像中的显著区域并修剪特征响应,它将输入特征图与注意力系数进行元素乘法来计算。通常情况下,我们可以通过特征图中的每个像素向量计算出该像素标量的注意值。当涉及多个语义类别时,可以学习多维的注意力系数,使得每个注意力门能够关注目标结构的1个子集,让每个像素使用1个门控向量来确定聚焦区域并采用加性注意力方法来计算门控系数。通过这些方式,我们能够建立适应不同场景和任务的注意力模式,并提高关键特征表达的准确性,加性注意力的公式如式(4)所示。

(4)

式中σ1表示激活函数,σ2表示Sigmoid函数,Θatt表示注意力门的一组参数,Wx、Wg和Ψ表示卷积操作,bΨ和bg表示对应卷积的偏置项。

3.2.2 将注意力门模块应用于U-Net

在标准的U-Net架构中融合了前文提出的注意力门模块,通过突出跳跃连接中的显著特征,将粗尺度提取的信息应用于门控中,消除跳跃连接产生的不相关和嘈杂的响应。在前向传播和反向传播过程中,使用注意力门对神经元的激活进行过滤和减少产生背景区域梯度的权重等方法,确保更新底层模型参数时更专注于与给定任务相关的空间区域。具体实现的结构如图6所示,对于l-1层的卷积参数,其更新方式如式(5)所示。

图6 Attention U-Net网络结构图

(5)

4 试验

4.1 试验环境

本试验采用Windows 10操作系统,11th Gen Intel(R) Core(TM) i5-11260H @ 2.60GHz (12 CPUs)处理器,NVIDIA GeForce RTX 3060显卡,深度学习框架为Pytorch1.3.0版本,采用CamVid(Cambridge-driving Labeled Video Database)数据集。从中随机挑选700张大小为970×720像素的图像测试数据作为试验的数据集,并按照5∶2的比例将数据集划分为训练集和测试集,训练集共500张,测试集共200张,通过生成对抗网络将数据集扩充到1 000张。

4.2 训练参数及评价指标

4.2.1 训练参数

使用本文的模型与U-Net+ +、FCN、DeepLabV1对扩充后的数据集进行训练,将Batch size(1次训练所抓取的数据样本数量)设置为4,迭代次数设置为100,当损失趋于稳定时停止训练。

4.2.2 评价指标

在深度学习领域,常用的语义分割评价指标有准确率(Accuracy)和损失值(Loss)等,这些参数可以较好地反映出一个模型的性能。准确率是用于评估分类模型的指标,简单来讲就是模型预测正确的数量占总量的比例,计算过程如式(6)所示。损失值是用于评价模型在训练集上的拟合程度的指标,通常情况下,损失越小则训练数据越拟合。

(6)

式中TP表示真正例,TN表示真负例,FP表示假正例,FN表示假负例。

4.3 试验结果及分析

从图7可以看出随着时间的增长,在80轮以后,准确率和损失值逐渐趋于稳定,结合表1可以得出在第100轮时准确率约为95.4%。从表2各种模型训练100轮的结果可以看出,在各自损失值趋于稳定的情况下,与其他模型相比,本文模型的训练准确率较高,损失值与U-Net++、DeepLabV1较为接近。

图7 训练结果图

表1 对模型进行100轮训练的结果

表2 与其他模型的对比结果

5 结语

针对当前深度神经网络在图像分割过程中普遍存在的处理耗时长、分割准确率不高的问题,本文在使用生成对抗网络扩充数据集后,提出了一种融入注意力机制的U-Net模型用于图像分割。具体来讲,就是以U-Net网络为基本框架,在每个解码器层中都添加1个注意力门模块,帮助模型更准确地区分前景和背景,提高学习效率,加强特征区分。通过对比试验得出,相较于各传统语义分割模型,如U-Net++、SegNet等,本文的模型训练时间较短,同时又具有较高的分割准确率。