【论文笔记】PnP-AdaNet:即插即用的对抗域适应网络,用于跨模态心脏分割的基准

论文地址

代码地址

摘要

        深度卷积网络在各种医学图像计算任务中表现出了最先进的性能。利用来自不同模态的图像进行相同的分析任务具有临床效益。然而,深度模型对不同分布测试数据的泛化能力仍然是一个主要挑战。本文提出了PnP-AdaNet(即插即用对抗域适应网络),用于适应不同形式的医学图像(如MRI和CT)之间的分割网络。我们建议通过以无监督的方式对齐源域和目标域的特征空间来解决显著的域转移。具体来说,域适应模块灵活地替代了源网络的早期编码器层,较高级的编码器层在域之间共享。通过对抗性学习,我们构建了两个鉴别器,它们的输入分别是多级特征和预测分割掩码。我们已经验证了我们的区域适应方法在非配对MRI和CT心脏结构分割。综合消融研究的实验结果表明,我们提出的PnP-AdaNet具有良好的效果。此外,我们在心脏数据集上引入了一种新的基准,用于无监督跨模态域自适应任务。我们将公开我们的代码和数据库,旨在促进未来对这一具有挑战性但重要的医学影像学研究课题的研究。

       索引词——域适应、对抗网络、跨模态图像、心脏分割、基准

I 介绍

       深度学习模型,特别是卷积神经模型网络(CNNs)在过去几年取得了显著的成功,在各种具有挑战性的医学成像问题[1]-[3]上取得了最先进的甚至人类水平的性能。通常情况下,深度网络在数据集上进行训练和测试,其中所有的图像都来自相同的分布。尽管存在过度拟合的风险,但这些模型能够对来自同一领域的新测试数据产生高度精确的预测。然而,人们经常观察到建立的模型在被使用在相关但不同目标域的样本时表现不佳[4]-[6]。对于医学图像计算,场景包括测试和训练图像来自不同的站点[7],[8]或不同的扫描协议[9],[10]甚至不同的成像方式[11],[12]。

       与一般由光学相机获得的自然图像不同,医学领域的一个典型情况是使用各种成像方式,捕捉不同的物理性质。这些不同的模态在疾病的临床诊断和治疗过程中起着互补的作用。例如,磁共振成像(MRI)和计算机断层扫描(CT)已经成为心脏成像不可或缺的工具。具体来说,MRI无电离辐射,在时间空间[13]中以高分辨率捕捉软组织之间的对比度。因此,它可以多参数评估心肌收缩和生存能力。相比之下,CT可以快速成像心脏形态、心肌活力和冠状动脉钙化,具有很高的空间分辨率[14]。在实践中,通常需要相同的图像分析任务,例如MRI和CT都需要分割或量化心脏结构。考虑到注释是非常耗时和昂贵的(例如,训练有素的操作员[15]从MRI或CT进行整个心脏分割需要8个小时),有效地将一种模式训练的模型适应于另一种模式具有临床效益。然而,心脏MRI和CT的表现有很大的不同,对比度和强度直方图明显,如图1所示。在如此重大的领域转移下的无监督领域适应是非常具有挑战性的,也是有待探索的。

图1:交叉模态图像中存在严重的畴移。心脏结构(AA:升主动脉,LV-blood:左心室血腔,LV-myo:左心室心肌)在MRI和CT图像上表现明显不同,但分割掩模相似。 

       早期的无监督域自适应研究集中在通过最小化从源域和目标域提取的特征之间的距离来调整特征空间的分布。例如,最大平均差异(Maximum Mean difference, MMD)与任务特定的损失一起被最小化,以 学习域不变和语义有意义的特征[16]。Long 等人[17]最小化了嵌入在再现核希尔伯特空间中的域特征的MMD。Sun等人[18]提出对齐域之间的特征协方差。最近,随着生成对抗网络(GAN)[19]及其强大扩展[20]、[21]的引入,跨领域的潜在特征空间可以通过对抗学习隐式对齐。值得注意的是,DANN方法通过在域[22]之间完全共享CNN编码器的权值来提取域不变特征。Tzeng等人[23]引入了一种更松散的对抗学习框架,称为ADDA,其中每个域在最后一个分类层之前都有一个专用的编码器。另外,在基于GAN的领域适应方面,另一种解决方案是利用图像到图像的转换,即,基于CycleGAN[21]基础,使用合成数据训练或测试网络。

       对于医学图像计算,对抗学习已经在各种各样的任务[24]-[28]上显示出令人鼓舞的功效。特别是在域适应方面,将学习到的模型推广到未知目标领域的研究越来越多。例如,Zhang等人[29]将x射线图像转换成类似于那些源合成x射线照片的图像,并直接用学习到的源模型对转换后的图像进行测试。同样基于CycleGAN, Jiang等人[12]提出了一种两阶段的方法,首先将CT图像转换为类似MRI数据,然后使用合成数据和有限数量的真实MRI进行半监督肿瘤分割。同时,遵循潜在特征空间对齐的思想,有一系列提取域不变表示的工作。Degel等人[30]使用域鉴别器最小化分割损失,以鼓励左心房分割的超声数据集的特征域不变性。Ren等人[31]利用对抗学习将目标图像的特征分布映射到源域,对不同染色程序获得的组织学图像进行分类。这些工作表明,在特征空间中施加对齐有助于将深度模型推广到多个数据。其中最相关的工作之一是Kamnitsas et al.[9],他们在多级特征空间中通过对抗学习进行无监督域适应,用于脑损伤分割。由于目标数据中使用了未见的MRI序列,因此实验设置具有挑战性。该方法通过共享编码器和对齐多层次特征,在目标区域取得了良好的效果。对于更具有挑战性的MRI和CT交叉模态图像分割,据我们所知,相关文献有限。Valindria等人[32]开发了一种使用未配对MRI和CT进行多器官分割的联合学习方法。Zhang等人[26]提出了跨模态图像翻译方法,以改进合成数据的心脏分割。然而,这些工作并不是针对CNN在跨模态医学图像上的无监督域自适应问题。

      本文研究了多类分割问题中具有挑战性的无监督跨模态域自适应问题。我们提出了一种即插即用的对抗域适应网络PnP-AdaNet,它有效地将目标域的特征空间与源域的特征空间对齐。具体来说,早期编码器被替换为目标域输入,更高的层在域之间共享。在对抗学习中,我们建立了两个域鉴别器,分别连接多层次特征和预测分割掩码。本文是我们先前工作[11]的一个实质性扩展。本文的主要贡献有:

  • 我们解决了医学图像分割中的无监督跨模态域自适应问题。提出了一种新颖的PnP-AdaNet,通过即插即用特征编码器实现对分割CNNs的灵活适应。
  • 通过对抗性学习,我们通过未配对的MRI和CT图像学习我们的模型。为了加强鉴别器的监督,我们在训练过程中聚合多层次特征和分割掩码预测。
  • 我们广泛地验证了我们的方法在多类心脏分割。四种结构的平均Dice回收率为13.2% ~ 63.9%。我们还对关键方法组件进行了综合消融实验。
  • 为了促进未来MRI和CT跨模态区域适应的研究,我们在心脏分割任务中引入了一个新的基准,展示了流行的区域适应方法的性能。

II 方法

        图2是我们提出的PnP-AdaNet方法的概述。通过CNN在源域上学习到的标准分割,我们用域适应模块替换其早期层,同时保留其更高的层,用于对目标域数据进行测试。因此,我们称我们的方法为即插即用域适应框架。自适应模块将目标图像映射到潜在特征空间中源域的分布。这一过程是以一种无监督的方式训练对抗损失的。

图2。我们提出的PnP-AdaNet(即插即用对抗域适应网络)的概览,包括一个源分割网络,一个域适应模块(DAM)和两个鉴别器。多层次激活和预测分割掩码聚合,以对齐潜在特征空间。域路由器用于测试。它选择哪一组早期层连接到更高的层进行分割任务。具体来说,在测试源数据时,它选择使用原始源早期层;在测试目标数据时,它选择使用DAM层。

A. 没有跳跃连接的分割网络

       我们提出的PnP-AdaNet的本质是为每个域建立一个独立的编码器,并在潜在空间中对齐它们的特征分布。考虑到只有早期的层编码器被更新,而更高的层是固定的,不同层的特征空间需要是自包含的,即不能相互混淆。这意味着使用跳过连接的网络架构,例如U-Net[33]和DenseNet[34],不是合适的选择。否则,即插即用设置就会有问题,因为那些特定于领域的低级特性会影响对齐的高级特性空间(应该跨领域共享)。

       为此,我们将我们的分割模型建立为一个扩展网络[35],它可以从一个大的接受域中提取有代表性的特征,同时也保持了特征图的空间灵敏度。局部范围内的残余连接用于简化梯度流动。具体来说,如图2所示,首先将输入图像馈送到一个卷积层,然后进行一系列残差模块(称为RM1-7,每个模块由堆叠的3x3个卷积组成),大小下采样8倍。接下来,为了扩大接受域以提取上下文信息,在DRM8中使用4个扩张卷积层,扩张因子为2。再经过两个卷积层Conv9和Conv10,我们进行上采样以得到分割任务的密集预测。然后立即进行5x5卷积运算,以平滑激活映射。最后,利用softmax层进行概率像素预测。

      形式上,我们用X^s=\left \{ (x_1^s, y_1^s),...,(x_i^s, y_i^s),...,(x^s_{N^s}, y^s_{N^s}) \right \}表示源的注释数据集,其中x_i^s表示图像的一个样本像素,y_i^s为其解剖结构的分类标签。我们通过监督学习来建立一个从输入图像到标签空间的映射M^s。通过最小化由多类交叉熵损失和Dice系数损失组成的混合损失L_{seg},对源域分割网络进行优化。此外,对于属于类别c\in C的样本x_i^s,我们将y^s_{i,c}表示为它的二进制标签,它的可能性预测值是\hat{p}^s_{i,c},标签的预测值是\hat{y}^s_{i,c},整体的分割损失函数是:

 第一项是针对多个解剖结构的Dice损失,第二项是针对单个像素的交叉熵损失。w_c^s是解决类别失衡问题的权重因子。我们结合这两种互补的损失来处理具有挑战性的心脏分割任务。在实践中,我们也尝试只使用其中一种,但性能不如同时使用两种。第三项是分段权重W的L2-正则化,\lambda\beta是权衡权值。

       为了便于标记,我们将在以下小节中省略下标索引i,直接使用x^sy^s表示来自源域的样本和标签。

B. 即插即用适应机制

        在获得在源域上训练的分割网络后,我们的下一步目标是以一种无监督的方式将其适应于目标域。在传统的迁移学习中,常用的方法是将预训练网络的最后几层更新为具有新的标签空间的新给定任务。支持的假设是网络的早期层提取视觉任务通用的底层特征。较高级的层更特定于任务,并学习语义级特征,以执行定义的预测[37]、[38]。相反,对于域适应,跨域定义的任务保持不变。这意味着源域和目标域的标签空间是相同的,例如,我们从MRI/CT图像中分割出相同的解剖结构。基本上,分布在跨模态域之间转移是主要是低级特征(例如,灰度强度)而不是高级特征(例如,几何或语义结构)。

       在这些方面,对于我们的模型,我们设计了一个即插即用的适应机制,即,一组早期的层被替换,而更高的层被重用为新的目标域。基本的直觉是,较高层与共享语义标签密切相关,而各自的早期层编码器在特征空间中执行分布映射,以实现我们的无监督域适应。形式上,得到的源分割模型M^s被看作是由叠加变换组成的分层特征提取器\left \{ M^s_{l_1}, ..., M^s_{l_n} \right \}。语义标签的预测表示为:

         对于来自目标域的输入x^t\in X^t,我们提出了一个表示为的域适应模块(DAM),该模块将x^t映射到与源域对齐的特征空间。我们用d表示适应深度,即在处理目标域图像时,将l_d之前和包括ld在内的层替换为DAM。同时,在领域适应学习过程中冻结源模型的上层,并将其用于目标推理。因此,对目标输入的预测为:

 其中代表DAM,它是一堆卷积层作为特征编码器。在实践中,我们将DAM层配置设置为与替换的源模型早期层集合相同,即\left \{ M^s_{l_1}, ..., M^s_{l_d} \right \}。这是一个合理和安全的实现选择,因为我们可以使用预先训练的源编码器来初始化DAM,而不是随机初始化。这有助于在对抗训练中更稳定的优化,特别是在我们的无监督学习场景中。

        总的来说,我们可以发现所提出的即插即用域自适应机制在测试中是优雅且相当灵活的。在对目标域进行推理时,DAM直接替换源域网络的前d层。对目标域的图像进行处理,并将其映射到源域的特征空间中。这些自适应特征对跨模态域移位具有鲁棒性,并能通过固定的高层正确地转换到标签空间中。还有必要提到的是,即插即用域适应过程不会损害源域的性能。源域的早期层编码路径被保留,较高层在学习时不受影响。因此,我们的PnP-AdaNet可以灵活地测试目标和源域数据,只需要选择输入路径。

C.特征空间对齐的对抗学习

       我们通过无监督的对抗学习来训练即插即用领域适应网络。在GAN的思想中,一个生成器和一个鉴别器组成了一个极大极小二人博弈。生成器的目的是捕捉真实数据的分布,而鉴别器应该识别所呈现的样本是来自真实的分布还是学习到的分布。在我们的PnP-AdaNet中,DAM作为生成器,将输入的目标图像映射到源域的潜在特征空间。域适应模块的目的是对与源域图像编码一致的表示进行编码。因此,可以重用源网络较高部分的固定层,对分割掩码进行语义级预测。由于我们对目标域没有注释,适应过程隐式地由鉴别器监督,即形成对抗学习博弈。

       在我们的框架中,我们建议使用两个鉴别器。具体来说,第一个鉴别器的输入(即图2中绿色部分)是分段器聚合的特征映射数组。该输入具有较高的维数和相对复杂的分布。这里一个自然的想法是,我们将DAM的输出特征连接到鉴别器中。而卷积神经网络具有层次化的结构,某一层的特征依赖于其上一层的激活,并且这些特征会影响到下一层的特征。如果我们只是监测从DAM立即获得的编码特征,潜在空间对齐可能是不稳定的。换句话说,我们不知道那些在适应深度之前的层中的激活是否一致。同时,在适应层l_d中可能仍然存在的小的变化将在更高的层中组合后被放大。为了克服这个问题,我们将来自多个层的激活聚合为到鉴别器的输入。从鉴别器的梯度可以通过多条路径流向DAM,这样可以更严格地监督特征空间对齐。这种学习方式在一定程度上具有深度监督的精神。

       在实践中,我们聚合来自多层的激活,并将它们重塑为相同的分辨率用于通道连接。形式上,我们将所选冻结层中的特征映射称为F_H(\cdot )集,其中H=\left \{ k,...,q \right \}为所选图层索引集。类似地,我们用表示DAM的选定特征图,其中A为选定的层集。这样,目标域的特征空间为,数组(M^s_A(x^s),F_H(x^s))是它们对应的源域。给定(M^s_A(x^s),F_H(x^s))\sim \mathbb{P}^s_{feature}的分布和,将要被最小化的这两个域分布之间的距离被表示为W(\mathbb{P}^s_{feature},\mathbb{P}^t_{feature})。对于稳定训练,我们使用两个分布的Wassertein距离[20],即:

        直接在鉴别器上输入高维激活来校准潜在特征空间是有效和必要的。这对于分类任务可能是好的,但对于需要用精细结构进行像素预测的分割任务可能不是最优的。早期的研究使用GANs进行分割应用(不一定在域自适应设置下),通常将预测的分割掩码输入鉴别器。当预测的分割掩码的形状或结构看起来扭曲(即,不像真正的掩码),鉴别器将施加惩罚。对于分割时的域自适应问题,我们也认为监测预测分割掩模的形状是重要的。

       为此,我们进一步在PnP-AdaNet中包含一个辅助鉴别器,它的输入是源域和目标域的预测分割掩码。在这种情况下,与第一个鉴别器相比,输入更紧凑,语义更明确。我们用S(x^s)\sim \mathbb{P}^s_{mask}S(x^t)\sim \mathbb{P}^t_{mask}表示目标域和源域的分割预测。根据式(4),我们还使用源和目标分布之间的Wassertein距离:

       鉴别器的详细网络结构如图2所示。对于模型配置,特征鉴别器相对于掩码鉴别器更深。

 D.损失函数和训练策略

       在对抗性学习中,DAM与具有上述两个鉴别器的对手进行对抗。我们用D_f代表第一个鉴别器,它以高维特征为输入;用D_m表示第二个鉴别器,它以紧凑预测分割掩模为输入。D_f&D_m对分别隐式估计了W(\mathbb{P}^s_{feature},\mathbb{P}^t_{feature})W(\mathbb{P}^s_{mask},\mathbb{P}^t_{mask})。在学习过程中,鉴别器将尝试从源域和目标域中区分输入。DAM的域适应不仅旨在消除早期层中域特定的模式,而且还禁止在更高的语义层中恢复它们。通过对抗损失函数,生成器 M (DAM)和鉴别器\left \{ D_f,D_s \right \}被共同优化。具体而言,DAM生成器的损失是:

 

 学习鉴别器D_fD_m的损失是:

 其中KD_f, D_m的Lipschitz约束。通过交替更新生成器和鉴别器,DAM可以更有效地从目标数据生成类源特征,以进行域适应。

       在实际应用中,我们首先用标准随机梯度下降在源域上有监督地训练分割网络。我们使用了批量大小为10的Adam优化器,学习率为1\times 10^{-3}。获得分割器后,我们用上述对抗损失训练\left \{ M,D_f,D_m \right \}以进行无监督域适应。首先,我们仅更新20k次迭代的鉴别器,批处理大小为6,作为预训练过程。接下来,我们交替更新生成器和鉴别器。根据训练WGAN[20]的启发式规则,我们更新生成器M一次,每20次迭代更新D_fD_m。在对抗学习中,我们使用RMSProp优化器,其学习率为3\times 10^{-4},每100次联合更新的步进衰减率为0.98。鉴别器权重的权重裁剪为0.03。鉴别器和发生器的损失按因子0.002进行缩放。在所有的卷积层中使用Dropout (rate = 0.25)和批归一化。

       由于对抗损失的学习是非常困难的,我们仔细地调整了实现设置,以稳定优化过程。具体来说,为了避免稀疏梯度导致的不稳定性,我们使用Leaky-ReLU作为激活函数,并在鉴别器内部使用跨步卷积进行下采样,而不是使用常见的ReLU和最大池化的做法。

  • 0
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值