TensorFlow入门课程：SAUNet详解

今天将分享Unet的改进模型SAUNet，改进模型来自2020年的论文《Shape Attentive U-Net for Interpretable Medical Image Segmentation》，简单明了给大家分析理解该模型思想。

1、SAUNet网络优点

心血管磁共振成像（CMR）是当前用作非侵入性评估各种心血管功能的金标准。其具有高空间分辨率，无电离辐射，相比核医学成像模式，具有多个优势，例如计算机断层扫描，正电子发射断层扫描（PET）和单光子发射计算机断层扫描（SPECT）。在疾病的诊断和预后治疗中，有很多地方需要估计双心室容积。

双心室体积估计的最大挑战之一是在收缩末期和舒张末期的左右心室心内膜的分割。在临床环境中，高度依赖人工用于估计心室容积。由于详细的轮廓绘制非常耗时。影像学的最新进展和计算急剧增加了在医学成像中使用机器学习。

卷积神经网络（CNN）已显示可学习从多维医学图像中获得抽象表达，学习人类难以定义的功能。这是为什么CNN在对象识别和分割方面表现出色的原因之一。许多神经网络都有一个缺点，包括U-Net缺乏可解释性。

因为这些神经网络与许多卷积层同时连接，将学习到的特征进行可视化变得非常具有挑战性。有效地使神经网络黑匣子，尝试会带来挑战找到错误分类的根本原因，并给出潜在的对抗攻击的优势。此外，CNN受密集像素值的影响很大，与物体的形状相比，这些特征不具有鲁棒性。从而，应该学习对象的形状以允许泛化与模型紧密相关透明度。

缺乏模型透明度和鲁棒性会阻碍其转化到临床环境中。虽然CNN在CMR图像心室分割中显示出很好的结果，但缺乏透明度，在分割过程中会限制这种技术进入临床环境的可解释性。因此迫切需要提高神经网络的透明度，文章提出一种实现透明度和鲁棒性的方法是强制执行模型以学习形状。通过提供更高的分割精度并验证算法不会保留偏差，可以创建工具来帮助解决众多临床医生面对医学图像分析的挑战。

目前为止，过去尝试在医学影像分割中的引入形状信息是使用新的损失函数。此外，关于医学成像模型可解释性的著作已经出版。因此，文章的贡献是：

1）、添加处理的辅助图像形状特征流与U-Net是并行的。建议不要构建新的损失函数，可以在内置在模型中学习形状特征。此外，形状流的输出是形状注意图，可以用于解释。

2）、空间和通道注意路径的使用，在U-Net的每个分辨率下解码器模块中用来对模型学习的特征进行可解释性。

2、SAUNet结构

提出新的可解释的图像分割网络叫形状注意力Unet（SAUNet）。SAUnet由两个流组成：纹理流和门控形状流。纹理流与Unet有相似的结构，但是编码器用密连接模块代替，解码器用双注意力编码模块。

图像信息和特征通过纹理流来学习得到，但这是学习不到形状特征的。门控形状流用来学习物体形状，因此这是可解释的。

Tensorflow入门教程（四十五）——SAUNet_编码器

2.1、门控形状流

门控形状流使用来自编码器的特征来处理和精炼相关的边界和形状信息。使用门控卷积层完成来自形状流的形状信息与纹理特征信息的融合。

根据论文中公式可以得到门控形状流的具体计算过程，如上图所示从E2出来的特征经过1X1卷积和残差模块输出得到Sl形状流特征图，E3的纹理特征经过1x1卷积后再于Sl拼接再经过1x1卷积和sigmoid函数产生alpha1形状注意力系数，将该系数于Sl元素相乘得到注意力卷积层的输出，后面再经过残差模块输出得到Sl+1形状流特征图，E4的纹理特征经过1x1卷积后再于Sl+1拼接经过1x1卷积和sigmoid函数产生alpha1+1形状注意力系数，将该系数于Sl+1元素相乘得到注意力卷积层的输出，后面再经过残差模块输出得到Sl+2形状流特征图，以此类推。

采用深监督方法来对金标准形状边界和预测形状边界进行损计算，损失函数采用二值交叉熵。该模型的目标是正确地学习类别的形状。由于整个门控形状流是可导的，梯度可以回传到纹理流编码器中，因此纹理流编码器还将学习到一些相关的形状信息。门控形状流的输出是感兴趣类别的预测形状特征图，按照通道方向与原始图像的Canny边缘拼接在一起的结果。

在最后一个纹理特征流规范化的3x3卷积层之前，再将该结果与纹理流的特征图进行拼接做为输出结果。

2.2、双注意力编码模块

解码器模块将通过跳跃连接的编码器的特征图与较低分辨率的上下文信息和空间信息的解码器块的特征图融合在一起。为了了解模型在这些模块中检测到的具体特征，以减少模型的黑匣子效应，提出了双注意力解码器块，由拼接特征图上的3x3卷积之后的两个新组件组成。两个新的组成是用于解释的空间注意力路径和用于改善性能的通道注意力路径。

1）、空间注意力路径非常简单，由两个串联1x1卷积构成，第一个1x1卷积将通道数减半，第二个1x1卷积再将通道数从一半降到1，然后经过sigmod函数，得到空间注意力系数，最后按照通道方向将空间注意力系数重复堆叠到输入通道个数得到空间注意力结果输出。

2）、通道注意力路径由squeeze和excitation模块组成，得到通道注意力系数，最后将输入特征图与通道注意力系数相乘就得到了通道注意力结果输出。

3）、通道和空间注意力，通过融合空间注意力和通道注意力的结果做为双注意力编码模块结果。为了放大特征而不是抑制特征，首先对空间注意力结果加1，再与通道注意力结果元素相乘得到结果。

2.3、双任务损失函数用来学习形状

论文提出的目标函数优化了精确的分割并促进了形状的学习。双任务损失由分割损失和形状边界损失构成——分割交叉熵损失，分割dice损失和二值交叉熵形状边界损失加权系数和。

3、实验设置与结果对比

3.1、训练数据采用的是SUN09和AC17分割数据集。数据预处理采用z-score归一化方式。数据增强采用旋转，水平垂直翻转，弹性形变，gamma位移。

3.2、硬件设备是英伟达Quadro RTX5000GPU。

3.3、代码地址 https://github.com/sunjesse/shape-attentive-unet

3.4、实现结果

文章模型在上面两个分割数据集上的结果相比于其它模型的结果，都是最好的。

Tensorflow入门教程（四十五）——SAUNet_卷积_02