11

AI学习者 · 2023年07月14日

SegNetr来啦 | 超越UNeXit/U-Net/U-Net++/SegNet

image.png

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

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

在本文中,作者重新思考了上述问题,并构建了一个轻量级的医学图像分割网络,称为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网络的效率和性能:

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

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

总之,本文的贡献可以总结如下:

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

2、本文方法

image.png

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

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

2.1、SegNetr Block

image.png

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

image.png

局部和全局分支最终通过加权求和进行融合,在此之前,需要通过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

image.png

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

image.png

3、实验

3.1、ISIC2017 和 PH2 SOTA对比

image.png

如表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的方法相比,对数据量的依赖性更小。

image.png

如表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个数据集上都显示出有竞争力的分割性能,表明作者的方法具有良好的泛化性能和鲁棒性。

image.png

此外,图3提供了定性示例,证明了作者提出的方法的有效性和稳健性。结果表明,SegNetr能够用较少的数据准确描述皮肤病变,并实现了最小化欠分割和过分割的多类分割。

3.2、消融实验

1、局部-全局交互作用的影响

image.png

局部全局交互在SegNetr中的作用可以从表3中理解。在没有局部或全局交互的情况下,网络的总体参数较少,但分割性能也受到很大影响。通过添加局部或全局交互,提高了网络对不同类别的分割性能。此外,通过串行和并行运行局部全局交互模块可以获得类似的性能,但串行连接导致计算效率较低,并影响运行速度。

2、patch size的影响

image.png

如表4(左)所示,不同的Patch大小显著影响模型的效率和参数。当在每个阶段中使用大小为2的Patch时,参数的数量达到54.34M,与使用大小为(8,4,2,1)的动态Patch相比,增加了42.08M。基于这项消融研究,作者建议在不同阶段使用[分辨率14]patch尺寸。

3、IRSC的影响

image.png

表4(右)显示了用IRSC替换UNeXit、U-Net、U-Net++和SegNet的Skip Connection的实验结果。这些方法在IRSC的帮助下得到了持续的改进,这清楚地表明了IRSC的有用性。

4、参考

[1].SegNetr: Rethinking the local-global interactions and skip connections in U-shaped networks.

作者: 小书童
文章来源:集智书童

推荐阅读

更多嵌入式AI干货请关注嵌入式AI专栏。欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
18756
内容数
1315
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息