ECCV 2022 | Med-DANet:用于3D医学图像的动态网络
-
论文标题:Med-DANet: Dynamic Architecture Network for Efficient Medical Volumetric Segmentation
论文地址:https://arxiv.org/pdf/2206.06575.pdf
代码地址:https://github.com/Wenxuan-1119/Med-DANet
摘要
- 现存问题: 对于三维医学图像(例如CT和MRI)分割,在临床病例中分割每个切片的难度差别很大。以往对逐层体积医学图像分割的研究通常使用相同的2D深度神经网络来分割相同情况下的所有切片,忽略了图像切片之间的数据异质性。
- 解决方法: 重点研究了多模式3D MRI脑肿瘤分割,并提出了一种基于自适应模型选择的动态架构网络Med DANet,以实现有效的准确性和效率权衡。对于输入的3D MRI体积的每个切片,提出的方法通过决策网络学习特定于切片的决策,以从预定义的模型库中动态选择合适的模型,用于后续的2D分割任务。
- 实验结果: 在BraTS 2019和2020数据集上的大量实验结果表明,提出的方法在模型复杂度大大降低的情况下,实现了与之前最先进的3D MRI脑肿瘤分割方法相当或更好的结果。与最先进的三维方法TransBTS相比,该框架在不牺牲精度的情况下将模型效率提高了3.5倍。
算法
如上图所示,整个模型框架由一个非常轻量级的决策网络(D)和一个模型库(B)组成,该模型库包含n个不同的2D医学图像分割网络(M1,M2…,Mn)。模型库中的模型应在参数数量和计算成本方面具有多样性。
粗略地说,决策网络将综合考虑每个模型的分割精度和效率,做出最合适的选择。对于模型库,可以包括任何2D网络,以满足各种精度和效率要求。
通过决策网络学习特定于切片的决策,以从模型库中动态选择合适的模型,用于后续的分割任务,如公式所示
其中x表示输入图像,y是相应的预测。D◦B表示在模型库B中取索引为D的匹配元素,计算细节将在下一小节中解释
Dynamic Selection Policy
减少了ShuffleteV2的通道大小,得到了一个非常轻量级的分类网络作为决策网络,因此其计算开销在整个框架中可以忽略不计。
决策网络承担一个n+1类分类任务,n+1类是指n个candidate网络和一个skip过程(可见上图)。因此,决策网络和模型库可以分别表示为
这里θ表示决策网络的所有参数,ˆD是D(x)的预测。M1∼Mn表示candidate网络,∅指示skip过程
具体来说,当遇到只有背景的切片(背景切片,病变被视为前景)时,决策网络将选择直接跳过后续分割过程。
在训练过程中,决策过程的监督GT信号计算如下:
其中Si和Fi分别是模型训练期间候选模型Mi的Dice分数和FLOPs。Pf表示前景像素数(分割目标的所有像素)。具体来说,如果前景像素数小于1(即Pf<1),则将认为当前切片没有任何病变区域,在推理过程中应直接跳过(即相应的监督为0)。
给定由D(x)预测的ˆD和由等式4计算的真值D,应用加权交叉熵损失来监督我们的决策网络(如上图所示)。这允许网络学习跳过数据集中的纯背景切片(将所有像素直接分配到背景类,而不经过分割模型),并全面测量不同模型对单个切片中分割目标的准确性(Si)和效率( Fi)。在实践中,跳过过程是必不可少的,并且可以广泛应用,因为一些几乎无法捕获任何图像内容的背景切片在医学卷中非常常见,在这些背景切片上投入太多计算是没有意义的。此外,分割目标的识别难度因切片而异,动态选择不同复杂度的分割模型更有效。
Training and Inference Strategy
训练
整个框架的训练过程包括两个步骤。
- 首先,分割模型的集成训练。为了节省训练时间成本,对n个分割模型进行联合训练,使所有模型的Dice损失平均值最小,并同步执行梯度反向传播
其中C表示数据集的分割类数,dicelossj是候选分割模型Mj(j)的Dice损失∈{1,2,…,n})和LossB是训练模型库时的 总损失。
- 之后,根据加权交叉熵损失训练决策网络:
其中,WCE是加权交叉熵的缩写,di和ˆdi分别是决策网络为模型候选i预测的gt和logit,wi表示相应的损失权重
Inference
在上述两步训练阶段之后,将经过良好训练的决策网络和预定义的模型库依次级联,以在推理阶段实现最终的模型结构。
给定一个2D切片作为输入图像,超轻量级决策网络将决定跳过当前切片,或根据当前切片的分割难度在模型库中选择最合适的分割网络。在决策网络做出特定选择后,将直接跳过当前切片(即输出相应的全零分割图)或由模型库中包含的单个激活分割网络进行分割。
这样,论文方法实现了一个动态切片相关框架,大大提高了效率。一方面,与之前提出的静态结构轻量级网络相比,Med-DANet对不同的输入进行了动态结构调整,而不是平等地对待所有输入。另一方面,与先前提出的利用预测置信度来确定级联架构是否需要提前退出的动态方法相比,高效Med-DANet可以以一次性方式实现精确分割。
实验
BraTS2019
BtaTS2020
消融实验