一种基于知识蒸馏和对抗训练的三维体素图像分割方法

技术领域

本发明属于图像分割领域，涉及一种基于知识蒸馏和对抗训练的三维体素图像分割方法。

背景技术

随着科技和成像技术的快速发展，三维图像使用越来越广泛，尤其是三维体素图像。体素是数字数据于三维空间分割上的最小单位，体素广泛应用于三维成像、科学数据与医疗影像等领域。在三维体素图像的使用场景上，有大量应用依赖于基础的三维体素图像分割这个基础任务。但是如果在三维体素图像上手动进行分割，非常复杂且低效率。这不仅仅是因为三维体素图像的体积大要在每个切面上都要进行分割，而且精确分割每张切面在对应使用场景中需要专业人士，这就导致会需要大量的高人才人力成本，这样显然不可行。

对于三维体素图像来说，通常有着单个数据大，数据总量少的特点，这也同时导致模型体积会受三维体素图像的影响变大，模型易过拟合。对于传统的机器学习方法来说，提取特征太复杂，通常不可行。基于深度学习的方法大多被应用于二维自然图像分割，如果将基于深度学习的二维自然图像分割方法应用到三维体素图像上，则效率和精度会较低。因为三维体素图像和二维自然图像的特点完全不同，二维自然图像往往单个数据小，数据总量大。所以亟需基于三维体素图像的特点，设计一种新的方法来应用于三维体素图像上。

Hinton等人于2015年首次提出了知识蒸馏的概念，其一般框架主要包括教师网络和学生网络。教师网络的软目标作为总损失的一部分被引入，以诱导对学生网络的培训并实现知识转移。由于硬目标是单热编码的，因此硬目标中包含的信息熵较低，而软目标中包含的信息熵较高，并且不同类别之间存在关系。温度T用于平滑分类结果，获得具有较高信息熵的软目标，并使用软目标监督作为损失的一部分来训练模型。

生成对抗网络(GANs)由Goodfellow等在2014年提出。GANs由两个神经网络组成，一个是生成器，另一个是鉴别器。生成器尝试捕获数据分布，鉴别器估计样本来自训练数据而不是生成的数据的概率。也就是说，生成器试图创建看起来“自然”的图像，并且要求它与原始数据分布尽可能一致。鉴别器的任务是确定给定图像看起来是否“自然”。Pauline Luc等人于2016年首次使用GANs训练语义分割模型，并对原始图像进行像素级密集分类。该方法使用鉴别器从正确标签中区分分割网络的结果，其中损失函数结合了传统的多类交叉熵损失和GANs损失。通过使用对抗网络来提高小规模数据下模型的鲁棒性并防止过度拟合，提高模型的鲁棒性，并使目标图像空间更加连续和平滑。

GANs在语义分割中的使用，真实标签和输出结果之间的差距非常大。因为标签是单热编码的，而输出结果是非单热编码，所以鉴别器可以轻松地直接判断标签和输出结果，因此GANs应用于语义分割的效果有限。为什么造成这种结果？理论上分析出这是由于“不自然”标签引起，所以需要像软目标这样的标签才能充当“自然”数据。

发明内容

为了克服现有技术的不足，本发明提供一种基于知识蒸馏和对抗训练的三维体素图像的分割方法，目的在于构建一种高效的、高鲁棒的三维体素图像分割方法。本发明使用知识蒸馏来扩充数据，针对于没有标签的数据可以获得一个软目标，这一定程度上缓解了三维体素图像数据总量少的问题。使用原始标签训练，这在一定程度上抵消了最终结果，可以弥补教师网络的准确性不足。使用知识蒸馏来进行知识迁移，将从多个大型模型的集合中学到的知识转移到小型模型，实现模型压缩。这样有利于投入实际生产中，且使用对抗训练的方式提升模型的鲁棒性，使得结果图像变得更平滑。知识蒸馏和对抗训练相结合，能改善对抗训练中判别器的易判断性，进一步提升模型的鲁棒性，使对于三维体素图像的分割方法模型，能更易于投入生产落地。

一种基于知识蒸馏和对抗训练的三维体素图像分割方法，分为两个大模块：

(1)训练多个教师网络进行集成学习，多个教师网络将作为先验知识获得软目标用于学生网络的训练；训练教师网络分为两个子模块，第一个子模块为数据预处理部分，主要是对三维体素图像进行随机剪裁、旋转、强度偏移进行数据增强，再用Z分数(Z-Score)标准化处理数据；第二个子模块为对教师网络进行监督训练，网络结构修改3DUNet作为教师网络，然后用图像标签进行监督；保存若干个训练有素的教师模型用作后续的学生网络训练；

(2)对抗训练学生网络，使用教师网络的软目标和图像标签进行监督训练，使用教师网络中间特征图监督，加入判别器构建GANs损失监督训练；训练学生网络分为四个子模块，第一个子模块为数据预处理部分，主要是对三维体素图像进行随机剪裁、旋转、强度偏移进行数据增强，再用Z-Score标准化处理数据；第二个子模块为前向传播到底层进行辅助监督训练；第三个子模块为前向传播最后输出用软目标和图像标签联合监督训练；第四个子模块为构建判别器利用GANs损失监督训练；同样修改3DUNet作为学生网络结构，学生网络结构比教师网络结构层数要少，进行了压缩。

为了达到上述目的，本发明采用的技术方案如下：具体步骤如下：

步骤(1)训练多个教师网络

(1.1)将原三维体素图像进行随机中心剪裁到合适尺寸，再随机旋转、强度偏移，最后用Z-Score标准化进行预处理；

Z-Score的公式如下：

Z-Score为零均值单位方差操作，可加快训练；X为输入图像，X'为输出图像，μ为所有图像均值，σ为所有图像方差；

(1.2)对教师网络进行训练，使用修改后3DUNet的网络结构，使用图像标签端到端的训练教师网络；3DUNet结构主要依赖于三维的卷积操作和三维的最大池化操作，每次进行完三维卷积操作之后都接一个组标准化操作和泄漏修正单元激活函数；不使用批量标准化的原因为三维体素数据较大，每次送入网络的批量只能很小，这时批量标准化几乎没有效果。图像每经过两次卷积操作，就用三维最大池化缩减图像分辨率为原来的一半；上采样使用三线性插值法，接同样的卷积操作，形成一个编码器解码器结构；最后的激活函数使用sigmoid激活函数最终输出分割图像；使用软骰子损失函数对教师网络进行监督训练，因为三维体素图像上通常有分割标签不均衡的情况；若使用softmax激活函数与多分类交叉熵损失函数的话，训练效果不是很好，所以本方法使用软骰子损失函数训练教师网络。最终保存若干个训练有素的教师网络作为先验知识为之后的训练学生网络做准备；

教师网络软骰子损失函数的公式如下：

教师网络软骰子损失函数，L

步骤(2)对抗训练学生网络，使用前文中的教师网络，图像标签，以及构建判别器监督训练；

(2.1)与(1.1)相似，同样使数据的预处理部分，不同的地方是，训练学生网络时，使用额外的无标签数据作为输入图像，因为使用之前得到的训练有素的教师网络得到无标签图像的软损失，这样就利用到无标签图像，一定程度上扩大了可使用的数据集；

(2.2)前向传播到网络底层进行辅助监督，因为教师网络与学生网络的结构非常相似，都是修改后的3DUNet网络结构，所以可以使用教师网络的中间结果辅助监督学生网络训练。因为学生网络比教师模型小，主要体现在层数上比教师模型小，可以通过辅助卷积将学生网络映射到更高的维度与教师网络的中间层特征图进行匹配，这样就可以得到辅助监督的损失函数，使用MSE损失函数辅助监督；

MSE的公式如下：

公式(3)描述为辅助监督的MSE损失函数，

(2.3)前向传播到最后端，使用教师网络软目标和图像标签联合监督训练学生网络。学生网络结构的3DUNet结构与教师网络的非常相似，主要区别就在于学生网络的层数比教师网络的少。通过编码器解码器结构最后输出学生网络的分割结果，对此使用软骰子损失函数衡量学生网络分割结果与教师网络软目标之间的损失，同时也使用软骰子损失函数衡量学生网络分割记过与图像标签之间的损失。前文中说明使用了无标签的图像数据，这里若无标签图像数据如何联合监督呢？通过动态调整联合损失函数的权重解决了这个问题，对于有图像标签的数据使用联合监督，若无图像标签的数据，只使用软目标监督；

学生网络的软骰子损失函数如下：

公式(4)中，

公式(5)中，

公式(6)中，L

(2.4)构建判别器来构造GANs损失，起到正则化效果。我们将原图像与学生网络分割结果或者教师网络软目标拼接起来输入判别器。判别器为一个全卷积网络，中间通过一个平均池化层来提取特征，最后输出二分类结果，代表图像为“真”还是“假”；

公式(7)中，

公式(8)中，

最终学生网络的损失L

本发明的有益效果：本发明使用知识蒸馏扩充了可使用的数据，并且使用知识蒸馏进行了知识迁移，实现了模型压缩。对于三维体素图像的深度学习模型来说，节省了参数，使其更容易部署。使用知识蒸馏和对抗训练结合，使正则化效果更明显，一定程度上解决了三维体素图像因为数据量稀少而造成的过拟合问题。

附图说明

图1为本发明的整体框架。

图2为本发明的教师模型结构示意图。

图3为本发明的学生模型结构示意图。

图4为本发明的判别器结构示意图。

图5为本发明的教师网络训练方法流程图。

图6为本发明的学生模型和判别器对抗训练方法流程图。

具体实施方式

以下结合附图和技术方案，进一步说明本发明的具体实施方式。

下面结合附图和具体实施方式对本发明做进一步的详细说明，目的是使得本领域的技术人员对本发明的构思、技术方案有更完整、准确和深入的理解。

本发明可用于多种三维体素图像的分割任务，本发明的整体框架如图1所示，细分到每个网络模型中，教师模型结构如图2所示，学生模型结构如图3所示，判别器结构如图4所示。

本实施方案应用到了三维核磁共振脑部成像图的肿瘤分割任务，所论述的具体实施方案仅用于说明本发明的实现方式，而不限制本发明的范围。

下面主要针对三维核磁共振脑部成像图的肿瘤分割任务对本发明的实施方案进行详细说明，任务主要分为教师网络的训练(如图5所示)和学生网络的训练(如图6所示)：

(1)教师网络的训练：

按照发明内容中的模块一进行数据预处理，将三维体素图像进行随机中心剪裁、随机旋转、强度偏移，再进行Z-Score处理输入为训练的教师模型中。输入的三维体素图像维度为4x128x128x128，最前面的4代表4种核磁共振的模态图，将其拼接得到。图像进行前向传播，经过编码器解码器得到分割图，将分割图与图像标签做软骰子损失，然后使用Adam优化器进行优化训练，从而实现端到端的训练。一直优化直到模型收敛，并用相同流程训练3个相同架构的教师模型，并保存。

(2)学生网络的训练：

按照发明内容中的模块二进行数据预处理，将三维体素图像进行随机中心剪裁、随机旋转、强度偏移，再进行Z-Score处理输入为训练的学生模型中，同时使用未标注的三维体素数据，扩大可使用的数据量。输入的三维体素图像维度为4x128x128x128，最前面的4代表4种核磁共振的模态图，将其拼接得到。对抗训练学生模型，具体来说为交替训练学生网络和判别器，先随机取得一个训练有素的教师模型用来得到教师网络中间结果和软目标。

训练学生模型时，输入原三维体素图像到学生模型和教师模型中，前向传播，前向传播过程中可以得到3个损失，在解码器末端将学生模型分割结果与教师模型软目标做软骰子损失得到loss1，在解码器末端将学生模型分割结果与图像标签做软骰子损失得到loss2，在学生模型和教师模型的最底层将学生模型的中间结果与教师模型的中间结果做MSE损失得到loss3。之后将学生网络分割结果与原图像拼接输入判别器中，前向传播，得到二分类结果，做GANs损失得到loss4，最后loss加权融合使用Adam优化器进行优化训练。这里会固定住教师模型中的参数，因为教师模型已经训练完成，不需要再次训练。

训练判别器时，拼接图像与软目标送入判别器中，前向传播，得到二分类结果，做GANs损失，使用Adam优化器进行优化训练。