code: yyliu01/TraCoCo: [TMI'24] Translation Consistent Semi-supervised Segmentation for 3D Medical Images (github.com)

paper:Translation Consistent Semi-supervised Segmentation for 3D Medical Images (arxiv.org)

Abstract

三维医学图像分割方法已经取得了成功，但由于获得这种注释的成本很高，因此需要解决对大量体素级注释数据的依赖。半监督学习(SSL)通过使用大型未标记数据集和小型标记数据集训练模型来解决这个问题。最成功的SSL方法是基于一致性学习，它最小化了从未标记数据的受干扰视图获得的模型响应之间的距离。这些扰动通常使视图之间的空间输入上下文保持相当一致，这可能导致模型从空间输入上下文而不是前景对象中学习分割模式。在本文中，我们引入了翻译一致性协同训练(TraCoCo)，这是一种一致性学习SSL方法，它通过改变输入数据视图的空间输入上下文来干扰输入数据视图，使模型能够从前景对象中学习分割模式。此外，我们提出用一种新的可信区域交叉熵(CRC)损失代替常用的均方误差(MSE)半监督损失，提高了训练收敛性，并保持了对共同训练伪标记错误的鲁棒性。我们还将CutMix增强功能扩展到3D SSL，以进一步提高通用性。我们的TraCoCo显示了具有不同主干的左心房(LA)、胰腺ct(胰腺)和脑肿瘤分割(BRaTS19)数据集的最新结果。我们的代码、培训日志和检查点可在 https://github.com/yyliu01/TraCoCo.

1. Introducetion

三维医学图像分割神经网络方法的训练需要大量的体素标注样本集。这些集是通过费力且昂贵的逐片标注过程获得的，因此人们寻求基于小标记集训练方法的替代方法。一个例子是半监督学习(SSL)，它依赖于一个大的未标记集和一个小的标记集来训练模型，一个特别有效的SSL方法是基于一致性学习，它最小化了从未标记数据的不同视图获得的模型响应之间的距离(Ouali等人，2020;Chen et al, 2021)。一致性学习方法的不同观点可以通过数据增强(Berthelot等人，2019)或从不同初始化网络的输出中获得(Tarvainen和Valpola, 2017;Chen等，2021;Ke et al, 2020)。Mean teacher (MT) (Tarvainen and Valpola, 2017;Yu et al .， 2019;Hang等，2020;Wang等，2021;Liu et al .， 2021)将这两种扰动结合起来，并在训练期间平均网络参数，为未标记的数据生成可靠的伪标签。引入了各种3D医学图像分割方案来提高师生方法的泛化性，包括不确定性引导阈值(Yu et al .， 2019;Hang et al, 2020)或多任务辅助(Wang et al, 2021;Luo et al .， 2021a)。然而，师生方案的特定领域转移(Berthelot et al, 2019)会导致两个网络收敛到相似的局部最小值，从而降低网络扰动的有效性。此外，一些硬医学分割案例被教师和学生模型一致地分割成相似的部分，这可能导致训练过程中的确认偏差。这个问题促使引入了共同训练框架，该框架涉及两个模型，这两个模型使用不同的参数初始化，并通过在训练阶段为未标记的数据生成伪标签来相互监督。与师生模型相比，这两个独立模型收敛到相同局部极小值的机会更小。最近的研究方法(Chen et al .， 2021;Ke et al .， 2020)表明，通过两个独立网络之间的交叉监督，协同训练提供了有效的一致性正则化。
即使成功，上述方法也可能无意中从训练数据的空间输入上下文而不是从待分割的前景对象中学习分割模式，这可能会为未标记的数据产生令人不满意的伪标签。例如，模型可以从背景模式中记忆分割特征，即使在传入数据显示变化时也可以进行一致的预测。这个问题在3D医疗数据中更为明显，训练样本数量较少，输入数据维数较大，会增加空间输入上下文与分割结果之间的依赖性，导致泛化效果较差。Lai等人(2021)通过多层感知器(MLP)提取中间嵌入，并通过对比学习对不同输入上下文施加一致性约束。然而，由于潜在的错误预测和缺乏网络扰动(我们在第4.4节中展示了这方面的结果)，这种隐式约束对于3D医学图像是不令人满意的，这对提高模型的泛化能力至关重要。
事实上，我们认为网络扰动结合空间输入上下文的扰动形成未标记样本的不同视图对于减少空间输入上下文与分割预测之间的依赖性非常重要。
本文提出了翻译一致性协同训练(TraCoCo)算法。TraCoCo在输入数据的翻译视图之间的相交区域内强制分割协议，其中应用了额外的对数似然正则化来平衡分割的视觉对象和背景体素的重要性。对于半监督一致性损失，我们提出了可信区域二元交叉熵(CRC)来约束基于模型预测为正的区域和模型预测为负的区域的学习的跨模型预测(Kim等人，2019;Chen et al .， 2020;Rizve等人，2021)，其目标是提高训练收敛性，同时保持对伪标签错误的鲁棒性。此外，当标注的训练集很小(Chen et al .， 2021)时，CutMix (Yun et al .， 2019)已经证明了2D数据的显着一般化改进，因此我们将