CN114332101A - 一种基于遗传算法的U-Net网络模型剪枝方法 - Google Patents

Abstract

本发明公开了一种基于遗传算法的U‑Net网络模型剪枝方法。本发明首先对数据集进行扩增处理，其次将遗传优化算法与U‑Net模型相结合，将U‑Net模型的损失函数引入到遗传算法的适应度指标中，而优化的目标则是U‑Net模型的各层通道数的压缩率组合；随后通过遗传算法获得所需的最佳U‑Net网络的压缩率组合，并对网络结构进行剪枝处理。本发明与剪枝前的U‑Net网络模型相比，在三个数据集以及各自的扩增数据集的分割结果上，模型剪枝前后对结果并没有较大差异，但整个U‑Net网络的参数量以及计算量等都有明显的降低，达到了网络模型剪枝的效果。

Description

一种基于遗传算法的U-Net网络模型剪枝方法

技术领域

本发明涉及一种U-Net网络模型剪枝方法，具体是将深度学习中的U-Net网络与遗传算法相结合进行剪枝的一种方法。该方法实际应用在了视网膜血管图像的分割中。

背景技术

视网膜血管的提取和形态学属性的表征，如直径、形状、扭曲度和分叉，可用于筛选、评估和治疗不同的眼部异常。随着近几年人工智能技术的快速发展，神经网络模型在图像分割领域得到了广泛的应用。但是现在主流的神经网络模型均有着很复杂的网络结构，因此模型的训练时间要大大增加。随着网络模型发展的越来越深，在网络性能得到提升的同时，随之而来产生了另一个问题：越来越复杂的网络结构带来了更大的计算资料的消耗，而且在卷积核数量的设计上的不合理设置更会导致卷积层参数未得到充分应用，最终浪费掉大量的计算资源。近年来用于医学图像分割领域的U-Net网络模型因其独特的结构有着不错的分割效果，但是对于U-Net的各层通道数设置，往往依赖于研究者的手工设置，这使得网络的分割效果在很大程度上依靠人工经验。针对这种现象，在分割视网膜血管的工作上，本文提出了一种利用遗传算法进行U-Net通道自动剪枝的方法进行处理，目的在于通过通道剪枝获得一个更加轻量化的U-Net网络结构，在性能上拥有不输于原始结构的效果。

发明内容

本发明针对现有U-Net网络模型在视网膜图像分割上存在的参数冗余问题，提出了一种基于遗传算法的U-Net网络模型剪枝方法，该方法与以往大多数通道修剪方法不同的是，所进行修剪的是U-Net网络结构中每层的通道数量，而不是对“非重要”通道进行选择性修剪。通过对每层设置一个具体的压缩率，利用遗传算法的快速随机性，能够尽可能地得到一个鲁棒性较好的最优压缩率组合。

按照本发明提供的技术方案，提出了一种基于遗传算法的U-Net剪枝方法，包括以下步骤：

步骤1、对训练集图像进行水平翻转、垂直翻转和多角度旋转；

步骤2、对彩色图像进行图像的预处理操作；

步骤3、对U-Net网络模型构建新的损失函数，使其更加适合于存在像素不平衡问题的视网膜图像数据集的训练；

所述的新的损失函数表达式为：

Loss＝L_dice+λL_r#(1)

L_r为交叉熵函数，表达式如下：

其中TP和TN分别是真阳性和真阴性像素的个数；N_p和N_n分别是目标像素和非目标像素的个数；y是标签值，y＝1,为分割目标，y＝0,为背景；p是像素的预测概率值；

L_dice为Dice系数表达式，如下：

其中N是像素的数量值；p(k,i)∈[0,1]，q(k,i)∈[0,1]分别是像素点k类的预测概率和真实标签，λ为系数；

步骤4、将U-Net网络模型对称的四层卷积层的数量作为优化目标，即遗传算法的优化结果为一组含四个压缩率数值的组合；采用遗传优化算法对该组合进行寻优，最终获得一个最佳的压缩率组合，每层卷积层数量进行相对应的修改操作后，便得到了一个剪枝后的U-Net网络模型；

步骤5、将U-Net网络的损失函数作为参数融合到遗传算法的适应度函数中，使两者存在一个反比关系，并通过适应度函数来扩大损失函数之间的差距；

步骤6、通过对遗传算法的不断迭代，得到所需的最优压缩率组合，并在三个数据集上分别进行U-Net网络模型剪枝前后的实验对比。

作为优选，对训练集图像进行水平翻转、垂直翻转使数据量扩增4倍；对训练集图像进行多角度旋转使数据量扩增8倍。

作为优选，所述的步骤4中，压缩率组合如下：

α＝[α₁，α₂，α₃，α₄]#(4)

式中α_β的取值范围为[0.1,0.2,…,0.8]；α作为遗传算法的输入参数。

作为优选，所述的U-Net网络结构如下：

在网络的编码结构，使用了四层编码块，其中每层编码块包括两个卷积层，一个批归一化层、LReLU激活层以及最大池化层；

在网络的解码结构，与编码结构相对应，使用了四层解码块，其中每层解码块包括两个卷积层，一个上采样层、BN层以及一个跳跃连接层。

作为优选，所述的步骤5中，遗传算法中的适应度函数表达式为：

Fitness＝e^{20*(0.3-Loss)}#(5)

其中Loss为U-Net网络模型的损失函数。

作为优选，遗传算法与U-Net网络模型结合的伪代码如下：

本发明与现有技术相比，有益效果如下：

相比于U-Net网络模型原始结构，本发明方法对模型结构进行了相应的剪枝处理，在降低了网络参数量的同时，在分割结果上并未造成明显的效果下降，达到了较好的剪枝预期；相比于传统手工设置U-Net各层通道数的方法，本发明方法具有自动优化获取最佳通道数的优点，并且经过遗传算法的一系列优化运算，所得结果更具有说服力。

附图说明

图1为本发明的流程图；

图2为本发明的U-Net网络模型结构图；

图3为卷积层剪枝示意图；

图4为适应度函数值变化图；

图5为U-Net剪枝前后在DRIVE扩增数据集上的结果对比图。

具体实施方式

下面结合具体实施例对本发明做进一步说明。以下描述仅作为示范和接收，并不对本发明作任何形式上的限制。

如图1所示，本发明的实施例实现步骤如下：

步骤1、对现有的公开数据集DRIVE、STARE和CHASE_DB1的训练集进行数据扩增操作，具体为对图像进行水平翻转、垂直翻转和多角度旋转，将数据量分别扩增到原来的4倍和8倍。该步骤的目的一方面是为了提供更多的训练集用于模型训练，另一方面也是为了验证在使用不同数量的训练集训练时，U-Net网络模型在剪枝前后是否仍然能保持分割结果的稳定性。

步骤2、对图像的预处理流程为：

对彩色图像进行通道分离，选取血管清晰度较好的绿色通道作为输入图像进行处理，图像尺寸为576×576；

步骤3、对U-Net构建新的网络模型损失函数是为了在训练中更好地反馈训练的过程，从而使模型能够获得较好的训练参数。为此本发明方法考虑到所用到的三个视网膜图像训练集均存在像素分布不平衡的特点，在传统的二分类交叉熵损失函数的基础上，引入了Dice系数，两者进行加权结合形成新的损失函数，具体表达式如式(1)所示，U-Net网络模型结构如图2所示。

步骤4、对U-Net网络模型的对称结构进行对称剪枝，剪枝的目标是每层包含大量参数的卷积层，降低卷积层的数量则是剪枝的重点。在原始结构数量的基础上，每层设置一个压缩率参数，则得到了一个包含四个参数的压缩率组合，后续采用遗传优化算法对该组合进行寻优，最终获得一个最佳的压缩率组合，每层卷积层数量进行相对应的修改操作后，便得到了一个剪枝后的U-Net网络模型。卷积层的剪枝操作如图3所示。

步骤5、为了评价不同的压缩率组合所产生的剪枝效果，将U-Net网络模型训练所得的最小损失函数作为参数输入到遗传算法的适应度函数中，使两者存在一个反比关系，并通过适应度函数来扩大损失函数之间的差距，以便于在遗传算法中能够较好地进行下一代的迭代计算。在U-Net模型的相关参数设置上，优化器使用了随机梯度下降算法(SGD)进行模型的参数训练优化，初始学习率设为0.1，随着训练迭代次数降低，每二十个epoch乘以0.1，epoch设为60。为了加快遗传算法寻优速度，将输入图像切分为48×48的像素块进行训练。

步骤6、对于遗传算法，将变异率和重组概率分别设为0.01和0.8，种群数量设为100，迭代次数为70。实验表明，压缩率组合在迭代45代以后便已经达到了算法的最优解，不再更新。具体的适应度变化如图4所示。

表1U-Net网络模型剪枝前后对比

表1列举了U-Net网络模型在进行剪枝前后的相关指标的对比，可以看出在参数量、计算量和模型大小等指标上都有了不同程度的降低，达到了剪枝的主要目的，降低网络模型训练中所需的计算和存储能力。另外，表2列出了在DRIVE测试集中，U-Net网络模型剪枝前后不同评价指标的分割对比，DRIVE4倍扩增数据集和8倍扩增数据集的分割对比图分别如图5(a)、(b)所示。

表2U-Net网络模型剪枝前后在DRIVE数据集上的分割结果

Claims (6)

1.一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于，包括如下步骤：

步骤1、对训练集图像进行水平翻转、垂直翻转和多角度旋转；

步骤2、对彩色图像进行图像的预处理操作；

步骤3、对U-Net网络模型构建新的损失函数，使其更加适合于存在像素不平衡问题的视网膜图像数据集的训练；

所述的新的损失函数表达式为：

Loss＝L_dice+λL_r#(1)

L_r为交叉熵函数，表达式如下：

其中TP和TN分别是真阳性和真阴性像素的个数；N_p和N_n分别是目标像素和非目标像素的个数；y是标签值，y＝1,为分割目标，y＝0,为背景；p是像素的预测概率值；

L_dice为Dice系数表达式，如下：

其中N是像素的数量值；p(k,i)∈[0,1]，q(k,i)∈[0,1]分别是像素点k类的预测概率和真实标签，λ为系数；

步骤4、将U-Net网络模型对称的四层卷积层的数量作为优化目标，即遗传算法的优化结果为一组含四个压缩率数值的组合；采用遗传优化算法对该组合进行寻优，最终获得一个最佳的压缩率组合，每层卷积层数量进行相对应的修改操作后，便得到了一个剪枝后的U-Net网络模型；

步骤5、将U-Net网络的损失函数作为参数融合到遗传算法的适应度函数中，使两者存在一个反比关系，并通过适应度函数来扩大损失函数之间的差距；

步骤6、通过对遗传算法的不断迭代，得到所需的最优压缩率组合，并在三个数据集上分别进行U-Net网络模型剪枝前后的实验对比。

2.根据权利要求1所述的一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于：对训练集图像进行水平翻转、垂直翻转使数据量扩增4倍；对训练集图像进行多角度旋转使数据量扩增8倍。

3.根据权利要求1所述的一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于：所述的步骤4中，压缩率组合如下：

α＝[α₁，α₂，α₃，α₄]#(4)

式中α_β的取值范围为[0.1,0.2,…,0.8]；α作为遗传算法的输入参数。

4.根据权利要求1所述的一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于：所述的U-Net网络结构如下：

在网络的编码结构，使用了四层编码块，其中每层编码块包括两个卷积层，一个批归一化层、LReLU激活层以及最大池化层；

在网络的解码结构，与编码结构相对应，使用了四层解码块，其中每层解码块包括两个卷积层，一个上采样层、BN层以及一个跳跃连接层。

5.根据权利要求1所述的一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于：所述的步骤5中，遗传算法中的适应度函数表达式为：

Fitness＝e^{20*(0.3-Loss)}#(5)

其中Loss为U-Net网络模型的损失函数。

6.根据权利要求1所述的一种基于遗传算法的U-Net网络模型剪枝方法，其特征在于：

遗传算法与U-Net网络模型结合的伪代码如下：

Images