[深度学习论文解读] TransBTS: Multimodal Brain Tumor Segmentation Using Transformer
基于Transformer的多模态脑肿瘤分割
论文:https://arxiv.org/pdf/2103.04430
代码:https://github.com/Wenxuan-1119/TransBTS
发表时间:Mar 2021
一、基本介绍
1.1胶质瘤
胶质瘤是最常见的恶性脑肿瘤,具有不同程度的侵袭性。在磁共振成像(MRI)上自动精确地分割这些恶性肿瘤对于临床诊断至关重要。卷积神经网络在分类、分割和目标检测等各种视觉任务中取得了巨大的成功。全卷积网络(FCN) 首次实现了端到端语义分割,取得了令人印象深刻的结果。U-Net 使用具有跳跃连接的对称编码-解码结构来改善细节,成为医学图像分割的主流架构。很多U-Net变体比如U-Net++ 和Res-UNet 进一步提高了图像分割的性能。虽然基于CNN的方法具有出色的表示能力,但由于卷积核的感受野有限,很难建立明确的长距离依赖关系。卷积运算的这种局限性对学习全局语义信息提出了挑战,而全局语义信息对于像分割这样的密集预测任务是至关重要的。
1.2 Transformer
在计算机视觉领域中,卷积神经网络(CNN)一直占据主流地位。不过,不断有研究者尝试将 NLP 领域的 Transformer 进行跨界研究,有的还实现了相当不错的结果。
受到 NLP 领域中 Transformer 缩放成功的启发,这项研究尝试将标准 Transformer 直接应用于图像,并尽可能减少修改。为此,该研究将图像分割成多个图像块(patch),并将这些图像块的线性嵌入序列作为 Transformer 的输入。然后用 NLP 领域中处理 token 的方式处理图像块,并以监督的方式训练图像分类模型。
详细可参考:NLP/CV模型跨界进行到底,视觉Transformer要赶超CNN?
1.3提出问题
但是,当直接将图像分割成补丁作为Transformer的token时,局部结构会被忽略。此外,对于超出2D的医疗体积数据(如3D MRI扫描),连续切片之间的局部特征建模(即深度维度)对于体积分割也是至关重要的。因此,不禁要问:如何设计一种神经网络,利用具有高度表现力的transformer,在空间和深度维度上有效地建模体积数据的局部和全局特征?
二、网络详解
2.1 解决思路
受自然语言处理中注意机制的启发,通过将注意机制与CNN模型融合来克服局限性。Transformer 被设计成对序列到序列任务中的长距离相关性进行建模,并捕捉序列中任意位置之间的关系。这种架构是完全基于self-attention而提出的,完全不需要卷积。与以往基于CNN的方法不同,Transformer不仅在全局上下文建模方面功能强大,而且在大规模预训练的情况下,可以在后半部分任务上取得优异的效果。
基于transformer的框架在各种计算机视觉任务上也达到了最先进的性能。Vision transformer(ViT)将图像分割成块,并使用transformer将这些块之间的相关性建模为序列,从而在图像分类上获得令人满意的结果。在Transformer的帮助下,将目标检测视为一个集合预测任务。TransUNet 就是一项类似工作,它将ViT用于医学图像分割。
2.2方法详解
Transformer利用sel-attention机制进行全局(长距离)信息建模,近年来在自然语言处理和二维图像分类领域取得了成功。然而,局部和全局特征对于密集预测任务至关重要,特别是对于三维医学图像分割。这篇文章首次利用三维CNN中的Transformer进行MRI脑肿瘤分割,提出了一种基于编码-解码结构的TransBTS网络。为了获取局部的三维上下文信息,编码器首先利用三维CNN提取三维空间特征图。同时,对特征映射进行了细致的改造,将标记输入到Transformer中进行全局特征建模。解码器利用Transformer嵌入的特征,进行渐进式上采样来预测详细的分割图。
2.3结构详解
① 大致结构
对于给定输入mr图像X∈RC×H×W×Dwi,空间分辨率为H×W,深度维度为D(切片数)和C(模态数),首先利用3D CNN生成捕获空间和深度信息的紧凑特征图,然后利用transformer编码器在全局空间中对长距离相关性进行建模。之后,重复叠加上采样和卷积层,以逐渐产生高分辨率的分割结果。
大致就是,利用3D CNN生成捕获空间和深度信息的紧凑特征图,下采样大家应该都很熟悉,再就是利用Transformer编码器在全局信息中对长距离相关性进行建模,最后重复叠加上采样和卷积层逐渐产生高分辨率的分割效果。
②网络编码器:
由于Transformer的计算复杂度相对于标记数(即序列长度)是二次的,因此将输入图像直接展平为序列作为Transformer输入是不切实际的。因此,ViT 将图像分割成固定大小(16 × 16)的补丁,然后将每个补丁重新整形为一个token,将序列长度减少了162倍。对于3D体积数据,直接的标记化,在ViT之后,将数据分割成3D面片。然而,这种简单的策略使得Transformer无法跨空间和深度维度对图像局部上下文信息进行建模以进行体积分割。
为了应对这一挑战,作者的解决方案是将3×3×3卷积块与下采样(步幅=2的卷积)进行堆叠,以此将输入图像逐渐编码为高级别特征表示F ∈ R**k×H/8×W/8×D/8(K = 128),这是H、W和D(总步幅(OS)=8)输入维度的1/8。这样,丰富的局部3D上下文特征被有效地嵌入到F中。然后,F被馈送到Transformer编码器,以进一步学习与全局感受野的长期相关性。
③ 网络解码器:
网络解码器相对简单一点,为了在原始3D图像空间(H × W × D)中产生分割结果,引入3D CNN解码器来执行特征上采样和像素级分割(上图右侧部分)。
2.3 与transunet的区别:
TransUNet也最近一项将Transformer用于医学图像分割的论文。作者强调了TransBTS和TransUNet之间的几个关键区别:
①TransUNet是一个2D网络,它以逐层方式处理每个3D医学图像。然而,我们的TransBTS基于3D CNN,一次处理所有图像切片,允许利用切片之间连续信息的更好表示。换句话说,TransUNet只关注标记化图像块之间的空间相关性,但我们的方法可以同时在切片/深度维度和空间维度上建模长期相关性,用于体积分割。
②由于TransUNet采用ViT结构,它依赖于大规模图像数据集上预先训练好的ViT模型。相比之下,我们的TransBTS具有灵活的网络设计,并在特定于任务的数据集上从头开始训练,而不依赖于预先训练的权重。
三、方法亮点
将transformer应用到医学图像分割领域,通过获取全局信息与3D CNN有效结合而得到了不错的分割效果
四、网络模型主要应用及结果
4.1 实验中使用的图像分割数据集
数据集:
使用的3D磁共振数据集由脑瘤分割2019挑战赛提供。它包含335例患者用于训练,125例用于验证。每个样本由四种脑磁共振成像扫描模式组成,即 T1、T1ce、T2和FLAIR。(具体可参考:多模态)
每个图片都是240 × 240×155,并排列在同一个空间中。标记包含4类:背景(标记0)、坏死和非增强肿瘤(标记1)、水肿(标记2)和增强肿瘤(标记4)。分割精度通过用于增强肿瘤区域(ET,标记1)、肿瘤核心区域(TC,标记1和4)和整个肿瘤区域(WT,标记1、2和4)的Dice评分和Hausdorff距离(95%)度量来测量。(Hausdorff)
4.2主要结果
在BraTS 2019数据集上的实验表明,TransBTS在3D MRI扫描分割上实现了优于最先进的脑肿瘤分割方法的性能。
首先对BraTS 2019训练集进行五重交叉验证评估,TransBTS 的平均dice分数分别为78.92%、90.23%和81.19%。还在BraTS 2019验证集上进行实验,并将TransBTS与最先进的(SOTA) 3D方法进行比较,结果如表所示。TransBTS在ET、WT、TC上分别达到78.93%、90.00%、81.94%的最佳dice分数。
在Hausdorff度量方面,分割也取得了相当大的改进。与普通的3D U-Net相比, TransBTS在这两个指标上都显示出巨大的优势,并有显著的改进。这清楚地揭示了利用Transformer建模全局关系的好处。
4.3视觉比较
为了定性分析,还显示了包括3D U-net、V-net、attention U-net和TransBTS在内的各种方法的脑肿瘤分割结果的视觉比较。由于验证集的ground truth不可用,对所有方法的训练集进行五重交叉验证评估。从图中可以明显看出,TransBTS可以更准确地描述脑肿瘤,并通过对每个体积之间的相关性进行建模来生成更好的分割mask。
4.4消融实验
①序列长度消融实验
②Transformer位置的消融实验(下采样四次效果最好)
③跳跃连接位置的消融研究
关于消融实验部分作者在论文中讲的很详细,这里就不多做赘述
六、个人思考总结
Transformer最近在深度学习领域都十分热门,从NLP到CV,归根结底结果优于一切,既然能够提升精度,我们又怎么不能去尝试一下呢,另一方面,还是之前提到的,各领域之间都有着某种程度上的联系,各领域的优秀成果都值得的我们去学习、研究。
总结:
提出了一种新的分割框架,该框架有效地结合了3D CNN和transformer,用于磁共振成像中的多模态脑肿瘤分割。最终的体系结构TransBTS不仅继承了3D CNN建模本地上下文信息的优势,而且利用了transformer学习全局语义相关性。在BraTS 2019数据集上的实验结果验证了所提出的TransBTS的有效性。