UNet家族迎来最小模型U-Lite | 800K参数实现性能极限超车

集智书童公众号

发布于 2023-09-04 10:45:39

1.2K0

发布于 2023-09-04 10:45:39

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

卷积神经网络（CNNs）和基于Transformer的模型由于能够提取图像的High-Level特征和捕捉图像的重要方面而被广泛应用于医学图像分割。然而，在对高精度的需求和对低计算成本的期望之间往往存在权衡。具有更高参数的模型理论上可以获得更好的性能，但也会导致更高的计算复杂性和更高的内存使用率，因此实现起来并不实用。在本文中，作者寻找一种轻量级的基于U-Net的模型，它可以实现几乎相当甚至更好的性能，即U-Lite。作者基于深度可分离卷积的原理设计了U-Lite，这样该模型既可以利用神经网络的强度，又可以减少大量的计算参数。具体来说，作者提出了在编码器和解码器中都具有7×7的轴向深度卷积，以扩大模型的感受野。为了进一步提高性能，作者使用了几个具有3×3的轴向空洞深度卷积作为作者的分支之一。总体而言，U-Lite仅包含878K参数，比传统U-Net少35倍。与其他最先进的架构相比，所提出的模型降低了大量的计算复杂性，同时在医学分割任务上获得了令人印象深刻的性能。代码：https://github.com/duong-db/ULite

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,H,W)

的输入图像通过3个阶段被馈送到网络：编码器阶段、Bottleneck阶段和解码器阶段。U-Lite遵循分层结构，其中编码器提取形状

(C_I,\frac{H}{2^i},\frac{W}{2^i})

中的6个不同Level的特征，其中

i∈\{0,1,2,…,5\}

。

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）。该算子的数学公式表示如下：

其中：

为输入特征，

为输出特征；

、

和

分别代表深度卷积、点卷积和批量归一化，

1×n

和

n×1

是卷积的kernel大小；

C_1

和

C_2

表示特征图的输入和输出通道的数量。在作者的实验中，

n＝7

。为了实现最小和灵活的设计，作者使用了一种独特的逐点卷积，而不添加残差连接，允许自适应地改变输入通道的数量。

2.2、编码器块和解码器块

编码器和解码器块的设计原理如下：

遵循深度可分离卷积架构。这是从头开始成功构建轻量级模型的重要关键。深度可分离卷积在使用较少参数的同时，提供了与传统卷积相同的性能，从而降低了计算复杂性，使模型更加紧凑。
限制使用不必要的操作op。只需使用普通的MaxPooling和UpSampling层。不需要诸如转置卷积之类的高参数消耗算子。逐点卷积算子可以同时扮演两个角色：沿着特征图的深度对特征进行编码，同时灵活地改变输入通道的数量。
每个编码器或解码器块采用一个批量标准化层，并以GELU激活功能结束。作者对批处理规范化和层规范化进行了性能比较，但没有太大区别。应用GELU是因为与ReLU和ELU相比，在使用GELU时证明了其在准确性方面的改进。

U-Lite的编码器和解码器结构如图4所示。