MALUNet:一种多关注,轻量级的皮肤病变分割UNet

Tittle:MALUNet: A Multi-Attention and Light-weight UNet for Skin Lesion Segmentation

摘要

本篇文章提出的目的是为了缓解计算资源的不足,轻量化模型提出的。简单来说本文提出了4个模块。

1)DGA:由扩展卷积和门控注意机制组成,提取全局和局部特征信息。(使网络模型更加关注目标区域,同时这一模块中的深度可分离卷积进一步减少了参数和计算复杂度)

2)IEA:基于外界关注来表征整体数据集,增强样本之间的联系。

3)CAB:由一维卷积和全连接层组成,对多阶段特征进行全局和局部融合,生成通道轴上的特征图

4)SAB:通过共享二维卷积对多阶段特征进行操作,生成空间轴上的特征图。

本文将这四个模块与UNet网络结合

Introduction

目前一些分割架构存在的问题是:

1)一些模型以大参数量的代价来换取模型的高性能,但是由于医疗设备内存有限无法很好的应用于现实的临床场景

2)样本数据之间的差异很小,但在医疗数据集内部中的差异很大,目前解决这一问题的是在自注意力机制及其改进结构上进行的。但是SA在建模样本之间关系方面效果不算好。

3)医学图像分割任务中,由于目标尺寸不同,必须充分利用多阶段,多尺度的信息。但是先前的一些网络结构忽略了这一点。

方法

1)Dilated Gated Attention Block-DGA

同时获取全局和局部信息对提高图像分割性能很重要。全局信息有助于模型了解病灶的整体结构及其与背景的关系,使得其更准确的定位到病灶区域。局部信息有助于获取病灶区域边缘和角落的细节,使得预测更加完整。

DGA模块的结构如下图所示:

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第1张图片

 

模块由两个子单元构成分别为:Split Dilated Conv Unit (SDC) and Gated Attention Unit (GA).

SDC将输入特征图沿通道维度分成四部分,通过使用不同扩张率的深度可分离卷积,得到全局(扩张率为5和7的卷积)和局部(扩张率为1和2的卷积)特征信息。然后在通道维度上进行Concate操作,恢复特征图的大小,然后进行卷积操作,使得全局信息和局部信息得以交互

GA通过使用深度可分离卷积生成与输入特征形状一致的特征图,抑制SDC传输的特征信息中不重要的区域,使得模型更加关注重要的信息。最后应用残差连接来获得输出。用公式表示为

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第2张图片

2)Inverted External Attention Block--IEA模块

External atttention(EA)称为外部注意力,它使用两个内存单元(两个共享参数的Conv1D操作)来表征整个数据集的特征信息。而本文提出的外部注意力块如下图所示:

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第3张图片

 给定一个输入X\epsilon R^{C*H*W}经过1*1卷积核Reshape运算以后大小变为X\epsilon R^{C*HW}。我们使用Memory Unit 1将特征图的通道数扩展成原来的4倍得到Att\epsilon R^{4C*HW}。然后再应用Memory Unit2恢复尺寸,再进行Reshape操作恢复到原始特征图尺寸,再进行Reshape操作恢复到原始特征图大小。最后通过1*1卷积得到输出特征图并加入残差信息。

优点在于Memory Unit是将输入映射到更高维的空间,使得内存单元更全面的描述数据集的整体特征信息。

3)Channel Attention Bridge Block

多阶段多尺度信息的获取对不同尺度目标的分割起着至关重要的作用。这个模块在通道维度上提出了一个桥式注意模块称为CAB。这个方法将通道轴上不同阶段的特征串联起来,生成通道注意力图,从而更好的整合信息。CAB结构图如下所示:

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第4张图片

其中GAP是全局平均池化,t_{i}为编码器获得的不同阶段的特征图,Contact是通道维度的级联运算。FC_{i}是第i阶段的全连接层。

如图所示,以五层模型(s = 5)为例,可视化CAB。注意,CAB将多阶段、多尺度的信息融合细分为局部信息融合(1D卷积运算)和全局信息融合(每个阶段不同的完全连接层),以提供更有信息量的注意特征图。

4)Spatial Attention Bridge Block--SAB

 CBAM模块认为通道注意和空间注意在抑制无关信息方面都发挥着重要的作用。因此在CABM的启发下,我们引入了另一种桥注意模块--SAB

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第5张图片

如图6所示。同时,以五阶段SAB为例进行了可视化分析。该模块在空间轴上融合多阶段、多尺度的信息,生成各个阶段的注意图。

首先,对于每一阶段的特征图,在通道维度上分别采用最大池化和平均池化操作,并将其串联得到具有两个通道的特征图。而宽度和高度保持不变。利用共享的扩张卷积运算(扩张率为3,)再使用Sigmod激活函数生成每个阶段的空间注意图。最后将生成的空间注意力图按元素乘上原始特征图,并添加上残差连接得到最终输出。

5)MALUNet

将以上各个模块和UNet网络结合起来得到如下架构:

MALUNet:一种多关注,轻量级的皮肤病变分割UNet_第6张图片

在这个网络结构中采用了六级U型结构,在UNet的基础上本文增加了一个阶段,将1-3阶段的卷积运算从两个减少到一个。

在阶段4到阶段6,我们使用IEA和DGA,它们是串联连接的。前者描述数据集的整体特征和样本之间的关系,后者获取样本中的全局信息和局部信息,此外我们使用SAB对两个桥注意模块进行了顺序排列。这两个注意力桥模块可以融合多个阶段,多个尺度的特征

阶段1-5,生成空间维度和通道维度的注意地图。然后,我们将桥式注意模块获得的特征与解码器部分的特征进行叠加,以减少编码器和解码器之间的特征语义差异,同时减轻采样过程中造成的信息损失。

总结

在本篇论文中,提出了四个注意力模块:1)获取全局和局部特征信息 2)描述这个数据集的样本特征,从而建立样本之间的关系 3)生成一个完整的融合局域和全局多阶段多尺度特征的通道 4)融合多阶段,多尺度的特征信息,得到空间轴上对应的空间注意图

你可能感兴趣的:(UNet各种变体记录,深度学习,人工智能)