基于模型压缩与重构U-net的胰腺分割

2022-07-21 04:11张俊然
计算机工程与设计 2022年7期
关键词:剪枝胰腺重构

马 豪,刘 彦,张俊然

(四川大学 电气工程学院,四川 成都 610000)

0 引 言

与其它腹部器官相比,胰腺拥有如下特性:其个体之间的形状、大小和位置差异很大,且与周围组织环境之间对比度低。传统的胰腺分割算法大概可分为如下两种模式:从分割模型架构上来讲往往采用双层分割网络来保障胰腺的分割精度;从特征图角度来讲可对胰腺提取不同维度的特征图信息以充分利用标签特征。其算法目的均是为了追求提升胰腺分割精度。但是这些算法缺点在于除了受胰腺自身特性影响,也受限于对分割模型卷积层输出的特征图信息利用不充分,致使最终的分割精度往往不能达到医生满意的程度,因此对于胰腺的分割仍是一个很大的挑战。最近由于深度学习中卷积神经网络对于硬件的苛刻要求和医疗设备算力薄弱的原因,一批专家学者开始应用模型压缩技术来降低卷积神经网络的计算需求与存储空间占用。融合了上述两种手段的实时医学图像分割课题,已然成为热门研究领域[1-6]。综上所述,可见医学上实时胰腺图像分割的研究目的为:在保障分割网络性能的前提下,尽可能减少分割模型对硬件资源的需求。针对上述问题,本文利用通道剪枝技术对分割模型进行充分压缩,同时利用重构U-net对胰腺进行图像分割处理。相较前人,在分割精度与分割模型存储空间占用上同时做出新的突破,更进一步推进了对于胰腺的医学上实时图像分割实现。

1 相关工作

1.1 胰腺分割

对于胰腺的分割算法研究经历了许多阶段,最早期的胰腺分割算法是基于模型形状统计设计,或者运用多图谱技术[7,8]。随后深度学习开始应用到胰腺分割领域上来。从对胰腺分割算法的结构设计方面讲,一开始的胰腺分割算法一般做法都是应用单个分割网络,但是由于胰腺图像所占CT图像比例非常小(约为0.5%),单个分割网络对于胰腺的特征信息提取不够充分,导致分割精度很低,所以后续的胰腺分割算法往往采用双层的粗细分割网络结构来处理胰腺数据。递归显著变换神经网络提出在粗细分割网络之间加入一个循环分割模块,用以将粗分割网络的权重转移到精细分割网络中去,使两个分割网络进行共同迭代优化直至训练结束,最终取得了Dice为84%的成绩[9,10]。

对于胰腺的分割网络从维度上分类,一般分为3种:3D/2D/2.5D分割网络。胰腺分割的一种3D分割网络为首先给3DFCN模型加入了Resnet连接,并提取不同维度特征图进行特征融合,再将其又分别反馈给对应位置初始的特征图,最后使用深度监督模块处理后得到最终分割结果,其在NIH胰腺公共数据集上取得的Dice参数为83.39%。但是问题在于3D的分割网络十分消耗硬件资源,所以现阶段2.5D分割网络在胰腺数据集上应用更为广泛。针对胰腺数据的长短期记忆网络为每次经过3张连续的特征图融合得到一张新的特征图,最终构成分割网络,相比于其它2.5D网络该分割网络所需的参数量更少[11-13]。

前人工作在只追求胰腺分割精度的前提下,目前取得的最好分割精度是由GAN网络实现的。GAN网络在胰腺分割上的应用方法是:用GAN网络融合膨胀卷积以获得更大的感受野,同时加入局部长短期记忆网络来保持相邻切片之间的上下文相关性,接着加入对抗模块用全局分布约束来保持全局空间平滑一致性。该文献在NIH数据集上最终取得的Dice分割精度为88.72%[14]。

1.2 模型压缩

卷积神经网络的数据压缩方法大概可以分为4类:量化、知识蒸馏、剪枝以及紧凑网络设计。其中最主流的两种方法是剪枝和量化。剪枝即基于卷积神经网络中参数的重要程度不同,可以只保留重要参数而剪去不重要的参数或通道,这样可以将网络进行高度稀疏化并获得相近的卷积神经网络效能。根据剪枝的细粒度不同,将剪枝方法分为5类:细粒度剪枝、矢量级剪枝、核级剪枝、组级剪枝和过滤级剪枝[15]。量化在卷积神经网络的压缩和时间加速方面有着巨大作用。对于神经网络参数的压缩又分为对权重压缩和对激活压缩。三元权值网络提出将权重参数约束在[-1,0,1]中取值,这样就摒弃了卷积神经网络参数中浮点数所占用的内存空间[16]。学者为彻底解决浮点运算的需求提出了一种二值化网络,将激活限制[-1,1]中,这样就将权重和激活的网络参数全部固定到一组固定点中,最终在ImageNet数据集上验证了模型的有效性[17]。

2 方 法

由于医学上实时胰腺图像分割的需求为:在不降低或少降低分割网络性能的前提下,尽可能减少分割网络对硬件资源的占用,使其到达一个新的最优平衡。针对上述需求,本文提出了两个创新:①在胰腺分割工作中,提出应用通道剪枝技术,并搭建了一个端到端框架(comprehension prune-reconstructed U-net network,CPRUN),实验结果表明,将模型大小压缩至极小情况下仍保证了非常高的胰腺图像分割精度;②为适配CPRUN模型设计了一种重构的U-net卷积神经网络,该重构模型优点是对其进行剪枝后,在其任一卷积层中所连接的任一输出特征通道都是对胰腺分割具有高贡献度的前提下,充分利用到各卷积层特征通道所提供的不同分辨率特征图信息,以保障胰腺分割精度。

本文提出一个基于通道剪枝技术与重构U-net的胰腺图像分割的端到端框架即CPRUN模型。算法的整体流程包括对于重构的U-net卷积神经网络进行预训练,然后通过通道剪枝算法对重构的U-net卷积神经网络进行模型压缩,接着生成CPRUN模型并进行微调即再训练,直至最终输出模型压缩后的CPRUN模型与胰腺分割结果。该CPRUN模型的运行流程如图1所示。

图1 CPRUN模型的运行流程

2.1 CPRUN模型的通道剪枝算法设计

剪枝有多种粒度不同的剪枝方法:细粒度剪枝往往是对于神经元之间的权重连接进行修剪,其进行剪枝后的模型精度往往最高;而较粗糙的粗粒度的剪枝方法,其缺点是不可避免的忽略相当数量的冗余参数。由于剪枝方法的粒度和与模型精度相互制约,为了使CPRUN模型在模型大小和分割网络精度间取得最优平衡,决定选取通道剪枝算法来对重构的U-net卷积神经网络进行修剪,使其一方面不会受限于细粒度剪枝所需的特定算法库或者硬件制约;另一方面,又避免了粗粒度的剪枝算法所导致的模型压缩效果不够。

本文的方法是定义一个全局的阈值因子α(α∈(0,1)), 当重构的U-net卷积神经网络预训练到接近收敛时,保留其卷积神经网络模型与对应模型参数。剪枝方法的第一步是,将连接在每个卷积层之后的BN(batch normalization)层的权重参数Wij集合起来组成一个集合γ,再对γ中所有元素γk取绝对值并按照大小排序,其数学表达式如式(1)所示

γ=sort({|Wij|})

(1)

其中, |Wij| 表示重构的U-net卷积神经网络中第i层卷积层的第j条输出特征通道所对应的BN层的权重参数再取绝对值, sort() 表示排序函数,所有排序过的 |Wij| 共同构成集合γ;第二步是统计集合中元素γk的个数numberBN, 将其与α相乘后得到集合剪切位置下标γthreshold, 其数学表达式如式(2)所示

γthreshold=α×numberBN

(2)

第三步是对集合γ中的全部元素γk进行处理,当k≥γthreshold时,在集合γ中保留元素γk。当k<γthreshold时,在集合γ中删除元素γk。 其数学表达式如式(3)所示

(3)

最后,对那些删除的集合元素γk,便要将其所代表的Wij所对应的卷积层输出特征通道包括其前后输入输出连接全部删掉。而集合γ中剩余的元素便代表着与其对应相连的卷积层输出特征通道保留满足α阈值因子的要求。最终重构的U-net卷积神经网络转化为CPRUN模型。

为将剪枝算法设计解释清楚,CPRUN模型的剪枝算法的流程见表1。

表1 通道剪枝算法流程

为直观表示通道剪枝算法的实际操作,CPRUN模型的通道剪枝算法设计如图2所示:其中左边虚线方框表示的是重构的U-net卷积神经网络中第i层卷积层中经过剪枝处理后需要被减掉的特征通道,左边实线方框表示剪枝后需要保存下来的卷积层特征通道;Wij表示的是BN层权重,其中虚线椭圆指的是经阈值因子α作用后需要被剪掉的BN层权重,相应的实线椭圆指的是保留下来的权重;Ci1至Cin表示的是第i个卷积层中的n条输出特征通道;图2中左边虚线和实线方框共同构成重构的U-net卷积神经网络的第i层卷积层,右边实线方框共同构成第j(j=i+1) 层卷积层。

图2 剪枝过程

BN层的作用是可以较好解决梯度消失和梯度爆炸问题,CNN(convolutional neural network)往往是在一个卷积层之后加入BN层,同时由前人实验指出,由于BN层的特性,BN层可能是最适用于考虑进行通道剪枝的区域。所以CPRUN模型决定在BN层区域完成通道剪枝操作[18]。在通道剪枝完成之后,经过模型压缩的卷积神经网络模型往往会有一定程度的精度衰减与性能下降,但是通过一定数量epoch的再训练与微调,新的剪枝模型往往可以达到甚至超越原有未压缩模型的精度。所以还需对剪枝后生成的CPRUN模型进行微调与再训练直至收敛并输出分割结果。

2.2 针对通道剪枝算法的重构的U-net卷积神经网络设计

传统U-net分割模型的设计亮点是:通过一种跳跃连接,将分割模型下采样环节中每两个卷积层后输出的特征图信息,连接至上采样环节对应位置的同层信息处,这样的好处是可以利用不同分辨率的同层位置的对应特征图信息来对目标的分割提供帮助。

CPRUN模型的设计目标指明,需要在保障模型压缩程度的情况下,尽可能稳定或略有提高分割网络的精度。针对上述目标,本文提出了一种重构的U-net卷积神经网络。从2.1节可知,由于CPRUN模型考虑将所有卷积层的特征通道按重要性从大到小排序,故最终所确定的剪枝方案对于每一个卷积层的每一个输出特征通道都有明确的删除或保留要求。这样经过通道剪枝算法进行模型压缩后,分割网络的所有卷积层的输出特征通道都是对于胰腺分割贡献度较高的。

至此,重构的U-net卷积神经网络的设计方法为:①为适配通道剪枝技术,重构的U-net卷积神经网络中卷积层的定义方法是:逐个卷积层指定其输入特征通道数和输出特征通道数;②将重构的U-net卷积神经网络中下采样环节中每个卷积层后输出的特征图信息,都连接至上采样环节中对应位置的同层信息处。

由图3展示为CPRUN模型适配设计的重构U-net结构框架:其中Conv(3×3),ReLU表示先进行2D卷积处理,后用ReLU激活函数进行激活。MaxPool2d表示2D池化层,sigmoid表示2D激活函数。其中copy and crop 表示对特征图的复制和裁剪。带CPRUN copy and crop指示的加粗箭头表示重构的U-net卷积神经网络加入的跳跃连接,带copy and crop指示的箭头则表示传统U-net分割模型的跳跃连接,带Upsample指示的箭头表示上采样。从图3不难看出,重构的U-net卷积神经网络的跳跃连接数目是传统U-net分割模型的两倍,更强程度上对重构U-net同层处特征图信息加以充分利用与提取。

2.3 评价参数

Dice系数是医学图像处理领域使用较为广泛的评价指标,它可以非常好体现出预测值和标签之间的相关关系,其定义表示如式(4)所示,这里P指的预测的分割结果,T指的是胰腺的标签

(4)

实验的评价参数主要是参考Dice,但是同时为了验证实验结果的有效性,我们同时采用了精确率Precision和感兴趣区域交并比Iou评价参数进行佐证。精确率Precision的定义如式(5)所示,其中TP(true positive)表示模型预测结果为正例,实际结果为正例;FP(false positive)表示模型预测结果为正例,实际结果为负例

(5)

感兴趣区域交并比Iou的定义如式(6)所示,其中P表示预测结果的集合,T表示真值的集合

(6)

为了表明压缩后的模型与未压缩的模型之间的压缩程度,定义一个评价指标即压缩率ratio, 其定义如式(7)所示

(7)

3 实 验

3.1 实验环境与训练过程

实验在Pytorch上验证了CPRUN模型的有效性,模型的预训练epoch次数定为30,微调epoch次数定为60。实验环境是CPU为I7-7800X;显卡为1080TI;64 G内存;Pytorch版本为1.5.1;CUDA版本为10.1;显卡驱动版本为430.50。

图3 为CPRUN模型适配设计的重构U-net模型

实验中模型的训练过程包括重构U-net的预训练与后续CPRUN模型的微调,预训练与微调过程中分割模型的损失函数设计均为-Dice, 其训练过程中的超参数设置见表2。

模型的训练过程中,重构U-net的精度收敛曲线与损失收敛曲线如图4所示,CPRUN模型微调的训练过程曲线如图5所示,其中图4和图5中训练精度表示训练过程中模型的精度,训练损失函数表示训练过程中模型的损失函数,训练次数指的是对训练集所有数据训练的次数。

表2 训练过程的超参数设置

图4 重构U-net预训练过程曲线

图5 CPRUN模型微调的训练过程曲线

3.2 重构U-net的实验设计

为了观察重构的U-net分割网络与传统的U-net分割网络之间的模型大小和性能之间的差距,设计了这两种模型的对照实验,其结果见表3。通过观察可以发现,重构的U-net分割网络优点是Dice大于传统的U-net卷积神经网络约一个百分点,缺点是重构的U-net分割网络模型大小略大于传统的U-net分割网络。虽然重构的U-net分割网络模型大小略有提升,但是通过后续应用剪枝技术,其模型大小会被压缩;同时2.2节阐述了重构的U-net分割网络是为了适配CPRUN模型的剪枝策略设计的,即增加的跳跃连接是为了充分应用后续的剪枝操作所带来的高贡献度的特征通道提供的不同维度特征图信息。故重构的U-net分割网络并不是只为了提高后续CPRUN模型的分割精度而忽视模型大小的增加。

表3 重构U-net分割网络的对照实验

鉴于前人的模型压缩工作中加入了L1正则化,用来帮助训练卷积神经网络的权重,同时对损失函数施加约束以便于挑选对标签特征贡献相对较大的特征通道[18]。故本文的CPRUN模型在重构的U-net分割模型进行初始训练时,就需要考虑是否为了后续通道剪枝的应用来加入正则化技术。正则化分为两种即L1正则化与L2正则化(也可称为L1范数与L2范数),L1正则化与L2正则化可以看作是加在损失函数后的惩罚项,其定义分别如式(8)和式(9)所示

(8)

(9)

但有以下两方面原因促使CPRUN模型在重构的U-net分割模型进行预训练时,并不采取正则化技术进行预训练:①CPRUN模型设计的目标中一部分的需求是尽量不给硬件增加额外的计算负担,模型的设计应该尽量精简;②CPRUN模型的实验结果显示在重构的U-net模型进行预训练时不论加入哪种正则化技术,对于分割精度的提高都没有任何增益,反而会拉低一些分割精度。

在实验阶段,针对于重构的U-net卷积神经网络,设计了应用不同正则化技术的对照实验,其结果见表4。观察表4可知,对重构的U-net卷积神经网络不采用L1或L2正则化,并不会影响其分割网络在NIH胰腺数据集的性能。所以最终CPRUN模型决定不采用任何正则化技术来对重构的U-net分割网络进行约束。

表4 正则化对重构U-net分割网络的影响

3.3 CPRUN模型的实验设计

为观察CPRUN模型在不同的阈值因子α下的性能表现,对CPRUN模型设计了应用不同α的对照实验。其结果见表5。观察表5可知:①CPRUN模型随着阈值因子α逐渐增大直至0.7,可以看出剪枝后模型大小也在逐步减小,同时Dice评价指标基本不变;②CPRUN模型在阈值因子α设定为0.7时,不但剪枝前后模型的Dice基本没有变化,并且模型大小压缩到比较小只有1.72 Mb,且压缩率ratio达到8.0%;③当阈值因子α增大至超过0.7时,虽然模型大小进一步压缩,但是,Dice评价指标由于此时的CPRUN模型卷积层特征通道数太少,对于胰腺特征的提取已经不够准确且开始出现了下降,但是下降幅度很小。

同时继续观察表5的剩余评价指标可以得出:①随着阈值因子α逐渐增大直至0.7,可以看出剪枝后模型大小也在逐步减小,同时Precision和Iou评价指标基本不变甚至略有增大;②当阈值因子α增大至超过0.7时,Precision和Iou评价指标也是呈现下降趋势,与Dice评价指标的变化基本保持一致。

表5 不同压缩程度的CPRUN模型表现

表6则详细阐述了当阈值因子α采取不同取值时,CPRUN网络的每层卷积层确切剩余的特征通道数。Ci表示图3虚线框中CPRUN模型的每个卷积层输出的特征通道数。需要指出的是表6中当α为0时表示CPRUN模型并未进行压缩(即重构的U-net卷积神经网络),此时的Ci表示卷积神经网络中每个卷积层的初始输出特征通道数。

表6 CPRUN模型的各卷积层剩余的特征通道数

CPRUN模型对于胰腺的分割效果如图6所示。

图6 分割效果

3.4 CPRUN模型与其它算法对比

表7展示了CPRUN模型在NIH数据集上与另外两种分割模型的结果比较,表7中(A6.0W0.4)表示采用量化进行模型压缩后,原始U-net分割网络的结构不变,只是在权重与激活参数的位数表示上进行约束。其中A6.0的意思是激活采用6个整数位表示,放弃小数位权重参数;W0.4代表权重采用4个浮点小数位表示,放弃整数位权重参数。

另外需要指出,本文实验阶段重现了fixed-point quantization(定点量化)中采用(A6.0W0.4)策略时的分割模型,其Dice大小为79.09%。另TernaryNet模型的权重参数大小应小于其实际模型大小[20]。

从表7我们可以得到如下结论:①从表7可以观察到当CPRUN模型的Dice和fixed-point quantization中采用(A6.0W0.4)策略时Dice大小没有明显差距时,可以明显看到我们的模型大小(841 Kb)远小于其模型大小(2.31 Mb),约为fixed-point quantization模型大小的1/3大小;②当CPRUN模型的大小(841 Kb)和TernaryNet的权重参数大小(0.66 Mb)差距不大时,可以看到CPRUN模型的Dice比其高出9个百分点左右;③当CPRUN模型的大小压缩至208 Kb时,可以看到CPRUN模型的Dice大小为78.31%,仍维持很高的分割水平,说明我们的CPRUN模型在模型大小和分割精度之间取得了新的最优平衡,真正做到了模型大小与分割精度共同兼顾。

表7 方法效果对比

4 结束语

本文提出了应用于医学图像分割中胰腺器官领域的CPRUN模型并在NIH胰腺数据集上做了验证。CPRUN模型一方面使用了重构的U-net卷积神经网络,另一方面集成了通道剪枝策略,从而实现了一个基于重构的U-net卷积神经网络与通道剪枝技术的胰腺分割端到端框架。实验结果表明,CPRUN模型可以在模型大小压缩至极小时保持较高的分割精度,真正在胰腺图像的分割模型大小与分割精度之前取得了新的最优平衡。进一步的工作将会在以下两方面进行:①未来会考虑在剪枝的基础上,融合模型压缩的量化技术对CPRUN模型进一步进行压缩;②同时也会关注于在时间压缩方面有所进步,真正的为推动实时医疗夯实基础。

猜你喜欢
剪枝胰腺重构
人到晚年宜“剪枝”
视频压缩感知采样率自适应的帧间片匹配重构
长城叙事的重构
《中华胰腺病杂志》稿约
基于YOLOv4-Tiny模型剪枝算法
高盐肥胖心肌重构防治有新策略
CT,MRI诊断急性胰腺炎胰腺内外病变价值比较
基于激活-熵的分层迭代剪枝策略的CNN模型压缩
北京的重构与再造
桑椹提取物对胰蛋白酶的抑制作用及对小鼠胰腺组织的损伤