[深度学习论文笔记] A Two-Stage Cascade Model with Variational Autoencoders and Attention Gates for MRI Brain Tumor Segmentation
磁共振脑肿瘤分割的变分自动编码器和注意门两阶段级联模型
论文：https://arxiv.org/abs/2011.02881
代码：https://github.com/shu-hai/two-stage-VAE-Attention-gate-BraTS2020
发表时间：MICCAI BrainLesion Workshop 2020
一、基本介绍
1.1问题动机
脑肿瘤根据其起源可分为原发性肿瘤和继发性肿瘤。胶质瘤是最常见的原发性脑肿瘤，可进一步分为低级别胶质瘤(LGG)和高级别胶质瘤(HGG)。HGG是一种恶性脑瘤类型，具有高度的侵略性，经常需要手术。通常，采集几个互补的3D磁共振成像(MRI)模式（t1，t2，flair，t1ce）来突出不同的组织属性和肿瘤扩散区域。与依赖医生专业知识和经验的传统方法相比，自动3D脑肿瘤分割是省时的，并且可以为进一步的肿瘤分析和监测提供客观和可再现的结果。近年来，基于深度学习的分割方法显示出比传统方法更好的性能。
多模态脑肿瘤分割挑战赛（BraTS）是一年一度的国际竞赛，旨在评估脑肿瘤分割的最新方法。组织者提供了一个3D多模态磁共振成像数据集，其中包含由医生和放射科医生注释的“基本事实”肿瘤分割标签。对于每个患者，提供四种3D磁共振成像模式，包括自然T1加权(T1)、对比后T1加权(T1c)、T2加权(T2)和T2流体衰减反转恢复(T2-FLAIR)体积。脑肿瘤分割任务集中在三个肿瘤子区域:坏死和非增强肿瘤(NCR/NET，标记为1)、瘤周水肿(ED，标记为2)和GD增强肿瘤(ET，标记为4)。下图示出了患者的图像集。这种分割任务的竞争方法的排名由指标决定，包括Dice分数、豪斯多夫距离(95%)、灵敏度和特异性，在ET,肿瘤核心(TC = ET+NCR/NET)和全肿瘤(WT=TC+ED)的测试数据集上进行评估。
思路来源：
在BraTS 2018中，Myronenko提出了一种非对称的U-Net（3d mri brain tumor segmentation using autoencoder regularization ），其中较大的编码器用于特征提取，较小的解码器用于标签重建，并赢得了挑战的第一名。该方法的一个令人鼓舞的创新是利用变分自动编码器(VAE)分支来正则化编码器并提高泛化性能。BraTS 2019冠军团队蒋等人提出了一个两阶段网络（Two-Stage Cascaded U-Net ），在第一阶段使用类似于Myronenko的非对称UNet来获得粗略预测，然后在第二阶段使用类似但更宽的网络来细化预测。在第二级网络的解码器中采用了一个附加分支来正则化相关的编码器。上述两种模型的成功表明了向解码器添加分支以减少过拟合和提高模型性能的可行性和重要性。
（脑肿瘤分割挑战赛冠军方法学习笔记）
与一般的计算机视觉问题相比，3D MRI图像分割任务通常面临两个特殊的挑战:训练数据的稀缺性和类别不平衡。为了缓解训练数据的不足，Isensee等人（nnunet）通过使用联合训练策略利用了额外标记数据的优势。（One-pass multi-task convolutional neural networks for efficient brain tumor segmentation）等人结合了一些提高性能的技巧，如引入焦点损失来缓解类别不平衡，以实现进一步的改进。
对于脑肿瘤的分割任务，另一个具有挑战性的困难是不同肿瘤发展阶段和不同病例的肿瘤形态和位置的可变性。为了提高预测精度，许多分割方法将任务分解为单独的定位和后续的分割步骤，并附加先前的目标定位模型。如Wang等人（Automatic brain tumor segmentation using cascaded anisotropic convolutional neural networks ）根据肿瘤分区层次顺序训练三个网络（可参考：基于级联各向异性卷积神经网络的脑肿瘤自动分割）。Oktay等人（attention unet）证明，在胰腺肿瘤分割任务中，通过在标准的卷积神经网络框架中引入注意门(AGs)可以实现同样的目标。
上述工作的启发，这篇文章提出了一种用于脑肿瘤分割的两阶段级联网络。借用Myronenko的网络结构作为第一阶段的网络，得到比较粗糙的分割结果。第二级网络将第一级网络的初步分割图与MRI图像的拼接作为输入，目的是细化NCR/NET和ET子区域的预测。使用AGs（attention unet）进一步抑制不相关背景区域的特征响应。第二阶段网络展示了以下能力:
(i)为模型集成提供更多具有竞争性能的候选模型，
(ii)稳定不同时代模型的预测，
(iii)提高每个单一模型的性能。
二、网络结构
2.1 主要方法
提出的两级网络结构由两个级联网络组成。第一阶段网络以多模态MRI图像为输入，预测粗分割图。将初步分割图与核磁共振图像的连接传入第二阶段网络，以产生改进的分割结果。
2.2第一阶段网络:具有VAE分支的非对称u-net
该网络架构包括较大的编码路径用于语义特征提取，较小的解码路径用于分割图预测，以及一个VAE用于输入图像重建。这部分与3d mri brain tumor segmentation using autoencoder regularization中提出的网络相同。
编码器 编码器由ResNet块组成，分四个空间层次，块数分别为1、2、2、4。每个ResNet块有两个卷积，分别是组归一化(Group Normalization)和ReLU，然后是一个附加的标识跳过连接。编码器的输入是一个大小为4×160×192×128的MRI裁剪，第一个通道涉及四种MRI模式。输入通过3×3×3的32个滤波器卷积层和速率为0.2的dropout层进行处理，然后通过一系列ResNet块。在每两个具有不同空间层次的块之间，采用步长为2的3×3×3卷积将特征图的分辨率降低2倍，同时使特征通道数增加一倍。编码器端点的大小为256×20×24×16，是输入数据空间大小的1/8。
解码器 解码器与编码器几乎是对称的结构，除了每个空间级别中的ResNet块数为1。在每个块之后，使用一个三线性上采样器来恢复2倍的空间大小，并使用一个1×1×1卷积来减少2倍的特征通道数量，然后从相应空间级的编码器输出进行加法跳跃连接。每个块中的操作与编码器中的操作相同。在解码器的末端，使用1×1×1卷积将特征通道的数量从32个减少到3个，然后使用sigmoid函数将将特征映射转换为概率映射。
VAE分支 此解码器分支接收编码器的输出，并产生原始输入的重构图像。首先，使用一个全连接层将解码器端点输出减少到一个256的低维空间，其中256代表了高斯分布的128个均值和128个标准差，从中抽取一个大小为128的样本。然后将绘制出来的向量映射回具有相同空间性质的高维空间，并按照与解码器相同的策略逐步重构为输入图像的维数。请注意，编码器和VAE分支之间没有附加的跳跃连接。
2.3第二阶段网络:具有VAE分支的注意门非对称u-net 上下文学习：
第二级网络的输入是基于第一级网络生成的分割图构建的。为了缓解标签不平衡的问题，将第一阶段网络的输出裁剪成128×128×128体素的空间大小，集中在肿瘤区域。裁剪后的分割图然后连接到原始的MRI图像(裁剪到相同的区域)。
编码器 第二级网络的编码器部分具有与第一级网络相同的结构，而输入有7个通道(3个用于分割maps，4个用于多模态MRI图像)，空间大小为128×128×128体素。
解码器 与第一级网络不同，在解码器部分加入了（attention unet）的AGs。在每个空间层次上，粗尺度的门控信号被传递到注意门以确定注意系数。AG的输出是通过跳跃连接和注意系数从编码器输入特征的哈达玛乘积。然后，在每个空间水平上的AG输出与粗尺度上的2倍上采样特征通过元素求和进行集成。网络结构的其余部分与第一级网络中的解码器保持相同。（哈达玛积）
2.3结构详解
Attention gate不是使用单一相同的标量值来代表每个像素向量的注意水平，而是计算一个门控向量gi来确定每个像素i的聚焦区域。在第l个空间级别内，AG公式如下:
在每个AG中，从粗尺度的门控信号gl+1 i中提取互补信息。为了减少计算成本，对输入特征xl i和门控信号gl+1 i进行线性变换WT x和WT g(1×1×1卷积)，分别将特征尺寸缩小2个，通道数量减少2个。因此，转换后的输入特征和门控信号具有相同的空间形状。通过元素方式求和的总和是由ReLU激活函数σ1和WT int低维空间映射的控制操作,其次是乙状结肠函数σ2和一个三线的上采样恢复注意力系数矩阵的大小αl i以匹配输入的分辨率特性。将输入特征xl i与注意系数矩阵αl i逐项相乘，得到AG的输出xˆl i。
损失函数：
对于这两个阶段，损失函数有3个部分:
Ldice是促使解码器输出Ppred匹配基本事实分割掩码Ptrue的soft dice损失:
Ll2为L2损耗，应用于VAE分支输出Ipred以匹配输入图像输入Iinput:
Lkl表示KL散度，它被用作一个VAE惩罚条款，用来诱导估计的高斯分布接近标准高斯分布:
其中N是体素的数量。将超参数权重设置为0.1，以在dice和VAE损失项之间达到良好的平衡。
实现细节：
网络在Pytorch中实现，并在四个NVIDIA P40 GPUs上进行训练。
优化：
使用初始学习率为lr0 = 0.0001的Adam优化器来更新权重。根据以下公式逐步降低学习率:
论文中训练epoch为300
数据预处理：
在将输入数据输入第一阶段网络之前，通过对每个患者的每个磁共振成像模态应用强度归一化来预处理输入数据。用平均值减去数据，再除以非零区域的标准偏差。在第二阶段，将来自第一阶段网络的分割图裁剪为每个患者的128×128×128大小的片，同时确保该片包括大多数肿瘤体素。补片与标准化的核磁共振成像图像连接(在数据增加后，在相同的位置裁剪)，并馈送到第二阶段网络进行训练。
数据增强：
为了降低过度拟合的风险，使用了三种数据增加策略。首先将训练数据随机裁剪成160×192×128的大小，然后送入第一级网络。此外，在这两个阶段中，将输入数据的强度随机移动每个通道标准偏差的[0.1，0.1]中的一个值，并将输入数据的强度随机缩放[0.9，1.1]中的一个系数。最后，在两个阶段中，以50%的概率沿着每个3D轴应用随机翻转。
训练数据扩展：
由于两个阶段的训练过程是独立的，所以可以选择几个第一阶段训练的竞技表现模型，并使用它们的分割结果作为训练第二阶段网络的训练数据。这种策略用更长的训练过程来换取更好的模型性能和结果的稳定性。具体来说，选择了6个单独的第一阶段模型(不同时期，具有不同的训练验证划分)，并将它们的分割结果组合成一个广泛的数据集，以训练第二阶段网络(图4)。请注意，训练验证划分是基于患者id的。属于同一患者的6个分割结果必然被分组到同一组中。还尝试使用单个模型的分割结果来训练第二阶段网络，但是与第一阶段网络相比，只获得了轻微的改进。
后处理：
观察到当预测的ET体积特别小时，算法倾向于将TC体素错误地预测为ET。在后处理中，根据作者总结的经验，当预测的ET体积小于500体素时，用TC代替ET。
Ensemble：
使用多数投票来进行模型组装。特别地，如果一个体素在多个类别中具有相等的投票，则基于每个类别的平均概率来确定该体素的最终预测类别。
三、网络模型主要应用及结果
4.1 实验中使用的分割数据集：
BraTS2020年训练数据集包括259个HGG案例和110个LGG案例。所有图像形态(T1、T1c、T2和T2-FLAIR)以240×240×155体素的图像尺寸和1毫米各向同性分辨率共同配准。训练数据带有注释，而验证数据集(125个案例)和测试数据集(166个案例)不带注释。参与者可以通过将预测的分割量上传到组织者的服务器来评估他们的方法。验证评估允许多次提交，而最终测试评估只允许一次提交。
4.2主要结果
BraTS 2020的验证数据集包括125例病例，未提供肿瘤亚型(HGG/LGG)或肿瘤子区域注释。表1报告了官方平台(https://IPP cbica upen edu/)评估的验证数据集的每类Dice得分和Hausdorff距离的分割结果。
通过比较两个阶段的第190代模型的分割性能，看到第二阶段网络的存在对准确度的改善对于TC比对于WT更明显，并且用扩展的训练数据训练第二阶段网络进一步提高了TC的dice分数。
第二阶段网络作为一种性能提升组件，可以将经过扩展数据训练的第二阶段网络添加到任何第一阶段模型中，以提高分割性能。具有扩展数据的第二阶段网络还减少了不同时期模型间的性能差异。表2显示，在第二阶段，TC的Dice得分和Hausdorff距离的标准差(SD)分别减少了68%和93%。SD的计算是基于所有训练的非集成模型的性能。还观察到，第二阶段网络显著降低了ET的Dice分数和Hausdorff距离的变化，但这种改善在后期处理后不再存在。
表2：验证数据的非集成模型之间的性能差异。
（注:变异度量是根据9个阶段1模型和37个阶段2模型的结果计算的，没有后处理。）
BraTS 2020测试数据集包含166个没有提供肿瘤注释的病例。我们对该数据集的分割结果如表3所示。
表3：对测试数据进行分割。
Attention Map：
最优尺度下的注意矩阵以热图的形式呈现，红色表示高权值，蓝色表示低权值。在前几个训练时期，观察到AGs掌握了肿瘤的位置，同时将高权值分配给灰质。随着训练的进行，分配给非肿瘤区域的权重逐渐减少。AGs还表明，该模型通过逐渐减少分配给肿瘤边界的权重，避免了肿瘤边界周围体素的误分类。
前三列分别显示了训练阶段3、20和115时的注意力图。第四列显示了从BraTS 2020训练数据集中提取的带有ground-truth注释的T2-形态示例图像。该模型逐渐学会分配较低的权值给非肿瘤区域和肿瘤边界。
六、总结：
本文提出了一种基于VAE和AGs的级联网络三维MRI脑肿瘤分割方法。结果表明，第二阶段网络改善并稳定了所有三个肿瘤亚区域的预测，特别是对TC和ET(后处理前)。第二阶段的网络还可以为进一步的模型集成提供更合格的候选模型。在本研究中，利用多个第一阶段模型的分割结果来训练第二阶段网络。虽然这有助于提高模型的预测性能，但作为权衡，它显著增加了训练时间。因此，这种技术可能不适用于计算资源和研究时间有限的场合。此外，从表1可以看出，即使扩展的训练数据不包括第一阶段第190个epoch模型的输出，仍然可以使用训练后的第二阶段模型获得比第一阶段模型更好的结果。这说明该策略训练的第二阶段网络在不同时期的模型之间具有通用性。