UNeXt：基于MLP的快速医学图像分割网络

摘要

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

引言

医学成像解决方案在医疗保健领域的诊断和治疗中发挥着关键作用。医学成像应用中的一个主要任务是分割，因为它对于计算机辅助诊断和图像引导手术系统至关重要。在过去的十年中，许多文献中的作品都集中在开发高效和鲁棒的分割方法。UNet [18]是一项具有里程碑意义的工作，它展示了具有跳过连接的编码器-解码器卷积网络对于医学图像分割的效率。近年来，UNet已经成为几乎所有领先的医学图像分割方法的支柱。在UNet之后，已经提出了一些关键扩展，如UNet++ [30]，UNet 3 + [13]，3D UNet [7]，V-Net [16]，Y-Net [15]和KiUNet [23]。最近，许多基于变换的网络已经被提出用于医学图像分割，因为它们学习图像的全局理解，这有助于分割。TransUNet [6]将ViT架构[10]修改为用于2D医学图像分割的UNet。其他基于变换器的网络，如MedT [22]，TransBTS [26]和UNETR [11]也被提出用于医学图像分割。请注意，几乎所有上述工作都集中在提高网络的性能，但没有太多关注计算复杂性，推理时间或参数的数量，这在许多现实世界的应用中是必不可少的。由于其中大多数用于实验室环境中的分析，因此使用具有高计算能力的机器（如GPU）进行测试。这有助于加快推理的速度，也有助于适应大量的参数。
近年来，医学成像解决方案已经从实验室转移到床边环境。这被称为即时成像，因为测试和分析是在患者身边完成的。床旁成像[24]可帮助临床医生扩展服务选项并改善患者护理。它有助于减少患者不得不前往放射中心的时间和程序。床旁成像技术的改进提高了患者的满意度。近年来，床旁设备的使用一直在增加。例如，床旁超声（POCUS）设备[1]已被证明可用于快速检查肺部胸膜不规则性、心脏血流动力学流量和自动膀胱容量计算。基于电话相机的图像也被用于检测和诊断皮肤状况[2]。磁共振成像（MRI）机器也已开发用于床边操作和快速分析[3]。这些最新的诊断发展有助于在护理点清晰快速地获取医学图像，如图1所示。分割、分类和配准等任务也与这些设备沿着集成，以帮助患者和临床医生加快诊断过程。针对这些任务的主要基于深度学习的解决方案（如UNet和TransUNet）具有固有的计算开销和大量参数，使得它们难以在床旁应用中使用。在这项工作中，我们专注于解决这个问题，并设计了一个高效的网络，具有更少的计算开销，少的参数，更快的推理时间，同时也保持了良好的性能。设计这样的网络对于适应医学成像从实验室到床旁的转变趋势至关重要。为此，我们提出了使用卷积网络和（多层感知器）MLP设计的UNeXt。
最近，基于MLP的网络[14，20，21，28]也被发现能够胜任计算机视觉任务。特别是MLP混频器[20]，一种基于全MLP的网络，其性能与变压器相当，计算量更少。受这些工作的启发，我们提出了UNeXt，这是一个基于卷积和MLP的网络。我们仍然遵循具有跳过连接的UNet的5层深度编码器-解码器架构，但改变了每个块的设计。我们在UNeXt中有两个阶段-卷积阶段和MLP阶段。我们在网络的初始和最终块中使用具有较少滤波器数量的卷积块。在瓶颈中，我们使用了一种新的令牌化MLP（TokMLP）块，它可以有效地保持更少的计算，同时也能够建模一个良好的表示。标记化的MLP将卷积特征投影到抽象标记中，然后使用MLP来学习有意义的分割信息。我们还在MLP中引入移位操作，以提取对应于不同轴向移位的局部信息。由于标记化特征的维数较低，并且MLP比卷积或自注意力和变换器更简单;我们能够显着减少参数数量和计算复杂度，同时保持良好的性能。我们在ISIC皮肤病变数据集[8]和乳腺超声图像（BUSI）数据集[4]上评估了UNeXt，并表明它比最近的通用分割架构获得了更好的性能。更重要的是，与TransUNet相比，我们将参数数量减少了72倍，将计算复杂性降低了68倍，并将推理速度提高了10倍，使其适用于即时医疗成像应用。
总之，本文做出了以下贡献：1）我们提出了UNeXt，这是第一个基于卷积MLP的图像分割网络。2)我们提出了一种新的标记化的MLP块与轴向移位，有效地学习一个很好的表示在潜在的空间。3)我们成功地提高了医学图像分割任务的性能，同时具有更少的参数，高推理速度和低计算复杂度。