SegNetr来啦 | 超越UNeXit/U-Net/U-Net++/SegNet，精度更高模型更小的UNet家族

集智书童公众号

发布于 2023-09-04 11:05:28

8910

发布于 2023-09-04 11:05:28

文章被收录于专栏：集智书童

最近，U-shaped网络由于其简单且易于调整的结构而在医学图像分割领域占据主导地位。然而，现有的U-shaped分割网络：

大多侧重于设计复杂的自注意力模块，以弥补基于卷积运算的远距离依赖性的不足，这增加了网络的总体参数数量和计算复杂度；
简单地融合编码器和解码器的特征，忽略它们的空间位置之间的联系。

在本文中，作者重新思考了上述问题，并构建了一个轻量级的医学图像分割网络，称为SegNetr。具体来说，作者介绍了一种新的SegNetr块，它可以在任何阶段动态执行局部全局交互，并且只有线性复杂性。同时，作者设计了一种通用的 Information Retention Skip Connection（IRSC），以保留编码器特征的空间位置信息，并实现与解码器特征的精确融合。作者在4个主流医学图像分割数据集上验证了SegNetr的有效性，与普通U-Net相比，参数和GFLOP分别减少了59%和76%，同时实现了与最先进方法相当的分割性能。值得注意的是，本文提出的组件也可以应用于其他U-shaped网络，以提高其分割性能。

1、简介

医学图像分割一直是开发自动化辅助诊断系统的关键方面之一，该系统旨在分离医学图像中的对象或结构进行独立分析和处理。通常，分割需要由专业医生手动执行，这既耗时又容易出错。相比之下，开发计算机辅助分割算法可以更快、更准确地用于批量处理。

U-Net表示的方法是一种用于医学图像分割的通用架构，它通过自上而下的编码器路径生成图像的分层特征表示，并使用自下而上的解码器路径将学习到的特征表示映射到原始分辨率，以实现逐像素分类。在U-Net之后，基于卷积神经网络（CNN）的U-shaped方法已扩展到各种医学图像分割任务。它们要么增强编码器-解码器的特征表示能力，要么精心设计注意力模块以关注图像中的特定内容。尽管这些扩展可以改进基准方法，但卷积的局部性质限制了它们捕获长期依赖性，这对医学图像分割至关重要。

最近，在Transformer的推动下，基于U-shaped网络的分割方法发生了重大变化。Chen等人提出了第一个基于Transformer的U-shaped分割网络。Cao等人将Swin Transformer 直接扩展到U-shaped结构。当特征图大小变大时，上述方法遭受高计算和存储器成本爆炸的影响。此外，一些研究人员试图通过结合CNN和Transformer的优势来构建混合网络，如UNeXt、TransFuse、MedT和FAT Net。

与这些工作类似，作者重新设计了基于窗口的局部全局交互，并将其插入到纯卷积框架中，以弥补卷积在捕捉全局特征方面的不足，并降低自注意力操作产生的高计算成本。

Skip Connection是融合U-shaped网络中浅层和深层特征的最基本操作。考虑到这种简单的融合并不能充分利用信息，研究人员提出了一些新的Skip Connection方法。UNet++设计了一系列密集的Skip Connection，以减少编码器和解码器子网络特征图之间的语义差距。SegNet使用最大池化指数来确定位置信息，以避免使用反卷积进行上采样时的模糊性问题。BiO-Net提出了双向Skip Connection，以循环方式重用构建块。UCTransNet设计了一种基于Transformer的通道特征融合方法，以弥合浅层和深层特征之间的语义鸿沟。作者的方法侧重于编码器和解码器的空间位置之间的连接，保留更多的原始特征，以帮助在上采样阶段恢复特征图的分辨率，从而获得更准确的分割图。

通过回顾以上基于U-shaped结构的多个成功案例，作者认为可以通过改进以下两个方面来提高U-shaped网络的效率和性能：

局部-全局交互。网络通常需要处理医学图像中不同大小的目标，而局部全局交互可以帮助网络更准确地理解图像的内容。
编码器-解码器之间的空间连接。使用编码器-解码器之间的空间信息可以获得语义上更强和位置上更准确的特征。

基于以上分析，本文对U-shaped网络的设计进行了重新思考。具体来说，作者构建了轻量级的SegNetr（带Transformer的分割网络）块，以在non-overlapping的窗口上动态学习局部全局信息，并保持线性复杂性。作者提出了Information Retention Skip Connection（IRSC），它专注于编码器和解码器空间位置之间的连接，保留更多的原始特征，以帮助在上采样阶段恢复特征图的分辨率。

总之，本文的贡献可以总结如下：

作者提出了一种计算成本更低、分割性能更好的轻量级U-shaped SegNetr分割网络；
作者研究了传统的U-shaped Skip Connection框架的潜在缺陷，并改进了具有信息保留的Skip Connection；
当作者将本文提出的分量应用于其他U-shaped方法时，分割性能得到了一致的提高。

2、本文方法

如图1所示，SegNetr是一个分层的U-shaped网络，其重要组件包括SegNetr块和IRSC。为了使网络更加轻量级，作者使用MBConv作为基础卷积构建块。SegNetr块在编码器和解码器阶段实现动态局部全局交互。Patch合并用于在不丢失原始图像信息的情况下将分辨率降低两倍。

IRSC用于融合编码器和解码器的特征，随着深度的加深，减少网络丢失的详细信息。请注意，通过更改通道的数量，作者可以获得SegNetr-S的较小版本（C=32）和SegNetr的标准版本（C=64）。

2.1、SegNetr Block

具有全局交互的自注意力机制是Transformer成功的关键之一，但计算整个空间上的注意力矩阵需要二次复杂度。受窗口注意力方法的启发，作者构造了只需要线性复杂性就能实现局部全局交互的SegNetr块。设输入特征图为

X∈R^{H×W×C}

。作者首先使用MBConv提取特征

X_{MBConv}∈R^{H×W×C}

，与通常的卷积层相比，它提供了非显式的位置编码。

局部交互可以通过计算non-overlapping的patch的注意力矩阵（P表示patch大小）来实现。

首先，作者使用无计算成本的局部分区（LP）操作将

X_{MBConv}

划分为一系列空间连续的Patch（

\frac{H×W}{P×P},P,P,C

）（图1显示了P=2的Patch大小）。

然后，作者对通道维度的信息进行平均，并对空间维度进行flatten，以获得（

\frac{H×W}{P×P},P×P

），将其输入到FFN中进行线性计算。由于通道方面的重要性在MBConv中进行了权衡，因此作者在执行局部交互时重点关注空间注意力的计算。

最后，作者使用Softamx来获得空间维度的概率分布，并对输入特征

X_{MBConv}

进行加权。这种方法不仅有利于并行计算，而且更纯粹地关注局部空间的重要性。

考虑到局部交互是不够的，并且可能存在不足的问题，作者还设计了并行的全局交互分支。首先，作者使用全局分区（GP）操作来聚合空间上的非连续Patch。GP将窗口位移的操作添加到LP中，目的是改变特征在空间中的总体分布（图中的全局分支1显示了位移后Patch空间位置的变化）。对于水平方向上的奇数patch，位移规则为向左一个窗口（对于向右的偶数patch，反之亦然），对于垂直方向上的奇patch，向上一个窗口，向下一个窗口。请注意，patch的位移不具有任何计算成本，并且仅发生内存变化。

与Swin Transformer的滑动窗口操作相比，作者的方法在本质上更具全局性。然后，作者将空间移位的特征图分解为2P（

\frac{H×W}{2P×2P},2P,2P,C

）块，并执行全局注意力计算（类似于局部交互分支）。尽管相对于局部交互操作，全局交互在更大的窗口上计算注意力矩阵，但所需的计算量远小于标准自注意力模型的计算量。

局部和全局分支最终通过加权求和进行融合，在此之前，需要通过LP和GP反转操作（即局部反转（LR）和全局反转（GR））来恢复特征图形状。此外，作者的方法还采用了 Transformer 的有效设计，如范数、前馈网络（FFN）和残差连接。

大多数Transformer模型使用固定大小的Patch，但这种方法限制了它们在早期阶段关注更广泛的区域。本文通过应用动态大小的Patch来缓解这个问题。在编码器阶段，作者依次使用（8，4，2，1）的Patch来计算局部注意力，全局分支将Patch扩展到（16，8，4、2）的大小。为了减少超参数设置，解码器的Patch与相应级的编码器Patch具有相同的大小。

2.2、Information Retention Skip Connection

图2显示了3种不同类型的Skip Connection。U-Net在编码器和解码器的相应阶段拼接通道维度，允许解码器在执行上采样时保留更高分辨率的细节信息。SegNet通过在编码器中保留下采样过程的位置信息来帮助解码器恢复特征图分辨率。

作者设计IRSC以考虑这两个特征，即在实现浅特征和深特征融合的同时保留编码器特征的位置信息。具体地说，编码器中的Patch合并（PM）操作将输入特征图

X_{in}∈R^{H×W×C}

的分辨率降低到原始分辨率的2倍，而通道维度扩展到原始维度的4倍，以获得

X_{PM}∈R^{H/2×W/2×4C}

。PM操作的本质是在没有任何计算成本的情况下将空间维度上的信息转换为通道表示，并保留输入特征的所有信息。

IRSC中的Patch反向（PR）用于恢复编码器的空间分辨率，它是与PM的倒数运算。作者交替选择XPM的一半通道数（即

H/2×W/2×2C

）作为PR的输入，一方面可以减少编码器中的冗余特征，另一方面可以对齐解码器中的特征通道数。与传统的上采样方法相比，PR在很大程度上减少了信息丢失的问题，同时提供了准确的位置信息。最后，将PR的输出特征

X_{PR}∈R^{H×W×C/2}

与解码器的上采样特征融合，用于下一阶段的学习。

3、实验

3.1、ISIC2017 和 PH2 SOTA对比

如表1所示，作者将SegNetr与U-Net和其他8种最先进的方法进行了比较。在ISIC2017数据集上，SegNetr和TransUNet获得了最高的IoU（0.775），比U-Net高3.9%。即使是参数数量较少的SegNetr-S也可以获得与其UNeXt-L对应方法类似的分割性能。

通过观察PH2的实验结果，作者发现基于Transformer的方法Swin-UNet分割性能最差，这与目标数据集的数据量直接相关。作者的方法在该数据集上获得了最佳的分割性能，并保持了较低的开销。尽管作者使用了基于窗口位移的注意力方法，但卷积神经网络具有更好的归纳偏差，因此与Swin-UNet或TransUNet等基于Transformer的方法相比，对数据量的依赖性更小。

如表2所示，SegNetr的IoU和Dice分别比双编码器FATNet高1.6%和0.8，而GFLOP则低32.65。在ACDC数据集中，左心室更容易分割，U-Net的IoU为0.861，但比SegNetr差1.1%。心肌位于左心室和右心室的中间，呈环状，作者的方法比专注于边界分割质量的EANet高0.6%的IoU。

此外，作者观察了4个网络UNeXt、UNeXt-L、SegNetr-S和SegNetr的分割性能，发现较小的参数可能会限制网络的学习能力。本文提出的方法在所有4个数据集上都显示出有竞争力的分割性能，表明作者的方法具有良好的泛化性能和鲁棒性。