UNeXt: MLP-based Rapid Medical Image Segmentation Network

UNeXt: MLP-based Rapid Medical Image Segmentation Network
来源：MICCAI 2022
文章为本人学习过程中翻译的相关文献，非本人发表。此外，如有侵权，请联系删除博客。
原文链接

Abstract

近年来，UNet及其最新扩展 (如TransUNet) 一直是领先的医学图像分割方法。但是，这些网络不能有效地用于护理应用中的快速图像分割，因为它们参数繁重，计算复杂且使用缓慢。为此，我们提出了UNeXt，它是一种基于卷积多层感知器 (MLP) 的图像分割网络。我们以一种有效的方式设计了UNeXt，其中具有早期卷积阶段和潜在阶段的MLP阶段。我们提出了一个tokenized的MLP块，其中我们有效地对卷积特征进行标记和投影，并使用MLP对表示进行建模。为了进一步提高性能，我们建议在输入到MLP的同时转移输入的通道，以便专注于学习本地依赖性。在潜空间中使用标记化的MLP减少了参数的数量和计算复杂性，同时能够产生更好的表示形式来帮助分割。网络还包括各级编码器和解码器之间的跳过连接。我们在多个医学图像分割数据集上测试了UNeXt，结果表明，我们将参数数量减少了72x，将计算复杂度降低了68倍，并将推理速度提高了10倍，同时还获得了比最先进的医学图像分割体系结构更好的分割性能。代码可在https://github.com/jeya-maria-jose/UNeXt-pytorch获得。

Keywords: Medical Image Segmentation, MLP, Point-of-Care

Introduction

医学成像解决方案在医疗保健领域的诊断和治疗中发挥了关键作用。医学成像应用中的一项主要任务是分割，因为它对于计算机辅助诊断和图像引导手术系统至关重要。在过去的十年中，文献中的许多作品都集中在开发有效且可靠的分割方法上。UNet [17] 是一项具有里程碑意义的工作，它展示了具有跳过连接的编码器-解码器卷积网络如何有效地进行医学图像分割。近年来，UNet已成为几乎所有领先的医学图像分割方法的骨干。继UNet之后，已经提出了许多关键扩展，例如UNet++ [29]，UNet3+ [13]，3D UNet [7]，V-Net [16]，Y-net [15] 和KiUNet [21,22]。最近，已经提出了许多基于transformer的网络用于医学图像分割，因为它们学习了对图像的全局理解，这可能有助于分割。TransUNet [6] 将ViT架构 [10] 修改为用于2D医学图像分割的UNet。还提出了其他基于transformer的网络，例如MedT [20]，TransBTS [25] 和UNETR [11]，用于医学图像分割。请注意，几乎所有上述工作都集中在提高网络性能上，但并没有过多关注计算复杂性，推理时间或参数数量，这在许多实际应用中都是必不可少的。由于其中大多数用于实验室环境中的分析，因此使用具有高计算能力的机器 (例如GPU) 对它们进行了测试。这有助于加快推理的速度，也有助于容纳大量的参数。

最近，医学成像解决方案已从实验室转换为床侧设置。这被称为即时成像，因为测试和分析是在患者身边进行的。护理点成像 [23] 帮助临床医生扩大服务选择和改善患者护理。它有助于减少患者必须去放射学中心的时间和程序。围绕护理点成像的技术改进正在提高患者满意度。近年来，即时护理设备的使用一直在增加。例如，即时超声 (POCUS) 设备 [1] 已显示出可用于快速检查肺部胸膜不规则性，心脏血流动力学和自动膀胱体积计算。基于Phonecamera的图像也用于检测和诊断皮肤状况 [2]。磁共振成像 (MRI) 机器也已开发用于床侧操作和快速分析 [3]。这些最近的诊断发展有助于在护理点清晰而快速地获取医学图像，如图1所示。分割、分类和配准等任务也与这些设备一起被整合，以帮助患者和临床医生加速诊断过程。这些任务的主要基于深度学习的解决方案 (如UNet和tranunet) 具有固有的计算开销和大量参数，使它们难以在护理应用中使用。在这项工作中，我们专注于解决这个问题，并设计了一个有效的网络，该网络具有更少的计算开销，更少的参数数量，更快的推理时间，同时还保持了良好的性能。设计这样的网络对于适应医学成像从实验室到床侧的变化趋势至关重要。为此，我们提出了使用卷积网络和 (多层感知器) MLPs设计的UNeXt。

Fig.1 UNeXt的动机: 随着医学成像解决方案在护理点变得越来越适用，重要的是要专注于使深度网络轻巧，快速，同时又有效。护理点医疗干预工作流程。(b) 最近的医学成像发展: POCUS设备 [1] 和 © 基于电话的皮肤病变检测和识别应用 [2]。

最近，基于MLP的网络 [27,19，14,18] 也被发现可以胜任计算机视觉任务。特别是MLP-Mixer [18]，这是一种基于全MLP的网络，相对于具有较少计算量的变压器，其性能相当。受这些工作的启发，我们提出了UNeXt，这是一个基于卷积和MLP的网络。我们仍然遵循带有跳过连接的UNet的5层深度编码器-解码器体系结构，但更改了每个块的设计。我们在UNeXt中有两个阶段-卷积阶段，然后是MLP阶段。我们在网络的初始和最终块中使用具有较少过滤器数量的卷积块。在瓶颈中，我们使用了一种新颖的Tokenized MLP (TokMLP) 块，该块在保持较少的计算量的同时还能够对良好的表示进行建模。标记化的MLP将卷积特征投影到抽象token中，然后使用MLP学习有意义的信息进行分割。我们还在MLPs中引入了移位操作，以提取与不同轴向移位相对应的局部信息。由于标记化特征的尺寸较小，并且MLP比卷积或自我注意和transformer的复杂性较小; 我们能够显着减少参数数量和计算复杂性，同时还保持良好的性能。我们在ISIC皮肤病变数据集 [8] 和乳腺超声图像 (BUSI) 数据集 [4] 上评估了UNeXt，并表明它比最近的通用分割架构具有更好的性能。更重要的是，与TransUNet相比，我们将参数数量减少了72x，将计算复杂度降低了68倍，并将推理速度提高了10倍，使其适用于护理点医学成像应用。

总而言之，本文做出了以下贡献: 1) 我们提出了第一个基于卷积MLP的图像分割网络UNeXt。2) 我们提出了一种新颖的具有轴向位移的标记化MLP块，以有效地在潜空间上学习良好的表示。3) 我们成功地提高了医学图像分割任务的性能，同时具有较少的参数，较高的推理速度和较低的计算复杂度。

UNeXt

Network Design

UNeXt是具有两个阶段的编码器-解码器体系结构: 1) 卷积阶段和2) Tokenized MLP阶段。输入图像通过编码器，其中前3个块是卷积的，接下来的2个是标记化的MLP块。解码器具有2个标记化的MLP块，然后是3个卷积块。每个编码器块将特征分辨率降低2，并且每个解码器块将特征分辨率提高2。编码器和解码器之间还包括跳过连接。每个块上的通道数是表示为C1到C5的超参数。对于使用UNeXt架构的实验，除非另有说明，否则我们遵循C1 = 32，C2 = 64，C3 = 128，C4 = 160和C5 = 256。请注意，这些数字实际上少于UNet及其变体的过滤器数量，这些过滤器的数量有助于减少参数和计算。

Fig.2 拟议的UNeXt架构概述。

Convolutional Stage

每个卷积块都配备有卷积层，批归一化层和ReLU激活。我们使用3 × 3的内核大小，步幅为1，填充为1。编码器中的conv块使用具有池化窗口2 × 2的最大池化层，而解码器中的conv块由双线性插值层组成，以对特征图进行上采样。我们使用双线性插值代替转置卷积，因为转置卷积基本上是可学习的上采样，并有助于更多可学习的参数。

Shifted MLP

在Shifted MLP中，我们首先在tokenize之前移动conv特征的通道轴。这有助于MLP仅专注于conv特征的某些位置，从而诱导块的定位。这里的目的类似于Swin transformer [5] ，其中引入了基于窗口的注意力，以向其他完全全局的模型添加更多的局部性。由于torkenized MLP块具有2个MLP，因此我们将特征先在宽度上移动，然后在高度上移动，就像在轴向注意中一样 [24] 。我们将特征划分为h个不同的分区，并根据指定的轴将它们移动j个位置。这有助于我们创建沿轴引入位置的随机窗口。

Fig.3 Shift操作。特征在tokenize之前在宽度和高度上顺序移动，以诱导网络中的窗口局部性。

Tokenized MLP Stage

在tokenized MLP块中，我们首先将特征移位并将其投影到token中。为了tokenize，我们首先使用3的内核大小，并将通道数更改为E，其中E是嵌入维度 (token数)，这是一个超参数。然后，我们将这些token传递到Shifted MLP (跨宽度)，其中MLP的隐藏尺寸是超参数H。接下来，将特征传递通过深度方向的卷积层 (DWConv)。我们在此块中使用DWConv有两个原因: 1) 它有助于对MLP特征的位置信息进行编码。[26] 中显示，MLP块中的Conv层足以对位置信息进行编码，并且实际上比标准位置编码技术表现更好。当测试和训练分辨率不同时，需要对ViT中的位置编码技术进行插值，通常会导致性能下降。2) DWConv使用较少的参数，因此提高了效率。然后，我们使用GELU [12] 激活层。我们使用GELU代替ReLU，因为它是一种更平滑的替代方案，并且被发现性能更好。此外，最近的大多数架构如ViT [10] 和BERT [9] 都成功地使用GELU获得了改进的结果。然后，我们将特征传递到另一个Shifted MLP (跨高度)，该MLP将尺寸从H转换为O。我们在这里使用残差连接，并将原始token添加为残差。然后，我们应用层归一化 (LN) 并将输出特征传递到下一个块。与BN相比，LN是优选的，因为它更有意义地沿着令牌进行归一化，而不是跨令牌化MLP块中的批次进行归一化。

Tokenized MLP块中的计算可以概括为:

其中T表示tokens，H表示高度，W表示宽度，DWConv表示深度卷积，LN表示层归一化。请注意，所有这些计算都是在嵌入维度H上执行的，该嵌入维度H明显小于特征图H/N × H/N的维度，其中N是取决于块的2的倍数。在我们的实验中，除非另有说明，否则我们将H设置为768。这种设计tokenized MLP块的方式有助于编码有意义的特征信息，并且在计算或参数方面没有太大贡献。

Experiments and Results

Datasets

为了使我们的实验尽可能接近护理点成像，我们选择国际皮肤成像协作 (ISIC 2018) [8] 和乳房超声图像 (BUSI) [4] 数据集来基准我们的结果。ISIC数据集包含相机获取的皮肤病学图像和相应的皮肤病变区域分割图。ISIC 2018数据集由2594图像组成。我们将所有图像的大小调整为512 × 512的分辨率。BUSI由正常，良性和恶性乳腺癌病例的超声图像以及相应的分割图组成。我们仅使用良性和恶性的图像，这导致总共647个图像的大小调整为256 × 256的分辨率。

Implementation Details

我们使用Pytorch框架开发了UNeXt。我们使用二进制交叉熵 (BCE) 和骰子损失的组合来训练UNeXt。预测 “y” 与目标y之间的损失L公式化为:

我们使用具有0.0001学习率和0.9动量的Adam优化器。我们还使用余弦退火学习率调度器，其最小学习率高达0.00001。批量大小设置为等于8。我们训练UNeXt总共400个时期。我们在数据集上进行了三次80-20随机分裂，并报告了均值和方差。

Performance Comparison

我们将UNeXt的性能与最近和广泛使用的医学图像分割框架进行了比较。特别是，我们与卷积基线进行了比较，例如UNet [17] ，UNet [29] 和recunet [28] 。我们还与transunet [6] 和MedT [20] 等最近的transformer基线进行了比较。请注意，我们专注于在分割性能 (F1分数和IoU) 以及参数数量，计算复杂度 (以GFLOPs为表示) 和推理时间 (以ms为表示) 方面与基线进行比较。

我们将结果列在表1中。可以观察到，与所有基线相比，UNeXt获得了更好的分割性能，而紧随其后的是transeunet。$p\le 10^{-5}$时，改善具有统计学意义。但是，这里要注意的最引人注目的一点是，与transunet相比，UNeXt的计算数量非常少，因为UNeXt没有任何注意力块。计算是根据浮点运算符 (flop) 的数量来计算的。我们注意到，与transunet的38.52和UNet的55.84相比，UNeXt的GFLOPs最少，为0.57。与所有基线相比，它也是最轻量级的网络。特别是，我们注意到，与transunet的105.32 M个参数相比，UNeXt仅具有1.58 M个参数。我们还介绍了在CPU上运行时的平均推理时间。请注意，我们已经在CPU而不是GPU中特别基准化了推理时间，因为护理点设备大多以低计算能力运行，并且通常不具有GPU的计算优势。我们对10个分辨率为256 × 256的图像进行前馈，并报告平均推理时间。用于基准标记的CPU是运行在2.30 GHz的Intel Xeon Gold 6140 CPU。可以注意到，我们对Swin-UNet [5] 进行了实验，但发现在小型数据集上存在收敛问题，导致性能不佳。然而，Swin-UNet具有41.35 M个参数，并且在11.46 GFLOPs的计算上也很复杂。

在图4中，我们绘制了F1-Score与GLOPs，F1-Score与推理时间以及F1-Score与参数数量的比较图。这里使用的F1-Score对应于ISIC数据集。从图表中可以清楚地看出，就细分性能而言，UNeXt和transunet是性能最佳的方法。但是，UNeXt在计算复杂性，推理时间和参数数量方面明显优于所有其他网络，这些都是护理点成像应用需要考虑的重要特征。在图5中，我们给出了UNeXt的样本定性结果以及其他基线。可以观察到，与其他方法相比，UNeXt产生了竞争性的细分预测。

Fig.4 比较图表。Y轴对应F1-Score (越高越好)。X轴对应于GFLOPs、推理时间和参数数量 (越低越好)。可以看出，与其他网络相比，UNeXt是最有效的网络。

Fig.5 定性比较。第1行-ISIC数据集，第2行-BUSI数据集。(a) 输入。(b) UNet © UNet (d) MedT (e) trancunet (f) UNeXt和 (g) ground truth的预测。

Discussion

Ablation Study

我们进行消融研究 (如表2所示)，以了解UNeXt中每个模块的个体贡献。我们首先从原始的UNet开始，然后只是减少过滤器的数量，以减少参数的数量和复杂性。我们看到性能下降，参数减少不多。接下来，我们减少深度并仅使用3级深度体系结构，这基本上是UNeXt的Conv阶段。这显著降低了参数的数量和复杂性，但也降低了4% 的性能。现在，我们介绍了tokenized MLP块，该块显着提高了性能，同时将复杂度和参数提高了最小值。接下来，我们添加 [26] 中使用DWConv的位置嵌入方法，并看到更多改进。接下来，我们在MLPs中添加移位操作，并表明在tokenize之前对功能进行移位可以提高性能，而无需添加任何参数或复杂性。由于移位操作对任何加法或乘法都没有贡献，因此不会对任何触发器进行添加。我们注意到，在两个轴上移动特征可获得最佳性能，这是具有最小的参数和复杂性的UNeXt的精确配置。请注意，以上所有实验都是使用ISIC数据集的单个折叠进行的。

Analysis on number of channels

通道数是UNeXt的主要参数，它影响参数的数量，复杂性和网络性能。在表3中，我们对ISIC的单折叠进行了实验，以显示UNeXt的两种更不同的配置。可以观察到，增加通道 (UNeXt-L) 进一步提高了性能，同时增加了计算开销。尽管降低它 (UNeXt-S) 会降低性能 (降低并不严重)，但是我们得到了一个非常轻巧的模型。

Difference from MLP-Mixer

MLP-Mixer使用全MLP架构进行图像识别。UNeXt是一种基于卷积和MLP的图像分割网络。MLP-Mixer专注于通道混合和token混合，以学习良好的表示。相反，我们提取卷积特征，然后对通道进行标记化，并使用新颖的tokenized MLP使用shifted MLP对表示进行建模。值得注意的是，我们尝试使用MLPMixer作为编码器和普通卷积解码器。性能对于分割不是最佳的，并且11 M左右的参数仍然很大。

Conclusion

在这项工作中，我们提出了一种新的深度网络体系结构UNeXt，用于医学图像分割，重点是护理点应用。UNeXt是一种基于卷积和MLP的体系结构，其中在潜在空间中有一个初始的conv阶段，随后是MLP。具体来说，我们提出了一个具有移位MLP的tokenized MLP块，有效地以最小的复杂性和参数对表示进行建模。我们在多个数据集上验证了UNeXt，在那里我们实现了更快的推理，降低了复杂性和更少的参数数量，同时还实现了最先进的性能。