[WACV2023] The Fully Convolutional Transformer for Medical Image Segmentation

文章提出了一种全卷积Transformer(FCT),用于医学图像分割,解决了Transformer模型在捕捉细粒度特征方面的不足。FCT结合了卷积神经网络和Transformer的优势,通过卷积注意力模块学习长范围语义依赖,Wide-Focus模块则利用多分辨率空洞卷积获取局部到全局上下文。实验结果显示,FCT在多个医学图像分割数据集上超越了现有模型,且参数更少。
摘要由CSDN通过智能技术生成

The Fully Convolutional Transformer for Medical Image Segmentation

摘要
  1. 医学图像分析的细粒度特性所带来的挑战意味着Transformer对其分析的适应仍处于初级阶段。UNet压倒性的成功在于它能够欣赏分割任务的细粒度性质,这是现有的基于Transformer的模型目前不具备的能力。

  2. 为了解决这个缺点,本文提出了全卷积Transformer(FCT),它建立在卷积神经网络学习有效图像表示的能力的基础上,并将它们与Transformer的能力相结合,有效地捕获其输入中的长期依赖关系。

  3. FCT是医学影像中第一个全卷积Transformer模型,能够分割不同形态的医学图像。它分两个阶段处理输入,首先,它学习从输入图像中提取长期语义依赖关系,然后学习从特征中捕获分层的全局属性。FCT结构紧凑、准确、健壮。

  4. 本文的结果表明,它在不需要任何预训练的情况下,在不同数据形态的多个医学图像分割数据集上,大大优于所有现有的Transformer架构。FCT在ACDC数据集上比其直接模型高出1.3%,在Synapse数据集上高出4.4%,在Spleen数据集上高出1.2%,在ISIC 2017数据集上高出1.1%,在dice metric上的参数少了多达5倍。在ACDC Post-2017MICCAI-Challenge在线测试集上,我们的模型在未见过的MRI测试用例上设置了新的最先进的技术,优于大型集成模型以及参数更少的nnUNet。

1 引言

医学图像分割是计算机辅助诊断的关键工具。它有助于检测和定位图像中病变的边界,有助于快速识别肿瘤和癌变区域的潜在存在。这有可能加快诊断,提高检测肿瘤的可能性,并使临床医生更有效地利用他们的时间,对患者的结果有利。传统上,现代医学图像分割算法构建为对称的自顶向下编码器-解码器结构,首先将输入图像压缩(编码)到潜在空间,然后学习解码图像中感兴趣区域的位置。将中间信号的水平传播(跳过连接)添加到这个垂直信息流中,我们就得到了UNet架构,这可以说是最近分割算法中最具影响力的飞跃。今天大多数现代分割系统都包括UNet或其变体。UNet成功的关键在于其完全卷积的性质。UNet在其结构中不估计任何非卷积可训练参数。

基于卷积神经网络(CNN)的UNet模型在医学图像分割任务中的准确性和性能方面取得了巨大的成功。然而,为了真正帮助临床医生进行早期疾病诊断,它们仍然需要额外的改进。卷积算子固有的局部特性是CNN的一个关键问题,因为它阻止了它们利用来自输入图像的长范围语义依赖。人们提出了各种方法来为CNN添加全局上下文,最引人注目的是引入注意力机制,以及扩大卷积核以增加核的感受野。然而,这些方法都有自己的缺点。Transformer在语言学习任务中取得了巨大的成功,因为它们能够有效地处理非常长范围的序列依赖。这导致它们最近适应了各种视觉任务。最近提出的架构,如ViT,已经超过了CNN在基准成像任务上的性能,而最近对ViT的许多改进,如CvT , CCT和Swin Transformer,已经表明Transformer不需要庞大的数据消耗模型,甚至可以处理少量数据,从而超过CNN的性能。通常,ViT风格的模型首先从图像中提取离散的非重叠patch(在NLP中称为token)。然后,他们通过位置编码将空间定位注入到这些patch中,并将此表示通过标准transformer层来建模数据中的长期语义依赖关系。

考虑到CNN和Transformer的明显优点,我们认为医学图像分割的下一步是一个完全卷积编码器-解码器深度学习模型,能够有效地利用医学图像中的长程语义依赖。为了实现这一目标,我们提出了第一个用于医学图像分割的全卷积Transformer。我们新颖的全卷积Transformer层构成了我们模型的主要构建块。它包含两个关键组件,一个卷积注意力模块和一个全卷积Wide-Focus模块。我们将我们的贡献形式化如下:

  1. 我们提出了第一个用于医学图像分割的全卷积Transformer,它超越了所有现有的基于卷积和Transformer的医学图像分割架构的性能,用于多个二分类和语义分割数据集。
  2. 我们提出了一种新型的全卷积Transformer层,它使用卷积注意力模块来学习长范围的语义上下文,然后通过宽焦点模块使用多分辨率空洞卷积创建分层的局部到全局上下文。
  3. 通过广泛的消融实验,我们展示了我们模型的各种构建块在其对模型性能影响的背景下的影响。
2 相关工作

早期的CNN和注意力模型:

​ UNet是第一个用于医学图像分割的CNN模型。最早将注意力模型引入医学图像分割的工作之一,是通过将门控函数应用于UNet的编码器到解码器的特征传播。FocusNet等方法采用双编码器-解码器结构,其中注意力门选学习将相关特征从一个UNet的解码器传播到下一个UNet的编码器。FocesNet++是在分组卷积的各种过滤器组中集成注意力机制的第一个作品之一。还有许多UNet的变体,它们使用不同的残差块来增强特征提取[32,28,33,20,16]。UNet++在编码器和解码器之间创建了嵌套的分层密集跳过连接路径,以减少它们之间学习特征的语义差距。作为最近最具影响力的UNet变体,nnUNet自动调整自身来预处理数据,并选择最适合任务的最佳网络架构,而不需要人工干预。

Transformer模型:

最初的Transformer架构彻底改变了自然语言处理任务,并迅速成为视觉理解任务的模型。Transformer在视觉方面工作得很好,因为它们能够创建长范围的视觉环境,但存在固有的缺点,不能利用CNN等图像中的空间环境。最近的工作转向了克服这一缺陷的可能解决方案。CvT , CCT和Swin Transformer都是在Transformer中集成足够的空间环境的尝试。在医学图像分割中,大多数现有研究着眼于创建用于特征处理的Transformer-CNN混合模型。与Attention UNet类似,UNet Transformer增强了CNN,在跳过连接内增加了多头注意。TransUNet是最早提出的用于医学图像分割的Transformer-CNN混合模型之一,它使用Transformer编码器馈送到级联卷积解码器。与TransUNet类似,UNETR和Swin UNETR在编码器上使用Transformer和卷积解码器来构造分割图。Transfuse运行双分支编码器,一个带有卷积层,另一个带有Transformer层,并将其特征与新颖的BiFusion模块结合起来。然而,这个模型的解码器是卷积的。

当前工作:

最近有一个转变,从创建混合Transformer-CNN模型,到改进Transformer块本身,以处理医学图像的细微差别。Swin UNet是第一个提出用于处理医学图像的纯Transformer的架构。这里的纯指的是仅由Transformer层提取和处理的图像特征,而不需要预训练的骨干网络架构。DS-TransUNet引入Transformer Interactive Fusion模块,以获得更好的表示全局依赖。这两个模型的计算核心都是Swin Transformer块。同时进行的工作,如nnFormer和DFormer,试图利用医学图像中的本地和全局上下文,通过特别制作的多头自注意力块来满足这一任务。这些模型的主要缺点是它们固有的注意力投射和特征处理的线性性质,FCT旨在缓解这一点。

现有的医学影像分割模型目前至少存在以下三个局限性之一。它们要么基于CNN主干网络,要么使用卷积层创建,因此限制了它们超越感受野以获得图像语义上下文的能力(早期CNN方法)。他们试图将Transformer集成到他们的特征处理管道中,以利用它们创建长期语义上下文的能力,但反过来,使模型庞大且计算复杂(混合Transformer-CNN)。他们试图通过创建用于分割的纯Transformer模型来减少这种计算负担,而不试图在低级特征提取阶段(并发工作)对局部空间上下文建模。与现有方法不同,我们的全卷积Transformer没有这些缺点,同时仍然是一个纯基于Transformer的医疗图像分割架构。补充资料中的表4额外总结了FCT与现有模型相比的主要差异。

3 全卷积Transformer

给定一个数据集{X, Y},其中,X是我们模型的输入图像,Y是相应的语义或二值分割映射。对于每个图像xi∈R H×W×C,其中H和W为图像的空间分辨率,C ={3,…, N}为输入通道数,我们的模型产生一个输出分割映射yi∈R H×W×K,其中,K∈{1,…D}。FCT的输入是从输入3D图像的每个切片中采样的2D补丁。我们的模型遵循熟悉的UNet形状,FCT层作为其基本构建块。与现有的方法不同,我们的模型既不是CNN-Transformer的混合,也不是Transformer - UNet的结构,它使用现成的Transformer层来编码或细化输入特征。它首先从图像中提取重叠的patch,然后创建基于patch的扫描嵌入,然后在这些patch上应用多头自注意,从而构建特征表示。然后通过我们的Wide-Focus模块处理给定图像的输出投影,以从投影中提取细粒度信息。图1显示了我们的网络体系结构的概述。

在这里插入图片描述

图1:用于医学图像分割的全卷积Transformer。网络(底部)遵循标准的UNet形状,显著的区别是它完全基于卷积Transformer。FCT层的第一个组件(顶部)是卷积注意力。在这里,投影层中的深度卷积消除了对位置编码的需要,从而得到一个更简单的模型。我们创建重叠的patch,其中patch的重叠程度是通过卷积投影层的步幅来控制的。为了利用图像中的空间上下文,我们的MHSA块用深度卷积代替线性投影。Wide-Focus模块在线性增加的接受野处对MHSA输出应用扩张卷积。

3.1 FCT层

​ 每个FCT层都从LayerNormalization-Conv-Conv-Maxpool操作开始。我们从经验上注意到,与直接先创建图像的patch-wise投影相比,在3×3内核大小较小的patch上按顺序应用这些连续卷积有助于更好地编码图像信息。每个卷积层后面都有一个Gelu激活函数。我们的FCT块与其他模型块不同的第一个实例是通过它对医学成像的卷积注意力应用。

​ MaxPool的输出被输入到转换函数T ( ⋅ )中,转换函数T ( ⋅ ) 将其转换为新的token映射。我们选择的T ( ⋅ ) 是深度卷积算子。我们选择一个较小的内核大小3×3, 步长为s×s和一个有效的填充,以确保:(1)与大多数现有工作不同,提取的patch是重叠的,并且(2)卷积操作不会始终改变输出大小。

​ 接下来是LayerNormalization操作。得到的token映射p i + 1 ∈ R W t × H t × C t ,创建了我们的patch嵌入式输入。

​ FCT层不同于现有的基于Transformer的医学成像应用方法,是通过它的注意力投影。所有现有模型都采用线性逐点线性映射来进行多头自我注意(MHSA)计算。这导致Transformer模型失去空间信息,这对成像应用非常重要。现有的方法试图通过卷积增强来缓解这个问题,使其适应成像任务。然而,这为所提出的模型增加了额外的计算成本。受[36]中提出的方法的启发,我们将MHSA块中的逐点线性映射替换为深度卷积以降低计算成本,并从图像中利用更好的空间上下文信息。patch嵌入和卷积注意力投影构成了我们的卷积注意力的组成部分。与[36]不同的是,我们注意到用LayerNormalization替换BatchNormalization有助于提高性能。此外,删除点卷积会导致一个更简单的模型,而不会损失任何性能深度卷积提供的空间上下文进一步消除了对位置编码的需求,位置编码用于在输入中插入空间信息,并顺序跟踪每个patch的位置,从而进一步简化了架构设计。

​ 一般的Transformer层遵循线性层MHSA块,因此丢失了图像中的所有空间上下文。直接用卷积替换这些线性层是一种相对简单的方法,可以缓解这个问题并提高性能。然而,医学图像需要细粒度的信息处理。记住这一点,本文采用了一个多分支卷积层,其中一层对MHSA输出应用空间卷积,而其他层应用空洞卷积,增加感受野,以获得更好的空间上下文。然后,通过求和来融合这些特征,并将它们传递到特征聚合层。这种特征聚合是通过另一个空间卷积算子完成的。我们称这个模块为Wide-Focus。残差连接用于增强整个层的特征传播。最后的特征被重新塑造,并进一步传播到下一个FCT层。图1(上)显示了FCT层。
在这里插入图片描述

3.2 编码器

我们的模型的编码器包含四个FCT层,负责特征提取和传播。对于第l个Transformer层,卷积注意力模块的输出为,z ‘ l = MHSA(zl−1)+ z q/k/v l−1,其中zl−1q/k/v = Flatten(DepthConv(ReShape(zl−1)))。MHSA(zl−1)= softmax(QKT/√d)V。然后由Wide-Focus (WF)模块处理z’l, zl=WF(zl)+z’l。我们进一步为编码器注入金字塔风格的图像输入,目的是在不同尺度上突出显示不同类别和更小的ROI特征。值得注意的是,即使没有这种多尺度图像金字塔输入,我们的模型也能够获得最先进的结果。数据的(瓶颈)潜在编码是使用另一个FCT层创建的。

3.3 解码器

解码器将瓶颈表示作为其输入,并学习从该信息中重新采样二进制或语义分割映射。为了在解码器层中创建更好的上下文相关性,还使用从编码器到解码器的跳过连接,其中来自编码器层的具有相同分辨率的特征映射与解码器层连接。解码器的形状与编码器对称。解码器中的层对应于编码器中的图像金字塔层,输出中间分割映射,提供额外的监督并提高模型的预测能力。上下文相关性是通过首先对特征量进行上采样,然后将其传递到FCT层以了解其最佳可能表示来创建的。我们没有在FCT的最低规模上采用深度监管,因此我们的模型不是“完全深度监管”。这是因为我们观察到,输入图像扫描中的感兴趣区域(roi)有时太小,无法在最低尺度(28 × 28)进行分割,这导致模型性能较差。这种低规模的输出在模型中增加了强烈的偏差,以预测一些输出roi作为背景类。

4 实验

数据集:

我们通过在四个不同模式的不同数据集上实现最先进结果的能力,证明了我们模型的有效性。我们使用来自(MRI)自动心脏诊断挑战(ACDC)、(CT)Synapse多器官分割挑战1、(CT)脾脏分割数据集和(皮肤镜)ISIC 2017皮肤癌分割挑战的数据。

ACDC数据集包含100个MRI扫描,其中包含左心室(LV)、右心室(RV)和心肌(MYO)的真实数据,train-val- test为70-10-20。

Synapse包含30名患者的CT扫描结果。我们对Synapse的实验设置和预处理类似于TransUNet。

脾脏分割数据集包含41个CT卷。我们对这个数据集的train-val- test分割为80-10-10。

对于ISIC 2017数据集,我们从训练数据集中的2000张图像中创建了70-10-20的train-val- test分割。

实验细节

使用Dice系数来衡量模型的性能。

模型输入有两种尺寸224 x 224384 x 384

使用TensorFlow 2.0运行所有的实验。使用一个NVIDIA A6000 GPU进行所有的实验。

我们的损失函数是交叉熵和骰子损失的同等加权组合。

使用Adam,其学习率为1e−3,通过监测验证损失在平台上降低。

先进行50个epoch的热身训练,然后再训练模型250个epoch。

数据增强如下:旋转(0◦到360◦),缩放范围(最大0.2),剪切范围(最大0.1),水平/垂直移动(最大0.3),水平和垂直翻转。FCT的默认设置是-每个阶段的过滤器数量16、32、64、128、384、128、64、32、16,每个阶段的注意头数量2、4、8、12、16、12、8、4、2。使用10个ACDC批量,4个Synapse,脾脏分割和ISIC 2017分割。我们从一个随机初始化的权重集训练我们所有的模型。

5 实验结果

我们的模型以更少的参数和Gflops在所有报告的基线中实现了最先进的结果。FCT包含3170万个参数和7.87 GFLOPs。

在ACDC数据集上,我们的模型大小比我们最接近的竞争对手nnFormer(15892万,157.88 GFLOPs)小五倍,优于所有现有的工作。我们在两种不同的图像大小上训练我们的模型,看看它对性能的影响。正如预期的那样,图像大小为384×384的FCT比输入图像大小为224×224的模型获得了更好的结果,因为增加的空间分辨率使模型能够更清楚地看到图像中的细粒度细节。与我们的模型相比,我们还测试了在每个尺度上进行深度监管与不使用深度监管的效果。表2总结了我们在ACDC数据集上的结果。这也表明,我们采用的深度监督设置是我们模型的最佳设置。

在这里插入图片描述

为了证明我们结果的统计显著性,我们还使用ACDC数据集进行了5倍交叉验证(CV),并计算p值,以表明我们的结果明显优于nnFormer。

我们使用FCT224进行这些实验。使用5-Fold CV,我们得到平均骰子分数为92.43±0.38。然后,我们在ACDC数据集上运行5次实验,并平均它们,得到92.88±0.09的骰子分数。

这两个结果对于数据集来说都是最先进的。与nnFormer(91.78±0.18)比较,p < 0.0001的ACDC挑战数据集中的100张图像,并报告我们对50个未提供Ground Truth掩码的未见测试用例的结果。我们在512 × 512的输入图像上训练我们的模型。为了考虑数据集中图像大小的变化,我们对图像进行裁剪和平铺,以获得512×512分辨率,并对蒙版应用相同的增强。为了生成最终的预测,我们删除了由于平铺作为后处理步骤而产生的这些额外的预测,通过平均平铺的预测来创建最终的输出。我们训练这个模型,如第4节所示。我们的研究结果的链接在网上,可以与以前最先进的研究结果进行比较。表4总结了前五名提交的结果与我们的结果进行比较(补充材料中的表5显示了所有类别的详细结果)。我们采用平均值来提供平均值,但是,详细的结果表可以在提供的链接中找到。

在这里插入图片描述

6 消融实验

我们主要通过消融实验了两个关键组件对模型性能的影响:去除编码器到解码器的跳过连接,以及我们新颖的Wide-Focus模块的不同设置。我们在ACDC数据集上进行消融。跳跃连接显然对我们的模型的性能很重要(见表3补充材料),最佳设置类似于原始的UNet。

为了创建Wide-Focus模块的最佳设置(见表1),我们观察了更宽的卷积分支和更大的膨胀率对模型性能的影响。

在这里插入图片描述

我们观察到,超过三个卷积分支,扩张率线性增加,模型的准确性开始饱和,并最终下降。我们认为这是由于扩张的核在更深的层次上无法近似全局核,导致扩张的接受域缺少关键特征信息。这也符合我们的发现,FCT块中更小的内核导致更好的性能。

7 结论

我们提出了全卷积Transformer,它能够准确地执行二进制和语义分割任务,参数比现有模型更少。FCT在参数数量上比nnFormer小5倍以上,比TransUNet和LeViT - Unet小3倍以上。FCT层由两个关键组件组成——卷积注意力和Wide-Focus。卷积注意力通过使用深度卷积为模型创建重叠的patch,消除了在patch创建阶段对位置编码的需求。我们基于深度卷积的MHSA块集成了空间信息,首次在医学成像背景下估计长期语义依赖关系。从我们的消融实验中可以看出,Wide-Focus有助于利用医学图像中存在的细粒度特征信息,并且是提高Transformer块性能的重要因素。我们通过在多个高度竞争的不同模式和维度的细分数据集上的最先进的结果证明了我们模型的能力。我们的FCT块是第一个为医学成像应用而提出的全卷积Transformer块,并且可以轻松扩展到医学成像的其他领域和应用。我们相信我们的模型可以作为未来分割任务的有效骨干,并为基于Transformer的医学图像处理的创新铺平道路。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值