一种医学图像分割模型压缩方法
技术领域
本发明涉及医学影像处理领域，尤其涉及一种基于神经网络搜索和知 识蒸馏的医学图像分割模型压缩方法。
背景技术
医学图像分割任务一直是计算机视觉和自然领域的研究热点。随着卷 积神经网络(Convolutional Neural Networks,CNNs)的快速发展和应用，越来 越多基于深度学习(Deep Learning,DL)的医学分割模型被提出，并在许多疾 病分割任务中取得了不错的成果。一方面，神经网络的层数越来越多；另 一方面，医疗设备的发展也提供了更高分辨率的数据。这使得医学图像分 割任务的效果越来越好，然而模型的体量却越来越庞大，这不便于模型的 应用和在硬件环境上的部署。因此，医学图像分割模型的体量和计算成本 有待通过模型压缩进一步优化。
一个好的网络结构是模型压缩的关键。目前医学分割模型的卷积核数 量和通道数量都是固定的，其中包含了大量的冗余。为了去除卷积通道数 量的冗余，直接人工选择减少卷积核的数量是一种可选方案，但这会使得 模型性能大大降低。而使用神经网络搜索(Neural Architecture Search,NAS) 则可以在基本网络结构提供的搜索空间内寻找整体性能最好的结构子网络。
目前医学图像分割任务中，最常使用对称的编码-解码结构，为了避免 在剪枝过程中丢失这种对称性，解码过程的特征图与编码过程的特征图会 在规模上一一对应。在模型压缩过程中，相比较保持这种对称性而言，丢 失网络结构对称性会大大地降低分割效果。
同时，由于搜索空间的结构子网络均为独立的深度学习分割模型，而 每个模型构建过程中有些局部信息是多次重复使用，因此可以使用权值共 享策略，让结构子网络共享一些权重，降低训练所需的时间成本和训练资 源。通过神经网络搜索得到的结构子网络需要训练才能用于医学图像分割 任务。然而，由于子网络的规模比基本网络规模要小，从头训练子网络是 很难达到和基本网络相当的结果。
发明内容
本发明针对上述技术问题，提供一种基于神经网络搜索和知识蒸馏的 医学图像分割模型压缩方法，引入对称-NAS编码-译码结构，保证搜索空 间中所有的子网络都具有这种对称性。同时，让子网络模型近似原来的基 本模型，这里基本模型作为教师模型，压缩子模型作为学生模型，使用知 识蒸馏的方法将医学图像分割模型的中间表示，从教师模型迁移到学生模 型，保证分割效果的同时减少医学图像分割模型的计算成本。
为了实现上述目的，本发明提供如下技术方案：
本发明提供一种医学图像分割模型压缩方法，包括以下步骤：
S1、收集医学影像数据库中的数据；
S2、进行数据预处理；
S3、针对一个医学图像分割基本模型，根据模型中每个位置所使用的 卷积核的数量来构建搜索空间，针对分割网络的编码-解码结构，使用对称 神经网络搜索在该搜索空间中寻找一个计算量小、分割精度高的子网络， 且其编码-解码结构均是对称的；
S4、在遍历整个搜索空间时，使用权重共享方法以减轻计算成本和训 练资源；
S5、在网络的训练过程中使用知识蒸馏方法，以基本模型作为教师模 式，压缩子网络作为学生模型，实现二者之间的知识转移。
上述的医学图像分割模型压缩方法，步骤S2中数据预处理包括运动校 正、空间标准化、灰度归一化、边缘去除、尺寸裁剪和中心裁剪。
上述的医学图像分割模型压缩方法，步骤S3中对称神经网络搜索的具 体步骤如下：
S301、在分割网络中根据编码过程中卷积层通道个数的选择来构建搜 索空间，搜索空间中每层中卷积通道数的配置为{c1,c2,...,cK}，其中K表示待 剪枝的第K个网络层，ci表示第i个网络层的卷积通道数；
S302、通过网络搜索得到的最优子网络的通道配置形式如下：其中，Ft是计算约束；
S303、利用编码过程中卷积通到个数与解码过程中卷积通到个数之间 的关系，在解码过程中生成相应数量的卷积通道，最终得到编码-解码结构 是对称的子网络。
上述的医学图像分割模型压缩方法，步骤S4中权值共享方法的具体步 骤如下：
S401、基于基本模型的个通道，给定网络结构配置为
S402、通过从一次即用网络(once-for-all,OFA)中的相应权重中提取 第一个通道作为结构子网络的权重；
S403、在训练过程的每次前向传播时，随机选择一个具有一定通道数 配置的子网络，计算其输出和梯度，根据学习目标来更新提取的权重，同 时其他部分的权值被冻结，不参与该次前向传播；
S404、网络结构训练完成后，遍历整个搜索空间进行测试，找到综合 性能最好的模型；
S405、用该子网络进行微调，得到最优结构子网络。
上述的医学图像分割模型压缩方法，步骤S5中知识蒸馏的具体步骤如 下：
S501、采用基本模型作为教师模型，压缩子模型作为学生模型；
S502、以知识蒸馏的最终目标来训练学生模型；
S503、将分割模型的中间表示从教师模型迁移至学生模型： L＝Lseg+λdistillLdistill，其中，Lseg是医学图像分割误差，Ldistill是教师模型和学生 模型迁移知识时的蒸馏误差，超参λ用于控制蒸馏目标的重要性。
与现有技术相比，本发明的有益效果为：
本发明提供的医学图像分割模型压缩方法，针对一个医学图像分割基 本模型，首先根据模型中每个位置所使用的卷积核的数量来构建搜索空间。 针对分割网络的编码-解码结构，使用对称神经网络搜索在该搜索空间中寻 找一个计算量小、分割精度高的子网络，且其编码-解码结构均是对称的， 以保证分割性能。其中，在遍历整个搜索空间时，使用权重共享策略以减 轻计算成本和训练资源。最后，在网络的训练过程中使用知识蒸馏方法， 以基本模型作为教师模式，压缩子网络作为学生模型，实现二者之间的知 识转移。本发明通过神经网络搜索和知识蒸馏，使得医学图像分割模型在 保证分割效果的前提下，构建网络的计算成本大大减少，优化模型结构， 并且该方法能够应用于各种医学图像分割模型。
附图说明
为了更清楚地说明本申请实施例或现有技术中的技术方案，下面将对 实施例中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附 图仅仅是本发明中记载的一些实施例，对于本领域普通技术人员来讲，还可 以根据这些附图获得其他的附图。
图1为本发明实施例提供的以Res-Unet网络为基本模型的模型压缩示 意图。
图2为本发明实施例提供的以脑肿瘤病灶分割为例的优化结果图。
具体实施方式
为了更好地理解本技术方案，下面结合附图对本发明的方法做详细的 说明。
本发明提供一种医学图像分割模型压缩方法，包括以下步骤：
步骤S1，收集医学影像数据库中的数据，本实施例中以脑肿瘤患者的 磁共振影像为例，主要为T1、T1c、T2和FLAIR四种模态。
步骤S2，进行数据预处理，包括运动校正、空间标准化、灰度归一化、 去头皮和颈部去除以及尺寸裁剪。然后对每个被试的3D MRI进行中心裁 剪，保留整个大脑区域，去除边界的黑色区域。
步骤S3，以Res-Unet作为网络的基本骨架，使用可分离卷积作为其中 的卷积层。Res-Unet是残差U型网络。Unet是医学图像分割的一种经典网络， 在本实施例中，在网络结构中引入Residual模块构成了Res-Unet网络。
首先，在该网络中根据编码过程中卷积层通道个数的选择来构建一个 搜索空间，在本实施例中卷积通道数的选择为8的倍数，可能的通道配置为 {c1,c2,...,cK}，其中K表示剪枝的网络层层数。最优子网络的通道配置如下：
其中，Ft是计算约束。
然后，利用编码过程中卷积通道个数与解码过程中卷积通道个数之间 的关系，在解码过程中生成相应数量的卷积通道。最终得到一个计算量小、 分割精度高，且其编码-解码结构是对称的子网络。
步骤S4，在训练网络过程中，使用权值共享方法遍历整个搜索空间。
具体来讲，每次前向传播时，随机选择一个子网络进行激活，而其他 权值被冻结。每个子网络均具有平等的机会被选择和训练。
假设基本模型有个通道，对于给定的通道数量配置通过从一次即用网络(once-for-all,OFA)中的相应权重张量中提取第一个通道，作为子网络的权重。
在训练过程的每次前向传播时，随机选择一个具有一定通道数配置的 子网络，计算其输出和梯度，根据学习目标来更新提取的权重，同时其他 权值被冻结，不参与该前向传播。网络结构训练完成后，只需要遍历整个 搜索空间进行测试，找到综合性能最好的模型即可。最后，利用该子网络 进行微调，实现神经网络搜索的目的。
步骤S5，使用知识蒸馏提升压缩子模型性能。采用基本模型作为教师 模型，压缩子模型作为学生模型，将分割模型的中间表示从教师模型迁移 至学生模型。知识蒸馏的目标为：
其中，St(x)和S't(x)是学生模型和教师模型中选择的第t层的中间特征， T是层数，ft是一个1x1的卷积层，用于将学生模型的特征映射到教师模型 对应特征中相同数量的通道。最终的目标为：
L＝Lseg+λdistillLdistill
其中，超参λ用于控制蒸馏目标的重要性。
以Res-Unet网络为基本模型的模型压缩示意图如图1所示。图1展示 的是两个Res-Unet网络，可以看到这是一个对称的编码-解码结构网络。上 方的Res-Unet网络是基本模型，作为教师模型；下方的Res-Unet网络是通 过权值共享和神经网络搜索得到的一个结构子网络，作为学生模型。通过 知识蒸馏的蒸馏损失最小为目标，将医学图像分割模型的中间表示从教师 模型迁移到学生模型。
以脑肿瘤病灶分割为例的优化结果如图2所示。图中共有4个病例， 每一个病例从左到右依次为脑肿瘤的核磁像、脑肿瘤病灶标签、Res-Unet 基本网络分割结果和模型压缩后的子网络分割结果。从图中可以看到，对 比病灶标签，基本网络和模型压缩后网络的分割性能近乎保持一致，但模 型压缩后的网络参数远远小于基本网络的网络参数，这表明本发明能够在 保证医学图像分割性能的前提下，大大减少模型体量。
本发明设计了一种基于神经网络搜索和知识蒸馏的医学图像分割模型 压缩方法。首先基于对称神经网络搜索寻找一个医学图像分割网络的性能 最佳对称结构的子网络，然后训练该网络时使用权值共享策略优化计算成 本，同时基于知识蒸馏将教师模型的中间表示迁移到学生模型上，进一步 优化计算性能。该方法综合了不同策略，能够在保证医学图像分割效果的 前提下，大大减少模型的计算量，优化模型结构，以便后续模型在实际应 用中的落地和部署。
以上实施例仅用以说明本发明的技术方案，而非对其限制；尽管参照 前述实施例对本发明进行了详细的说明，本领域的普通技术人员应当理解： 其依然可以对前述各实施例所记载的技术方案进行修改，或者对其中部分 技术特征进行等同替换，但这些修改或者替换，并不使相应技术方案的本 质脱离本发明各实施例技术方案的精神和范围。
- 上一篇：石墨接头机器人自动装卡簧、装栓机
- 下一篇：基于人工智能技术的新冠肺炎数据处理系统
