摘要
在本文中提出了一种名为空间注意力UNet(SA-UNet)的轻量级网络,它不需要大量的带注释的训练样本,可以以数据增强的方式来更有效的使用可用的带注释的样本。SA-UNet引入一种空间注意模块,这个模块沿着空间维度推断注意图,并将注意图与输入特征图相乘,进行自适应特征性细化。此外,这个网络采用结构化的Dropout卷积块代替U-Net的原始卷积块,防止过拟合。
尽管UNet变体的性能很好,但它们不可避免的使得网络更复杂。为了解决这些问题,在UNet中加入了空间注意力,并提出了一个轻量级网络模型称之为SA-UNet。SA-UNet首先采用了一种集成DropBlock和批处理归一化(BN)的结构化dropout卷积块的变体来取代原有的U-Net卷积块。
方法
网络结构
网络结构如下图所示:
编码器的每一步都包括一个结构化的Dropout卷积块和2x2的最大池化操作。每个卷积块的卷积层之后是DropBlock,批归一层(BN)和整流线性单元(ReLU),然后利用最大池化操作进行下采样,步幅大小为2。在每个下采样步骤中,我们将特征通道的数量增加一倍。解码器中的每一步都包括2×2转置卷积操作,用于上采样和减半特征通道数量,与编码器中的相应特征映射连接,然后是结构化的dropout卷积块。在编码器和解码器之间增加了空间注意模块。在最后一层,使用1×1卷积和Sigmoid激活函数来获得输出分割映射。
结构化Dropout卷积块
DropBlock是Dropout的一种结构化形式,可以有效的防止卷积网络中的过拟合问题。它与dropout的主要区别在于它丢弃了一个层的特征图中的连续区域,而不是丢弃独立的随机单位,在这个基础上,本文构造了一个结构化的Dropout卷积块,也就是在每个卷积层后面都有一个DropBlock,一层批处理归一化以及一个Relu激活函数。
空间注意力模块(SAM)
引入空间注意力模块作为卷积注意模块的一部分,空间注意力利用特征之间的空间关系生成空间注意图,为了计算空间注意图,空间注意力SA首先沿通道轴应用最大池化和平均池化操作,并将他们连接起来生成一个高效的特征描述符。如下图所示:
经过最大池化和平均池化处理之后得到输出特征和然后使用上述的卷积层和Sigmod激活函数来生成空间注意力图总之空间注意模块的输出特征用公式表示为
式子中的代表核大小为7x7的卷积运算
总结
本文使用了一个集成了DropBlock和批处理的结构化dropout卷积块取代了U-Net的卷积块。此外,在视网膜眼底图像中,血管区域与背景的区别并不明显,尤其是边缘和小血管。为了帮助网络学习这些,我们在主干网的编码器和解码器之间添加了一个空间注意力模块,并提出了空间注意力U-Net (SAUNet)。空间注意可以帮助网络集中在重要的特征上,抑制不必要的特征,提高网络的表示能力。