Tensorflow入门教程（四十一）——SAU-Net

今天将分享Unet的改进模型SAU-Net，改进模型来自2020年的论文《SAU-Net Effient 3D Spine MRI Segmentation Using Inter-Slice Attention》，简单明了给大家分析理解该模型思想。

1、SAU-Net优点

医疗3D分割一般是通过3D卷积神经网络（CNN）利用空间信息来进行连续和准确的分割。但是3D CNN受限于更高的计算成本，内存和过拟合的风险，特别是对于医疗而言标记数据数量是有限。为了解决这些问题，论文作者采用注意力机制来利用切片间信息并基于2D卷积网络，提出基于空间注意力的密连接的U-Net（SAU-Net）。首先初始3D图像分解为一堆2D切片，然后通过密连接的U-Net结构来获取切片内粗糙的分割概率结果，最后，附加了切片间注意模块以捕获和融合3D切片间空间信息和2D上下文信息，以此来完善3D分割结果。论文实验结果表明了ISA的有效性以及与其他深度学习方法相比分割结果有更高准确性和有效性。论文主要贡献如下所述：

A、运用注意力机制在基于2D卷积网络的3D分割任务中利用切片间信息并提出切片间注意模块（ISA）。

B、提出了一种新颖的结构，称为基于空间注意力的密连接U-Net（SAU-Net）从3D MR图像进行有效和准确的脊柱分割。

2、SAU-Net结构

SAUNet结构如下所示，首先将3D体数据按照深度分成多个2D序列图像，然后用密连接的Unet网络来获取切片内信息得到每个切片大致分割概率结果，最后通过切片间注意力模块对2D上下文信息来捕获和融合3D空间信息，产生最后的精致分割结果。

2.1、密连接Unet结构

卷积层之前和之后的短连接可以合并上下文的信息，使得卷积网络更加准确和高效。在经典Unet结构上引入密连接模块可以更好地从不同切片中提取更好的特征。下采样模块包括卷积层，BN层，Relu，密连接和最大池化层，上采样模块包括拼接层，卷积层，BN层，Relu，反卷积层。

2.2、切片间注意力模块（ISA）

由于空间连续性，各层的分割结果在空间上与上层和下层是相关的。因此，相邻切片的信息用于优化分割图像的单个切片是有用的。为了解决这个问题，提出了切片间注意模块（ISA），利用相邻切片之间的上下文信息来增强分割结果的连续性。对于分割任务，通常通过在特征图上创建表示信息区域的模板来获取注意力，从而突出显示最突出的区域并抑制无关的区域。ISA结构如下所示，为了利用空间连续性信息，使用相邻切片的特征图通过生成注意力模板并融合到切片的特征图中来进行分割，从而获得了精确的分割结果。

ISA细节结构如下所示，例如对于切片i的分割，切片i + 1和切片i-1的特征图（第一个和最后一个切片只有一个相邻切片)通过使用1×1卷积和sigmoid函数生成注意力模板。生成的模板表示基于相邻切片信息进行切片i分割的注意力区域。然后将注意力模板与切片i特征图融合。模板的大小与切片i的特征图相同，通过切片i特征图对应模板元素相乘并将相乘后的结果与切片i特征图相加，得到注意力融合输出，在模型的最后使用sigmod函数得到最后的精细分割结果。对于第一个和最后一个切片，注意力融合的过程减半。

3、实验与结果

3.1、数据与评价指标

采用SpineT2W数据集来训练和验证，包含195例T2加权脊柱MR图像，采用开源的MR数据集来测试，包含23例T2加权MR图像。

Dice相似系数，Jaccard系数，Hausdorff距离，精度和召回率。

3.2、实现细节

采用TeslaV100GPU（32G显存），采用二值交叉熵损失函数，Adam优化器，学习率为0.0001.实验结果通过5折交叉验证的方式来验证。

3.3、实验对比

与2DUnet和3Dunet结果进行比较，SAUNet在所有的指标上都是最好的结果。