基于注意力机制的U2-Net图像语义分割
2024-01-16邓晓冰杨火祥
刘 帅,邓晓冰,杨火祥,柳 伟
(1. 深圳信息职业技术学院计算机学院,广东 深圳 518172;2. 深圳大学应用技术学院,广东 深圳 518118;3. 广东技术师范大学电子与信息学院,广东 广州 510665;4. 深圳大学电子与信息工程学院,广东 深圳 518060)
随着深度学习的不断发展,卷积神经网络(Convolutional Neural Network, CNN)以其优秀的特征提取能力,被广泛应用于图像分类、目标检测以及语义分割等任务中。其中,全卷积神经网络(Fully Convolutional Network, FCN)的提出标志着图像语义分割技术迈向了新的阶段。Long等[1]在2015年提出全卷积神经网络FCN,实现了端到端的图像语义分割。此后,越来越多的研究人员开始利用CNN模型进行图像语义分割,并取得了较好的分割效果。虽然 CNN 能够有效地提取图像特征,但是其卷积和池化等操作会导致感受野和分辨率受到限制,进而影响了语义分割的准确性。因此,越来越多的改进模型应运而生。
自编码器(Autoencoder, AE)因其简单且有效的网络结构受到研究者们的关注。其主要思路是将输入数据通过多个卷积层提取特征后进行降采样,再通过反卷积层将特征图上采样至原始尺寸。Ronneberger等[2]基于FCN和自编码结构提出了一种改进的图像语义分割模型U-Net。它的主要特点是网络结构完全对称,并且在解码器部分加入了更多的卷积操作和跳跃连接,以此将局部细节的损失降到最小,从而提高图像分割的效果。随后,Qin等[3]在2020年提出了U-Net的改进版本U2-Net模型。该模型不仅可以用于显著目标检测(SOD),还可用于图像分割等其他任务。
除了上述利用自编码结构进行模型改进的方法外,还有一些研究者通过在卷积神经网络中添加注意力机制来进行模型优化[4-6]。注意力机制[7]的思想源于人类头脑的思考过程,即通过专注于自己关注的事物,来提高对这些事物的认知和理解。在图像语义分割任务中,使用注意力机制可以根据图像的特征自适应地选择感兴趣的区域,突出重要特征,抑制不重要信息,从而提高图像语义分割的精度。
因此,为了解决现有深度学习模型面对一些存在遮挡、噪声的图像容易预测出错的问题,本文在U2-Net模型的基础上,通过引入注意力机制对U2-Net模型进行改进和优化。通过在U2-Net模型中插入同时包含空间注意力和通道注意力的CBAM[8]模块,自适应的调整每一个通道的权重,捕捉特征图中的空间相关性,有效地减少信息丢失和冗余。改进后的模型能对不同通道和空间位置的特征图进行更细致的加权融合,以提高特征图的质量和准确性,进而提高模型分割结果的精度和准确性。在增广后的VOC2012数据集上的实验结果表明,在增加CBAM模块后,U2-Net模型的MIoU及准确率分别提高了8.21%和4%。
1 相关工作
1.1 图像语义分割
图像语义分割技术按照技术方法的不同,可以分为传统的图像语义分割和基于深度学习的图像语义分割。
传统的图像分割算法包括基于阈值、边缘检测、区域、主动轮廓以及图论等。这类方法通过手工方式根据颜色、形状、纹理、灰度等特征对图像进行划分,以完成对物体的分割[9]。然而,传统方法需要花费大量人力,并且在图像信息处理和实时场景理解方面的效率和精度都难以达到要求。相比之下,基于深度学习的图像语义分割技术具有更高的效率和精度。
基于深度学习的图像语义分割方法无需手动设计特征,而是利用大量的原始图像数据,通过卷积神经网络模型自动学习并提取高层次的特征,最终生成带有像素类别标签的分割图像。2015年,Long等人提出全卷积神经网络模型FCN,该模型可以适应各种尺寸的输入图像,是首个支持端到端训练的语义分割模型。2016年,Ronneberger等人提出经典的医学图像分割模型U-Net,U-Net采用U型自编码结构,并在对称的编解码阶段引入了跳跃连接,有效地避免了信息丢失的问题。该网络因其简单有效的结构,不仅在医学领域有着广泛应用,通过不断改进也可应用于其他场景[10-11]。同年,Google基于DeepLabv1[12]提出了DeeplabV2[13]模型,使用空洞卷积代替部分池化操作进行下采样,并引入空洞卷积金字塔模块进行多尺度特征提取。2017年,DeepLabV3[14]对空洞卷积金字塔模块进行了改进,引入轻量卷积结构,避免模型的参数量过大。2018年,DeepLabV3+[15]将整个DeepLabV3作为编码器进行特征提取,以获得更好的分割效果。2020 年 Qin等提出的一种既简单又强大的深度网络架构U2-Net。U2-Net在U-Net的基础上增加了U型残缺块(RSU模块),该模块的结构如图1所示。RSU模块由三个部分构成,包括输入卷积层、U型结构和残差连接层。其中,输入卷积层是对输入图进行通道转换,U型结构可以提取多尺度上下文信息,而残差连接层则将输入层和中间层通过残差块进行连接。
图1 RSU结构图
U2-Net整体结构如图2所示,主要由三个主要组成部分构成,包括6个编码模块,5个解码模块以及特征融合输出模块。每个模块都由RSU模块组成,以逐步地实现特征提取和上下采样处理。
图2 U2-Net模型结构
在编码模块En_1至En_4中,分别包含了7至4个RSU模块,每个编码模块中RSU模块的个数根据输入特征图的空间分辨率进行逐层减少。但在En_5与En_6中,特征图的分辨率相对较低,进一步下采样会导致上下文信息严重丢失。因此,在En_5与En_6使用RSU-4F模块,其中F表示对RSU模块的扩展,即使用膨胀卷积来代替上采样和下采样操作,使得RSU-4F模块中所有特征图的分辨率相同。
U2-Net模型可以实现多尺度特征提取和高效率计算,具有较好的分割准确性和高效率,在多个场景中都有改进和应用[16,17]。
1.2 CBAM注意力模块
CBAM(Convolutional Block Attention Module)注意力模块是通道和空间注意力机制的组合。通过通道注意力机制,网络可以学习到各个通道的重要性,并强制模型关注较重要的通道,同时抑制不重要的通道的影响,从而提高网络的性能。而空间注意力机制能够学习到每个空间位置的上下文信息,使得网络能够更好地提取图像中的语义信息。CBAM结构如图3所示。
图3 CBAM结构图
CBAM注意力机制利用了特征通道和空间上两个维度的信息,可以更加准确地学习到特征图的各个通道和位置的权重信息,从而提高了卷积神经网络的性能还有泛化能力。
1.2.1 通道注意力模块CAM
CAM模块主要用于获取通道注意力特征,将输入的特征图F分别进行全局最大池化与全局平均池化操作,在空间上对其进行压缩,生成两个大小为1×1×C的特征图,然后,将这两个特征图输入到一个包含两层ReLU激活函数的多次感知机(MLP)中进行处理。接着,将得到的输出特征逐元素相加,并通过Sigmoid函数生成通道注意力特征图FC。最后,将FC和输入的特征图F逐元素相乘,得到SAM模块的输出特征图F*。CAM结构如图4所示。
图4 CAM结构图
1.2.2 空间注意力模块SAM
空间注意力模块通过对特征的空间相关性进行处理来提取注意力信息,它与通道注意力模块有所区别却也可互补。空间注意力模块的输入特征图F*是通道注意力图FC和通道注意模块的特征输入逐元素相乘的结果,对其做最大池化操作和平均池化操作,得到两个H×W×1的特征图并将它们在通道上进行拼接,然后通过卷积核和Sigmoid激活函数来生成空间注意力特征图FS。最后将其与上一步获得的特征图F*逐元素相乘,得到最终的特征图F**。 SAM结构如图5所示:
图5 SAM结构图
2 基于CBAM的U2-Net图像分割
2.1 基于CBAM的U2-Net模型设计
尽管U2-Net模型在图像语义分割领域中表现很出色,但该模型在面对一些存在遮挡、噪声的图像时,依然容易出现分割不准确的情况,导致模型整体的分割性能下降。CBAM注意力机制能够在空间和通道维度上对特征图进行加权,提取图像的全局信息和局部信息,能更好地适应不同尺度和形状的目标,也能使模型更好地聚焦于感兴趣的区域。此外,CBAM还能够更好地指导模型从图像中提取到有关目标的更多信息,并减少噪声的干扰。因此,为解决U2-Net模型分割目标物体形状复杂时模型性能下降的问题,本文在U2-Net模型的基础上,通过引入注意力机制对U2-Net模型进行改进和优化。通过在U2-Net模型中插入同时包含空间注意力和通道注意力的CBAM模块,自适应地调整每一个通道的权重,捕捉特征图中的空间相关性,有效地减少信息丢失和冗余。本文提出的基于CBAM的U2-Net图像分割模型的结构如图6所示,本文模型能对不同通道和空间位置的特征图进行更细致的加权融合,以提高特征图的质量和准确性,进而提高模型分割结果的精度和准确性。
图6 基于CBAM的U2-Net模型设计
2.2 基于CBAM的U2-Net模型性能分析
为了验证提出的基于CBAM的U2-Net模型的有效性,本文将通过对比改进模型与原始模型的训练结果和在公共数据集上的测试结果。在训练模型前,先对实验所用到的数据集及实验环境进行介绍。
2.2.1 数据集
本文先使用数据集SBD[18]对VOC2012[19]进行扩充,然后去除两个数据集中的重复部分,得到增广后的VOC2012数据集,共包含12 031张图片。此外,本文还对增广后的VOC2012数据集进行了划分,将其中90%(共10 827张)作为训练集,剩下10%(共1 204张)作为验证集。
2.2.2 训练参数设置
本文所有实验均基于深度学习框架Pytorch[20]进行实现,训练过程中batch统一设置为32,最大学习率设置为1e-4,模型使用Adam优化器[21]进行训练,并设置优化器内部动量参数为0.9,权重衰减因子为0。模型总共进行100个epoch的训练,每5个epoch保存一次权值,并使用cos学习率下降策略对学习率进行调整。为了公平比较,本文模型和原U2-Net模型统一采用预训练的VGG模型作为主干网络,并基于预训练权重对模型进行训练。
2.2.3 训练结果对比
深度学习模型的损失函数是衡量模型性能的关键指标之一,损失函数值越小,则代表模型的预测值和真实值的差异越小,即模型性能越优。因此,本文通过对比U2-Net模型在增加CBAM注意力模块前后损失函数值的变化来评估模型的性能改善程度。图7、图8分别展示了原始U2-Net模型和基于CBAM的U2-Net模型在增广VOC2012数据集上的训练过程和损失变化。对比图7和图8,可以发现,基于CBAM的U2-Net模型的损失值下降速度更快,即收敛速度更快。具体而言,优化前的U2-Net模型需要经过60次迭代,损失函数值才会下降到0.3左右,而优化后的模型只需要经过40次迭代。表明优化后的模型更快地学习到数据集的特征,且优化后的模型的损失值比原模型的损失值小,证明了优化后模型的收敛速度和性能均有所提升。
图7 U2-Net模型的训练过程和损失变化
图8 基于CBAM的U2-Net模型的训练过程损失变化
3 实验与结果分析
3.1 实验评价指标
本文采用语义分割领域常用的两个评价指标—平均交并比MIoU和准确率Precision,来衡量模型性能和预测结果的准确性。
平均交并比MIoU是指真实分割结果与预测分割结果的平均交集占真实分割结果和预测分割结果的平均并集的比例。它首先计算每个类别的交并比IoU,然后将所有结果求平均值。IoU的计算方式如公式(1)所示:
上式中TP(True Positive)、FP(False Positives)和FN(False Negatives)分别表示“实际上是正样本,预测为正样本”“实际上是负样本,预测为正样本”和“实际上是正样本,预测为负样本”的数量。假设共有k+1个不同的分类,分别用L0到Lk进行标记,Pmm表示类别m的像素被预测为类别n的数量。此时,Pmm就等价于TP,而Pmn和Pnm则分别等价于为FP和FN,则MIOU的公式表示可以简化成以下形式:
准确率Precision的含义是在所有被预测为正样本的样本中,实际为正样本的概率,用P来表示准确率,计算方式如公式(3)所示:
3.2 实验结果分析
为验证所提方法的有效性,本文在增广后的VOC2012数据集上进行了实验,比较了添加CBAM模块前后U2-Net模型的性能变化。同时,为了进一步凸显CBAM模块的有效性,本文还与基于SENet[22]的U2-Net模型进行了比较,对比测试结果如表1所示。
表1 基于CBAM的U2-Net模型优化结果
由表1可以看出,本文所提基于CBAM的U2-Net模型在增广后的VOC2012数据集上的性能优于基准模型U2-Net,所取得的MIoU及准确率比基准模型分别提高了8.21%和4%,分别达到了67.7%和81.82%。由此可见,所提模型具有良好的图像语义分割性能,且通过在U2-Net模型中插入同时包含空间注意力和通道注意力的CBAM模块,可以帮助模型自适应地调整每一个通道的权重,捕捉特征图中的空间相关性,有效地减少信息丢失和冗余。此外,与仅提取通道注意力信息的SENet模块相比,基于CBAM的U2-Net模型的MIoU和准确率也分别提高了0.33%和0.37%。表明引入CBAM模块后,模型能对不同通道和空间位置的特征图进行更细致地加权融合,以提高特征图的质量和准确性,进而提高模型分割结果的精度和准确性。
为了进一步体现所提方法的有效性,本文将所提方法和基准模型U2-Net模型在增广后的VOC2012数据集上的部分图像分割结果可视化,如图9所示。从图9中的飞机机翼分割结果可以看出,本文所提模型可以完整的分割机翼部分,原始U2-Net无法完整的分割出机翼部分,增加SE模块的U2-Net模型的分割结果相较于U2-Net而言更加完整,但仍然存在部分区域不能正确分割。表明了在U2-Net模型中添加CBAM模块可以准确地捕捉特征图中的空间相关性,有效地减少信息丢失,提高模型的分割准确性。
图9 部分图像分割结果对比
图10 直观地展示了增加注意力机制后的U2-Net模型在MIoU及Precision两个评价指标上的性能变化。不难看出,通过在U2-Net中添加注意力机制是一种简单且有效的模型优化方法,它可以显著提高模型的性能。
图10 增加不同模块的网络优化分析
4 结论
本文针对现有深度学习模型面对一些存在遮挡、噪声的图像时容易分割出错的问题,提出通过引入注意力机制来提高模型的鲁棒性和泛化能力。本文以U2-Net模型为基准,通过在U2-Net模型中插入CBAM模块来调整特征图中每一个通道的权重,并捕捉特征图中的空间相关性,有效地减少信息丢失和冗余。通过实验对优化后模型的性能进行了综合评估,表明CBAM模块可以显著提高模型的性能和泛化能力,证明了所提方法的有效性。