MobileUNETR 一种用于高效图像分割的轻量级端到端混合视觉 Transformer !

图像处理机器学习数据库

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

移动UNETR是一种轻量级端到端混合视觉Transformer,旨在克服CNN和Transformer之间的性能限制,同时最小化模型大小,从而朝着高效图像分割迈出有前景的一步。移动UNETR的三个主要特点如下:

  1. 移动UNETR包含一个轻量级的混合CNN-Transformer编码器,旨在有效平衡局部和全局上下文特征提取;

  2. 一种新颖的混合解码器,在解码阶段同时利用不同分辨率的低级和全局特征进行精确分割;

  3. 超越大型和复杂的结构,移动UNETR在仅300万参数和1.3 GFLOP计算复杂度下实现卓越性能,分别比参数量和FLOPS减少了10倍和23倍。

在四个公开可用的皮肤病变分割数据集上进行了大量实验,包括ISIC 2016,ISIC 2017,ISIC 2018和PH2数据集。

代码将在https://github.com/OSUPCVLab/MobileUNETR.git上公开提供。

1 Introduction

皮肤癌是世界范围内最普遍且快速增长的癌症类型之一,对全球健康构成重大挑战[55]。由于患者之间皮肤癌的多种表现形式、不同程度以及主观视觉过程的影响,准确识别和分类皮肤病变变得非常复杂。诊断这类癌症的主要困难在于病变的视觉检查。视觉过程的主观性受照明条件、个人专业知识以及皮肤癌在不同患者中的固有变量变化的影响,使其成为一个困难任务。为了提高诊断精确度,皮肤科医生使用皮肤表面显微镜检查非侵入性技术。尽管这一进步无疑提高了人类视觉分析的准确性,但并未完全消除与人类主观性相关的挑战。尽管有皮肤镜图像的帮助,皮肤科医生在解释皮肤病变方面仍然可能存在差异。医务行人之间对皮肤病变诊断的不一致强调了需要提供客观和标准化评估的额外工具。

CAD系统的设计用于皮肤癌分段最初是通过对复杂的多步骤图像处理 Pipeline 进行方法,。这些早期的方法包括颜色空间转换、主成分分析以及使用手工艺特性,仅举几例。尽管这些方法在医学诊断方面取得了进步,但它们在准确界定病变区域方面仍存在困难。基于规则和手工艺的系统通常过于简化复杂的、可变形的皮肤病变,包括来自体毛的伪影和噪声。

深度学习的发展和采用是增强CAD系统效率和准确性的关键步骤。这些系统利用先进的神经网络来界定病变的边界,允许更精确地评估它们的特点。深度学习算法具有直接从数据中自动学习复杂模式和特征的能力,在分割皮肤病变方面表现出色。这些算法能够辨别颜色、纹理和形状的微小变化,随着不同个体之间皮肤癌多样性的变化而动态适应。

深层学习在医学图像分割领域的成功离不开编码器-解码器架构的引入。通过全卷积神经网络(FCNNs)实现的编码器-解码器架构在许多分割任务中表现出色,成为了当前最先进(SOTA)的状态。尽管这些方法取得了很高的成功,但FCNN/CNN方法的一个主要缺点是它们的长期上下文理解能力不足。尽管卷积神经网络在捕捉图像内的局部特征方面表现出色,但它们天生难以收集更广泛的信息或获取不同元素之间的全局关系。特别是在皮肤癌的案例中,病变在患者之间有很大差异,全局理解对于帮助模型解决歧义至关重要。

为了解决CNN的上下文限制,研究行人采取了更大、更深层次的模型来通过纯粹的卷积提高整体感受野,从而改进全模型 。然而,这也带来了一系列挑战。更大模型需要更多的计算资源,使其既昂贵又难以训练和部署。此外,通过扩大模型大小仅仅获得更大感受野可能带来收益递减,强调了需要更有效的方法。将Transformer [63]架构中引入的自我注意力机制与卷积层相结合,被提出作为一种增强非局部建模能力的方法,并为许多下游任务提供了期望的长程上下文理解优势。

最初用于自然语言处理(NLP),Transformer架构已经在许多计算机视觉任务中大规模应用。随着最初允许Transformer执行图像分类的视觉Transformer [12]的出现,研究行人得到了一种能够捕捉长期依赖性和在模型每个阶段收集全局上下文的架构。然而,为了换取自注意力机制,Transformer架构的核心,在计算上过于昂贵的部分。此外,ViTs产生的单一尺度的特征与CNN模型通常生成的多尺度特征形成对比。在资源受限的实时应用中采用Transformer架构,全球化意识和计算效率之间的权衡带来了一个重大的挑战。

为了克服广泛采用CNN和Transformer架构中的瓶颈,作者提出 MobileUNETR,一种新颖的端到端Transformer基于编码器-解码器架构,用于高效的图像分割。在最高层,复杂和具有挑战性的图像分割任务往往从考虑特征提取阶段内的局部和全局上下文信息中受益。

然而,分割方法通常仅关注优化特征提取器,而忽视了开发新的解码策略的重要性。在医学影像中使用的方法通常利用复杂的CNN和/或Transformer结构,并一般排斥基于Transformer的解码器,而选择纯CNNs 。这个选择可以归因于,尽管它们在获取全局信息方面做得很好,但Transformer却无法获取复杂的局部细节,这在生成准确的分割 Mask 时非常有益。

作者在四个公开可用的皮肤病变分割数据集上展示了MobileUNETR的优点,包括模型大小、运行复杂性和准确性。在将模型大小和复杂性降低10倍和23倍的同时,作者显著提高了所有数据集和先进架构以及训练方法的表现。

作者的主要贡献可以总结如下。

  1. 作者提出了一种新颖的轻量级和高效的端到端Transformer基于的混合模型,用于皮肤病变分割,其中在每个阶段都加强了局部和全局上下文特征,以保留对给定场景的整体感知。
  2. 为了克服对基于CNN的解码策略的过度依赖,作者引入了一种新颖的Transformer基于的混合解码器,它同时利用不同分辨率的低级和全局特征进行高度精确且对齐良好的 Mask 生成。
  3. 所提出的架构在只有300万个参数和1.3 GFLOP计算复杂度的情况下,超越了大的、高度复杂的CNN、Transformer和混合模型,在分割方面分别减少了计算复杂度的10倍和23倍,与目前的SOTA模型相比,具有显著的优势。

2 Related Works

皮肤病变分割在自动化皮肤病诊断中至关重要;然而,由于病变的多样性和图像中的噪声,使其难以操作。传统图像处理方法已经被先进的深度学习系统所取代,特别是在卷积神经网络(CNNs)和基于Transformer的方法的出现,这些方法大大提高了分割的准确性和可靠性。

CNN Based Methods

神经网络(DNNs)和卷积神经网络(CNNs)的日益普及,使它们成为皮肤病变分割任务的必备工具。它们巧妙地解决了特征鉴别和数据变化管理的问题。该领域取得了显著的发展,例如由[3]引入的多级全卷积网络(FCN),该网络采用并行集成方法来提高皮肤病变边界的分割精度。[71] 做出了类似的贡献,他们创建了一个专门优化皮肤病理图像分析的改进卷积-反卷积网络,并结合各种颜色空间以更好地诊断病变。 [29] 继续了这种建筑创新趋势,他们的DoubleU-Net结合了多个U-Net结构以提高分割精度。

与此同时,开发自动化检测系统的工作也取得了突出成果。 [50] 在使用膨胀卷积在多个架构如VGG16、VGG19 [56]、MobileNet [26] 和 InceptionV3 [57] 上的早期恶性皮肤病变检测,以及使用HAM10000数据集 [62] 进行训练和测试。预训练网络和深度学习模型的使用也体现在2018年ISIC挑战赛 [2] 等多项获胜解决方案中,其中许多人基于DeepLab [7]架构使用PASCAL VOC-2012 [14]的预训练权重构建模型,并采用了集成策略,包括使用VGG16、U-Net、DenseNet [28] 和 InceptionV3 等模型,通过额外的训练迭代对这些模型进行微调以实现最先进的表现。

[58] 和 [60] 证明了这些模型在多种情境感知设置下的灵活性,通过改进CNN的特征提取,前者通过修改跳过连接,而后者通过多级UNets实现。此外,[1] 引入了一种新的focal Tversky损失函数,该损失函数可以解决医学影像中的数据不平衡问题,从而改善了小病变结构的精确度与召回率之间的平衡。

2019年ISIC挑战赛也促使了 several 项关于使用CNN进行皮肤病理医学影像的新研究。 [47], [45] 和 [61] 使用了各种CNN架构和不同的数据增强方法。这些研究表明CNN能够局部分割皮肤病变,但它们的性能缺陷可以归因于它们无法提取有价值的全局上下文信息。

Transformer Based Methods

限于仅使用局部特征,研究行人被迫寻求新方法。这导致了朝着使用基于全局特征的工具的演变。这一演变体现在从标准卷积神经网络(CNN)技术向使用 Transformer 和自注意力机制的全新方式转变上。[37]首次在皮肤病变分割中使用了密集解卷积网络(DDN),通过使用密集层和链式残差池来捕捉长程关系,这与先前的方法有显著的差异。[69] 调查了对抗学习与SegAN,通过巧妙地捕获细致关系来提高分割精度,这在皮肤病学图像方面是一个显著的发展。[40]和[15]以新的方法显著提高了皮肤病变分割。Mirikharaji和Hamarneh在完全卷积网络中实现了一个星形先验(SSP),通过惩罚非星形区域提高准确性和可靠性,同时保留全局结构。本研究证明了使用形状先验分割复杂皮肤病变模式的有效性。[15] 补充了CPFNet,它使用金字塔模块在特征图上收集全局上下文,成功管理皮肤病变的可变性和复杂病变模式的界定精度。[15] 还引入了BAT,为皮肤病变的分割提供了一个具有边界感知自注意力机制的 Transformer 结构,以有效地在皮肤病变影像中收集全局和局部信息。FAT-Net,一个自适应特征 Transformer 网络用于分割皮肤病变,是由[65]提出的。FAT-Net通过将额外的 Transformer 分支集成到 standard encoder-decoder 结构中,恰好保留了皮肤病变的可变性和不规则性,并改善了黑色素瘤的分析。

3 Methodology

在本节中,作者介绍了MobileUNETR,这是作者用于皮肤病变分割的高性能、高效和轻量级的架构。如图2所示,MobileUNETR的核心架构由两个主要模块组成:

(1)首先,一个轻量级的混合编码器有效地生成粗略的高层次和细粒度的低层次特征;

(2)其次,一个新颖的轻量级混合解码器有效地结合多层特征,同时考虑局部和全局上下文线索,以生成高精度的语义分割 Mask 。

picture.image

Model Complexity

医学影像领域的总体目标是追求特定任务(如皮肤病变分割)上的性能优于复杂性。移动UNETR架构的的一个重要贡献是表明,精心构建的轻量级和高效模型可以比大型计算密集型架构提供更好的性能。如图3中所示,该架构在皮肤病变分割方面与最先进架构相比,大小缩小了10倍,计算效率提高了23倍,同时生成更好的结果。简化模型不仅提高了小数据集上的训练和性能,还便于在资源受限的环境中部署。

picture.image

编码器

在医学视觉研究中,深度学习架构主要分为两大类,卷积神经网络(CNNs)和 Transformer (Transformer),各自具有其优劣势。CNNs因其效率、自然的归纳偏见和能够分层编码特征而在许多医学视觉应用中成为默认方法。然而,尽管取得成功,纯CNN based特征编码器仍无法有效地获得给定场景的全局上下文理解。许多自制的技巧已经被提出来帮助CNNs获得更大的感受野,例如膨胀卷积 [8]和更深层次的模型,但是图像大小和计算复杂性限制了进一步的研究,以帮助提高整体性能。与CNNs不同,Transformers旨在真正全面地理解场景。然而,在大空间分辨率下,它们存在着计算约束,这阻碍了它们在有效深度学习应用程序中的采用。

通过利用CNNs和Transformer架构的天然优缺点,作者提出的编码器在最大化特征表示能力的同时,显著降低了计算复杂性和参数数量。从高层次来看,特征提取模块可以分为两个阶段:首先,基于CNN的局部特征提取和降采样;其次,混杂的Transformer/CNN局部和全局表示学习。

基于CNN的局部特征提取:用于计算机视觉的端到端 Transformer 模型,例如ViT及其变体 ,由于每个输入图像生成的序列长度很长,导致计算非常复杂。通过将Transformer的序列长度瓶颈与ViTs在早期层学习低级特征的自然倾向相结合,可以添加基于CNN的早期特征提取替换,从而显著降低架构的计算复杂性。

具体而言,提出移动UNETR架构中使用的MobileNet Downsampling blocks用于最小化低级特征提取阶段的计算复杂性,同时不牺牲学习的特征表示。此外,基于CNN的特征允许模型更好地结合空间信息,与纯ViT based方法相比,有效地降低了输入数据的空间维度,使下游Transformer层能够高效地学习全局特征表示。

混杂的Transformer/CNN块:一旦通过CNNs实现有效的降采样以降低大空间分辨率计算复杂性,移动ViT块可用于同时提取局部和全局表示。移动ViT块使作者能够在保持空间顺序和局部归纳偏差的同时,集成Transformers的长程上下文优势。该操作可以分解为两个主要部分,如图2中所示。

首先,应用卷积神经网络中的深度可分卷积 [26]编码空间信息并将特征映射到高维空间。最后,以模型长期依赖性,将张量展开为非重叠的扁平patch,并应用自注意力层来捕捉相邻patch的关系。这种组合允许每个特征图在各个阶段具有对场景的局部和全局理解,从而提高其对场景的上下文理解。

Decoder

对于大多数分割模型来说,编码器阶段的重要性被强调,因为这样可以保证学习、压缩并传递相关特征到下一阶段。大多数使用卷积神经网络(CNN)、 Transformer (Transformers)或CNN/Transformer双编码器进行特征提取的解码器方法,严重依赖纯卷积来映射提取到的特征到最终分割 Mask 。这种方法的一个缺点是,通过在解码器中使用纯CNN层,作者使得模型只依赖于瓶颈处提取的信息来学习确保局部连续性而不具有使用全局上下文信息进行模型校准的能力。此外,盲目堆叠CNN层可能导致解码器模块变得庞大,从而增加解码器架构的整体计算复杂性。作者的新异解码器架构是一种快速、计算高效的轻量级方法,允许模型分层构建最终分割 Mask ,同时确保在解码过程中的每个阶段都使用局部和全局上下文特征。所提出的解码模块在150万参数上有效地将CNN和 Transformer 架构的优点结合在一起,作为一种基于CNN的解码方法。

Simple Hybrid Decoder:在医学成像中典型的CNN解码器模块[23, 24, 55]通过转置和标准卷积结合提取和优化编码器的特征。使用这种结构使得模型可以递归地增加空间分辨率,并在每个层利用通过跳跃连接提供的特征优化特征。尽管它们取得成功,但仅依赖于CNN的解码器在动态适应自己以确保每个阶段学习的特征全局对齐方面存在挑战。

所提出的轻量级解码器在每一阶段执行三个操作,以确保每一解码阶段提取的特征在局部和全局上都对齐。首先,通过转置卷积将前一阶段的特征图上采样。接下来,作者通过结合相应的跳跃连接信息对人体上采样特征进行局部细化。最后,作者使用Transformer/CNN混合层让模型根据远距离全局上下文动态调整自身。将局部细化与全局细化阶段结合,使解译器生成能通过改进局部和全局边界来提高分割结果特征的特征。

4 Experimental Results

要展示MobileUNETR作为一款高度竞争力的分割架构的有效性,作者在广泛流行的皮肤病变分割数据集上进行多个实验,并将其与高性能分割模型进行了对比和分析。

Dataset

为了评估作者高效轻量级模型 MobileUNETR 的性能,作者在分割皮肤病变的四个公开数据集进行了测试。国际皮肤影像合作组织(ISIC)已经开发并发布了三个广泛使用的数据集:ISIC 2016、ISIC 2017 和 ISIC 2018,用于皮肤病变分割任务。此外,作者还评估了作者的模型在葡萄牙佩德罗·胡西班牙医院皮肤病服务提供的 PH2 数据集上的性能。数据集的 breakdown 如下所示。

Implementation Details

如图4所示:在ISIC 2016数据集上的定性结果。(a) 原始皮肤镜像输入;(b) 真实 Mask ;(c) 预测 Mask (蓝色)和真实 Mask (红色)叠加在原始图像上。定性结果补充了定量结果,以说明分割性能。叠加图像中红色的显示表示预测和真实之间的分割差异。

picture.image

作者提出的MobileUNETR及其配套实验在配备有CPU和RTX 3090 GPU的服务器上使用PyTorch进行训练和评估。所有模型都遵循简单的训练程序,使用AdamW优化器,参数为和,批次大小为8。实验设置包括线性升温阶段,覆盖40个周期,在此阶段,学习率从0.0004/40逐渐增加到0.0004。然后,在一个400周期的余弦退火调优器中,衰减学习率。按照既定的实践,作者使用PyTorch中可用的直接数据准备和增强技术进行操作,确保了结果的可访问性和可重复性。

Results on ISIC 2016

ISIC 2016数据集代表了第一个标准化皮肤病变分割任务,包含900个训练图像和300个测试图像。作者提出的MobileUNETR在包括FCNN、Attention-augmented FCNNs、基于GAN的生成对抗网络(GAN)和Transformer的七种不同架构进行了比较。表1汇总了跨七个指标的性能结果,与IoU和Dice指标相比,分别提高了2.17%和1.21%。

picture.image

Results on ISIC 2017

如图5:在ISIC 2017数据集上的定性结果。(a)原始皮肤镜像输入。(b)真实 Mask 。(c)在原图像上叠加的预测 Mask (蓝色)和真实 Mask (红色)。定性结果在补充定量发现的基础上,阐述了分割性能。叠加图像中红色部分表示预测和真实之间的分割差异。

picture.image

ISIC 2017通过扩充数据集大小,拓宽了皮肤病变分割的范围。该数据库包含了2500张训练图像和600张测试图像。作者提出的MobileUNETR模型在12种不同的架构上进行了系统性的比较。作者展示出,无论在所有架构上,作者的模型都能保持一致的改进,在IoU、Dice和准确性指标上,同时保持轻量级和高效的设计。结果在表2中展现,作者在IoU和Dice指标上,分别实现了2.47%和1.84%的提升。

picture.image

Results on ISIC 2018

如图6所示:ISIC 2018数据集的定性结果。(a)原始皮肤镜图像。(b)真实目标掩图。(c)预测掩图(蓝色)和真实目标掩图(红色)在原始图像上的叠加。定性结果补充了定量结果,以显示分割性能。在叠加图像中,红色表示预测和真实目标之间的分割差异。

picture.image

ISIC 2018数据集是常用皮肤病变分割数据集中最全面的一个。该数据集包括2694个训练图像和1000个测试图像。与之前的实验一样,提出的MobileUNETR在包括广泛架构的十个不同架构的比较中进行了基准测试。ISIC 2018的七个指标的性能结果总结在表3中。作者的结果一致显示,所有架构中的IOU和Dice指标的改进范围从2.54%到1.71%,同时保持轻量级和高效的设计。###ISIC PH2结果

picture.image

最后,作者使用PH2数据集对MobileUNETR的性能进行了评估。与早期的ISIC数据集不同,PH2是一个相对紧凑的数据集,提供了一个机会来突出作者在处理较小数据集中的混合架构的泛化能力。与作者的先前的实验一致,作者将提出的架构与九个不同的架构进行了比较,并呈现了性能结果表4。作者的结果一致显示,IOU和Dice指标的改进范围分别为2.68%和1.3%。在PH2上的成功实验表明,作者提出的模型在涉及稀疏数据集的应用中具有适应能力。

picture.image

Comparison to Advanced Training Techniques

为了展示即使模型架构紧凑,作者的模型也能实现与使用高级训练技术的大型架构相当的结果,作者采用了称为参数有效微调(PEFT)的一类先进的训练技术。为了证明这一点,作者将作者的方法与最近使用这些先进技术的解决方案进行比较。表5展示了对轻量级架构设计合理性的有效性,证明它们可以与大型复杂模型相媲美,并强调过度参数化并不是现代深度学习的未来。

picture.image

5 Conclusion

编码器-解码器架构为医学图像分割的研究者提供了一个强大的结构范式。尽管它已被成功用于推动医学图像分割的边界,但更大更复杂的编码器-解码器范式版本可能不是现代深度学习架构的解决方案。本文介绍了MobileUNETR,这是一个创新且高效的多层次混合Transformer架构,专为图像分割而设计。

与现有方法不同,MobileUNETR在编码器和解码器阶段有效地整合了局部和全局信息,利用了卷积和变换器的优势。这种整合使得编码器在编码阶段能够提取局部和全局特征,同时允许解码器重构这些特征,确保最终分割掩码中的局部和全局对齐。通过在每个层次上结合局部和全局特征,MobileUNETR避免了需要大型、复杂和过度参数化的模型。

这不仅提升了性能,还显著减少了模型的大小和复杂度。作者进行了广泛的实验,将提出的医学图像分割方法与四个广泛使用的公共数据集(ISIC 2016、ISIC 2017、ISIC 2018和PH2数据集)进行了比较和对比。与最先进方法的比较分析证明了 MobileUNETR架构的有效性,展示了在模型训练和推理方面的优越准确性表现和卓越效率。

在所有数据集上,MobileUNETR在Dice和IoU指标上实现了1.3%至2.68%的提升,并与当前SOTA模型相比,参数和计算复杂度分别减少了10倍和23倍。

作者希望该方法能够成为医学成像研究的一个坚实基础,因为MobileUNETR在图像分割中的应用是无限的。

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论