卷积神经网络(CNNs)和基于Transformer的模型由于能够提取图像的High-Level特征和捕捉图像的重要方面而被广泛应用于医学图像分割。然而,在对高精度的需求和对低计算成本的期望之间往往存在权衡。具有更高参数的模型理论上可以获得更好的性能,但也会导致更高的计算复杂性和更高的内存使用率,因此实现起来并不实用。
在本文中,作者寻找一种轻量级的基于U-Net的模型,它可以实现几乎相当甚至更好的性能,即U-Lite。作者基于深度可分离卷积的原理设计了U-Lite,这样该模型既可以利用神经网络的强度,又可以减少大量的计算参数。
具体来说,作者提出了在编码器和解码器中都具有7×7的轴向深度卷积,以扩大模型的感受野。为了进一步提高性能,作者使用了几个具有3×3的轴向空洞深度卷积作为作者的分支之一。
总体而言,U-Lite仅包含878K参数,比传统U-Net少35倍。与其他最先进的架构相比,所提出的模型降低了大量的计算复杂性,同时在医学分割任务上获得了令人印象深刻的性能。
1、简介
人工智能(AI)最近在发达国家的许多医院得到了实际应用。据《健康公平》杂志(2018)报道,人工智能有助于有效存储和访问大量信息。目前,医疗信息量每三年翻一番。据估计,如果一名医生想了解所有医学新闻,他必须每天阅读29个小时,这是不可能的。
另一方面,深度学习模型也可以帮助诊断医学图像。IBM Watson医疗保健系统(美国)显示,他们的人工智能系统提出的90%的建议与医学专家的建议一致,但只需40秒即可完成所有流程。意识到人工智能在医学领域的重要性,人们做出了许多研究努力来提高深度学习模型的性能。
计算机视觉是人工智能在医学领域最突出的应用之一。现代医学除了根据临床症状诊断疾病外,还根据从医疗设备获得的图像中的亚临床症状来诊断疾病。因此,开发了深度学习模型来分割肿瘤和细胞或异常区域,从而初步支持医生进行疾病识别和诊断以及疾病的严重程度。
2015年,Ronneberger等人引入U-Net作为一种高效的医学图像分割模型。U-Net成功后,许多后续工作对U-Net进行了不同的优化研究,并提供了许多具有更高性能的变体,如UNet++、ResUNet++、Double UNet、Attention UNet等。总的来说,它们都是基于神经网络开发的深度学习模型。不可否认,神经网络的出现开创了计算机视觉领域的一场伟大革命。
近年来,视觉Transformer和MLP-Like架构(MLP)得到了广泛的应用,并成为计算机视觉中的一个新的事实标准。视觉Transformer将图像的每个Patch都视为一个Token,并通过多头自注意力机制为它们提供信息,就像它们在自然语言处理(NLP)中成功地处理句子一样。
在医学分割任务中,TransUNet可以被认为是精度和效率较高的模型之一。继TransUNet的成功之后,基于Transformer的模型继续被开发。金字塔视觉Transformer(PVT)被用作许多高性能模型的Backbone,如MSMA-Net、Polyp PVT。
同时,MLP-Like的体系结构也是研究的重点。MLP利用传统MLP的优点来沿着其每个维度对特征进行编码。AxialAtt MLP Mixer通过应用轴向注意力力来代替MLP Mixer中的Token混合,在许多医学图像数据集上提供了非常好的性能。与神经网络不同,基于Transformer或MLP的模型主要集中于图像的全局感受野,因此计算复杂度高,训练过程过于繁重。
为了在实践中成功实现,机器学习模型首先需要实现高精度,其次,它应该足够快速和紧凑,以便集成到移动医疗设备中。尽管如此,在对高精度的需求和对低计算成本的期望之间往往存在权衡。上述研究理论上可以取得令人印象深刻的性能,但由于参数数量庞大,其中大量研究可能会带来繁重的运算和缓慢的计算速度。
为了解决这个问题,可以提到一些轻量级架构的尝试,如Mobile UNet、DSCA-Net和MedT。在本文中,作者重新思考了一种用于医学分割任务的高效轻量级架构,以进一步探索一种能够有效解决这一问题的高性能模型。
简而言之,本文的主要贡献有3个方面:
- 基于深度可分离卷积的概念,提出了轴向深度卷积模块的使用方法。该模块帮助模型解决每一个复杂的体系结构问题:扩大模型的感受野,同时减少沉重的计算负担。
- 提出U-Lite,一种基于CNN的轻量级、简单的架构。据作者所知,U-Lite是为数不多的在性能和参数数量方面超过最近高效紧凑型网络UneXt的型号之一。
- 作者已经在医学分割数据集上成功地实现了该模型,并取得了可观的效果。
2、本文方法
作者提出的U-Lite模型的概述如图1所示。作者遵循U-Net的对称编码器-解码器架构,并以一种有效的方式设计U-Lite,以便该模型能够利用CNN的强度,同时保持计算参数的数量尽可能少。
为此,作者深思熟虑地提出了一个轴向深度卷积模块,如图2所示。描述U-Lite的操作,形状为的输入图像通过3个阶段被馈送到网络:编码器阶段、Bottleneck阶段和解码器阶段。U-Lite遵循分层结构,其中编码器提取形状中的6个不同Level的特征,其中。
Bottleneck和解码器参与处理这些特征,并将它们放大到原始形状以获得分割Mask。作者还在编码器和解码器之间使用 skip connections连接。值得注意力的是,尽管U-Lite的设计很简单,但由于轴向深度卷积模块的贡献,该模型在分割任务上仍然表现良好。
2.1、轴向深度卷积模块
视觉Transformer的成功推动了研究和改进这种特殊结构的各种工作。Swin-Transformer通过将自注意力计算限制在大小为7×7的非重叠局部窗口,降低了Transformer的计算复杂性。ConvNext实现了这一修改,并在CNN架构中采用了kernel大小为7×7的卷积,使ResNet在ImageNet上的最高精度达到86.4%。
同时,最近的一个新范式,Vision Permutator利用线性投影来沿着高度和宽度维度分别编码特征表示。这种类似MLP架构的变体被认为很容易在计算机视觉中获得有希望的结果。作者的探索是由一个自然的问题驱动的:
如果作者用局部感受野版本取代视觉Transformer的十字形感受野,就像Swin Transformer对视觉Transformer所做的那样,会发生什么?
为了给出一个简单的答案,作者提出了轴向深度卷积模块,作为Vision Permutator和卷积设计的组合(图2)。该算子的数学公式表示如下:
其中:为输入特征,为输出特征;、和分别代表深度卷积、点卷积和批量归一化,和是卷积的kernel大小;和表示特征图的输入和输出通道的数量。在作者的实验中,=。为了实现最小和灵活的设计,作者使用了一种独特的逐点卷积,而不添加残差连接,允许自适应地改变输入通道的数量。
2.2、编码器块和解码器块
编码器和解码器块的设计原理如下:
- 遵循深度可分离卷积架构。这是从头开始成功构建轻量级模型的重要关键。深度可分离卷积在使用较少参数的同时,提供了与传统卷积相同的性能,从而降低了计算复杂性,使模型更加紧凑。
- 限制使用不必要的操作op。只需使用普通的MaxPooling和UpSampling层。不需要诸如转置卷积之类的高参数消耗算子。逐点卷积算子可以同时扮演两个角色:沿着特征图的深度对特征进行编码,同时灵活地改变输入通道的数量。
- 每个编码器或解码器块采用一个批量标准化层,并以GELU激活功能结束。作者对批处理规范化和层规范化进行了性能比较,但没有太大区别。应用GELU是因为与ReLU和ELU相比,在使用GELU时证明了其在准确性方面的改进。
U-Lite的编码器和解码器结构如图4所示。
2.3、Bottleneck Block
为了进一步提高U-Lite的性能,作者将kernel大小的轴向扩展深度卷积应用于Bottleneck块(图4)。应用的空洞率为。作者使用具有大小为3的kernel的轴向扩张卷积,原因有两个:
- 大小为3大小的kernel更适合底层特征的空间形状,其中这些特征的高度和宽度减少了多次,
- 当使用具有不同空洞率的空洞卷积来捕获后面阶段的High-Level特征的多空间表示时,它给出了更好的性能。
为了进一步减少可学习参数的数量,在Bottleneck块的开头采用了逐点卷积层。这有助于在将最后一层特征提供给轴向扩展深度卷积机制之前缩小其通道尺寸。
3、实验
3.1、SOTA对比
4、参考
[1].1M PARAMETERS ARE ENOUGH? A LIGHTWEIGHT CNN-BASED MODEL FOR MEDICAL IMAGE SEGMENTATION.