1. 引言
医学影像技术作为现代医学诊断与临床决策的重要支撑,在疾病筛查、病灶定位以及治疗方案制定等方面发挥着不可替代的作用。随着医学成像设备分辨率的不断提升以及影像数据规模的持续增长,如何从复杂的医学影像中快速、准确地提取关键结构信息,已成为医学图像分析领域亟需解决的核心问题之一。医学影像图像分割通过对感兴趣区域进行精确划分,为病灶检测、器官测量和病理分析提供了可靠依据,因此在临床辅助诊断中具有重要研究价值和应用前景。
近年来,随着深度学习技术的快速发展,基于卷积神经网络的医学影像分割方法取得了显著进展。其中,U型网络由于其编码器–解码器对称结构以及跳跃连接机制[1],能够在有效提取高层语义信息的同时保留图像的细节特征,在医学影像分割任务中表现出较高的分割精度和良好的泛化能力。基于U型网络的分割模型已被广泛应用于脑部影像、肺部影像以及肿瘤分割等多个医学场景[2] [3],并逐渐成为医学影像分割领域的主流方法。然而,此类模型通常具有网络结构复杂、参数规模较大以及计算量高等特点,对计算平台的存储资源和浮点运算能力提出了较高要求[4]。
随着智慧医疗和数字医疗的不断推进,医学影像分析系统逐步从传统的中心化计算模式向更加灵活、高效的边缘计算模式转变[5]。在临床应用场景中,实时性和稳定性往往具有较高要求,而依赖云端服务器进行集中处理不仅会引入额边的通信时延,还可能受到网络带宽和数据隐私等因素的制约。边缘计算通过将数据处理与模型推理任务下沉至靠近数据源的边缘设备,有效缩短了数据传输路径,降低了系统整体时延,同时提升了数据安全性和系统响应能力,为医学影像智能分析提供了新的技术支撑。
边缘计算在医学影像应用中具有明显优势,但受限于边缘设备算力和存储资源的限制[6],现有医学影像分割模型难以直接部署和高效运行。因此,如何在保证分割精度的前提下,对模型进行结构优化和推理加速,成为医学影像分割算法工程化落地过程中面临的重要挑战。本文基于此以瑞芯微RK3588平台作为边缘计算硬件基础[7],对U型网络医学影像分割模型进行轻量化设计和量化部署,实现医学影像分割模型在RK3588平台上的高效部署与运行。
2. U型网络的轻量化设计和量化部署方法
2.1. U型网络的轻量化设计——MALUNet网络
MALUNet [8]是基于U型网络架构进行轻量化设计后得到的医学影像分割模型。考虑到实际应用中计算资源的限制,尤其是对于移动医疗设备等低功耗平台,MALUnet在保持UNet核心架构的基础上通过集合四个轻量化的注意力机制模块,实现了模型的高效性和高性能之间的平衡。该模型在设计上注重减少参数量与计算复杂度,同时最大化保留分割精度,尤其是在处理复杂的医学影像数据时。接下来的部分将详细介绍这四个模块的设计理念与实现方法,分别是:空洞门控注意力块(Dilated Gated Attention Block, DGA)、反转外部注意力块(Inverted External Attention Block, IEA)、通道注意力桥接块(Channel Attention Bridge Block, CAB)和空间注意力桥接块(Spatial Attention Bridge Block, SAB)。实验证明MALUNet能够在尽可能少的计算资源下,达到较为卓越的医学影像分割效果,为低功耗平台上的医学影像智能处理提供了可行的解决方案。MALUNet架构图如下图1所示:
Figure 1. The illustration of MALUNet architecture
图1. MALUNet架构示意图
医学影像分割任务属于密集预测问题,因此模型需要具备对整体结构和局部细节信息的理解能力。全局信息有助于模型从整体层面把握目标区域与背景之间的空间关系,而局部信息有助于提高对边界细节和形态特征分割的准确率。
基于此本文引入DGA模块[8],通过融合空洞卷积与门控注意力机制,实现对全局信息与局部信息的特征提取。如图2(a),本模块由两个单元组成,分离空洞卷积单元首先利用不同膨胀率的卷积操作对特征图进行处理,使网络能够同时捕获大尺度结构信息和小尺度细节特征。而门控注意力单元对提取到的特征进行自适应调制,使模型能够抑制冗余信息并强化对目标区域的关注。IEA是在外部注意力机制的基础上引入倒置结构设计的一种改进模块[8] [9],如图2(b),能够实现在保持全局建模能力并进一步降低计算复杂度,增强特征表达的稳定性。如图2(c),CAB用于在通道维度上实现多阶段特征的信息交互与有效融合。CAB将各阶段特征首先通过全局信息提取操作进行压缩[8],然后不同阶段的通道特征在通道维度上进行融合,并通过轻量化的映射生成对应的通道注意力权重。最后使得模型能够识别哪个阶段的拓展通道是重要的。而SAB用于在特征融合过程中学习空间位置之间的依赖关系[8],如图2(d),以增强关键区域的表达能力。
Figure 2. The illustration of the attention module architecture
图2. 注意力模块架构示意图
2.2. 轻量化部署方法
随着深度神经网络的迅速发展,其模型规模与计算复杂度也随之显著增长。但在边缘设备等资源受限平台上,深度模型实现端侧部署面临着存储开销、推理延迟与能耗问题等问题。基于此,学界逐步形成了以剪枝、蒸馏和量化压缩为代表的三类轻量化部署方法体系。
模型剪枝是最早被系统研究的轻量化手段之一,其基本思想是通过消除网络中的冗余参数或冗余结构,降低模型复杂度。LeCun等人最早提出基于二阶泰勒展开的Optimal Brain Damage方法[10],通过近似估计单个权重对损失函数的影响程度,从而实现精细化的权重剪枝。Hassibi与Stork在此基础上提出的Optimal Brain Surgeon方法[11]考虑到权重之间的相互依赖关系,提高了剪枝精度。Han等人提出了基于权值幅值的剪枝方法[12],该方法通过假设小幅值权重对模型输出贡献有限,在训练完成后直接移除低于阈值的权重,并辅以微调过程以恢复性能。近年来,研究重点逐步从非结构化剪枝转向以通道、卷积核或网络层作为剪枝对象的结构化剪枝,可以更好地适配实际硬件架构,本文即将提到的一次性层剪枝就是一种结构化剪枝。
知识蒸馏通过模型间的知识迁移实现轻量化部署。该思想最早由Hinton等人提出[13],其核心在于利用性能较强的教师模型输出的软标签,引导学生模型学习更丰富的类别关系信息,从而在较小模型规模下获得接近教师模型的性能。Romero等人提出的特征蒸馏方法通过约束中间特征表示实现更深层次的知识传递[14]。
量化压缩则从数值表示的角度降低模型的存储与计算成本,其核心目标是使用低精度数据表示网络参数与中间激活。早期量化研究主要集中在权重量化领域,其核心思想是将原本以浮点数表示的网络权重映射为低比特整数表示。随着深度学习模型规模的不断扩大,Courbariaux等人提出了二值化与低比特网络[15],表明在一定条件下低精度表示仍可维持可接受的模型性能。Jacob等人提出了后训练量化方法[16],通过在模型训练完成后对权重和激活进行线性映射,实现无需重新训练的快速部署。
2.2.1. 一次性层剪枝
一次性层剪枝的核心思想在于度量每一网络层对整体模型性能的相对重要性[17],并在满足约束条件的前提下移除冗余层。设原始神经网络由L个可剪枝层组成,记第l层的参数集合为
,网络整体参数为
,训练目标函数可表示为:
(1)
其中
表示网络映射,
为损失函数。
一次性层剪枝的关键在于构建层级重要性评分函数
,用于衡量移除第l层对损失函数的影响程度。常见做法是基于敏感性分析或近似误差评估,将剪枝问题转化为对损失变化的估计。在忽略高阶项的情况下,移除第l层可视为对参数
施加零化操作,其引起的损失变化可近似表示为:
(2)
反映了当前参数状态下第l层对目标函数的瞬时贡献程度,可作为层重要性的近似度量。基于所有层的重要性评分,一次性层剪枝通常在全局约束条件下完成剪枝决策。在给定参数预算C的条件下,剪枝问题将转化为如下优化目标:
(3)
其中
为二值掩码变量,
表示保留第l层,
表示剪除该层,
表示对应层的计算或参数开销。通过求解该优化问题,可以在单次决策中确定最终的网络结构。
2.2.2. 归一化蒸馏
归一化蒸馏的核心原理是通过对特征蒸馏损失进行归一化处理[17],使学生模型能够更加充分地吸收教师模型在不同层次上的表征信息,适用于剪枝后结构发生变化的网络重训练阶段。
设教师模型与学生模型在第
个阶段输出的特征表示分别为
与
,传统特征蒸馏通常通过最小化二者之间的欧氏距离来实现知识约束,其损失函数可表示为:
(4)
上式表示不同阶段特征的数值尺度决定了其在总损失中的相对贡献,当某些特征幅值显著大于其他阶段时,其损失项将主导梯度更新过程,从而削弱整体蒸馏效果。因此,归一化蒸馏通过计算教师模型在各阶段特征的范数信息构造归一化系数,对特征蒸馏损失进行重加权处理。归一化特征蒸馏损失可形式化表示为:
(5)
其中
表示第
个阶段对应的归一化权重。
在整体训练目标中,归一化蒸馏通常与任务损失及输出蒸馏损失共同作用,从而在性能保持与结构压缩之间实现平衡。通过这种方式,学生模型在学习任务目标的同时,能够在不同层级上稳定地对齐教师模型的中间表示,为剪枝或结构简化后的模型提供更加可靠的性能恢复路径。
2.2.3. 量化压缩
量化压缩的核心在于构建连续浮点空间到离散整数空间之间的映射关系。设某一层的浮点权重为
,其量化表示为
,常见的线性量化过程可表示为[18]:
(6)
其中
表示量化尺度因子,用于控制浮点值到整数值的映射比例,
表示零点偏移,用于保证零值在量化空间中的精确表示。
在推理阶段,量化模型以整数形式执行乘加运算,由于整数运算在RK3588平台有更高的吞吐效率,量化压缩能够减少模型体积并降低推理延迟。在量化感知训练场景下,为保证梯度的可传播性,量化函数通常采用直通估计方式对不可导算子进行近似,使模型参数能够在低精度约束下稳定收敛。
3. RK3588开发板编译流程
Figure 3. Physical view of the RK3588
图3. RK3588实物图
随着深度学习技术对边缘计算需求的日益增长,国产芯片产业持续发展,打破了长期的技术壁垒。瑞芯微电子(Rockchip)推出的新一代旗舰级芯片RK3588展现出了卓越的性能表现。RK3588内置了针对深度学习加速的三核NPU,能够提供6TOPS的混合精度算力支持,具备强大的通用计算能力。因此,RK3588很适合应用于对实时性与精度要求严苛的医学影像分割任务,为深度学习算法在端侧部署场景提供了硬件基础。RK3588开发板实物图见图3。
3.1. ONNX模型
ONNX模型是一种针对机器学习模型设计的开放式文件格式标准,它定义了一套通用的算子集与数据流图规范。本项目中ONNX模型为连接训练环境与端侧部署环境的中间介质,旨在将基于特定框架训练得到的医学影像分割模型转化为一种硬件无关的标准表达形式。
3.2. RKNN模型
RKNN模型作为瑞芯微平台专属的深度学习推理格式,是实现算法在RK3588芯片NPU上高效运行的核心载体。利用配套的RKNN-Toolkit2工具链,我们将上一阶段的ONNX模型进一步编译为RKNN文件,这一转换过程深入执行了图融合、内存复用以及算子重排等硬件级优化策略。基于此,模型能够充分释放NPU的异构计算潜能,为最终的端侧高效推理提供即插即用的执行文件。
3.3. 模型量化部署
针对RK3588开发板上的轻量化U型网络的部署和推理过程,本文按照如下流程对其进行量化部署(见图4):首先对原始U型网络进行轻量化设计,得到MALUNet模型,并进行归一化蒸馏,一次性层剪枝获得轻量化网络;随后利用RKNN-Toolkit2生成RKNN模型文件;最后将生成的RKNN模型部署至RK3588开发板,实时输出分割结果,从而在保证分割精度的前提下实现了算法在嵌入式平台上的高性能运行。
Figure 4. Deployment workflow of the image segmentation model
图4. 图像分割模型部署流程图
4. 实验设置
4.1. 数据集
本实验选用ISIC2017 (International Skin Imaging Collaboration 2017)皮肤病变数据集作为模型训练与评估的数据集。ISIC2017数据集共包含2150张皮肤镜图像,每张图像均配有对应的像素级分割掩码标签,可用于监督式图像分割任务。图像内容涵盖多种皮肤病变类型,具有病变形态多样、边界复杂等特点,对模型的分割能力提出了较高要求。实验中,按照7:3的比例对数据集进行划分,其中训练集包含1500张图像,测试集包含650张图像。
4.2. 实验细节
MALUNet网络的训练在单张NVIDIARTX A4000 GPU上完成。所有图像进行归一化并调整尺寸至256 × 256,同时采用包括随机翻转、颜色变换和添加噪声等数据增强策略。损失函数采用BCEDiceLoss,其表达式如公式(7)所示,学习率调度器采用CosineAnnealingLR,优化器使用AdamW,初始学习率为1e−3,最大迭代次数为300,最低学习率为1e−5。训练周期设置为300,批次大小为16。
(7)
其中,
是一个超参数,用于调节
和
之间的权重。
对于MALUNet教师模型的归一化蒸馏,所有学生模型的训练均在单张NVIDIA GeForce RTX4060 GPU完成。损失函数采用复合蒸馏损失,其表达式如公式(8)所示,学习率调度器采用CosineAnnealingLR,优化器使用AdamW,初始学习率为1e−4,最大迭代次数为50,最低学习率为1e−6。蒸馏训练周期设置为50,批次大小为8,温度参数
,输出蒸馏权重
,特征蒸馏权重
。
(8)
其中,
为原始分割任务的损失函数,表达式如公式(7)所示;
为输出蒸馏损失,即通过KL散度计算教师模型和学生模型输出之间的差异;
为归一化特征蒸馏损失,表达式如公式(5)所示。
MALUNet的八种学生模型的单样本推理全部在Rockchip RK3588 SoC的NPU上完成,采用与训练相同图像预处理和后处理方法,采用ISIC2017训练集的前150个样本构成量化校准数据集。
4.3. 学生模型的设计
MALUNet的通道数配置为[8, 16, 32, 64, 128],总计参数量为177,943,本文以MALUNet为基础,设计了八种不同轻量化程度的学生模型,具体设计的八种学生模型变体及其详细配置详见下表1。
Table 1. System resulting data of standard experiment
表1. 标准试验系统结果数据
模型名称 |
配置细节 |
参数量 |
参数压缩比 |
MALUNetSmall |
各层通道数设置为[4, 8, 12, 16, 24, 32] |
36,765 |
4.84× |
MALUNetMedium |
各层通道数配置为[8, 16, 20, 28, 40, 56] |
132,991 |
1.34× |
MALUNetShallow |
将编码器与解码器构建为五层结构 |
109,991 |
1.62× |
MALUNetNarrow |
保持深度但减少每层复杂度 |
63,133 |
2.82× |
MALUNetEfficient |
全面采用深度可分离卷积以降低计算成本 |
18,803 |
9.46× |
MALUNetAttLight |
全面简化注意力机制模块以降低计算成本 |
126,092 |
1.41× |
MALUNetGlobalAtt |
使用全局注意力机制代替原有的多尺度注意力维度 |
126,718 |
1.40× |
MALUNetMobile |
结合MobileNetV3的逆残差结构与注意力机制进行重构 |
118,257 |
1.50× |
上述八种变体分别代表了不同的轻量化优化方向,构成了完整的实验对照组。其中,MALUNetSmall 与MALUNetMedium主要通过调整通道数量来直接压缩参数量,旨在考察模型宽度对特征提取能力的影响;MALUNetShallow通过减少编码器与解码器的层数来降低计算深度;MALUNetNarrow标准卷积替换复杂注意力模块、使用BatchNorm替代GroupNorm,全面降低模型复杂度;MALUNetEfficient全面采用深度可分离卷积以降低计算成;MALUNetMobile引入了MobileNetV3的逆残差结构;MALUNetAttLight与MALUNetGlobalAtt分别尝试了简化注意力桥接结构与引入全局注意力机制。本文通过后续实验对比这八种模型在端侧的实际表现,将筛选出最适合该开发板部署环境的模型架构。
4.4. 评价指标
本文采用五类指标衡量分割与实时推理性能:平均交并比(mIoU)、Dice系数、单样本推理时间(
)、单次推理浮点运算量(FLOPs)。
(9)
(10)
其中,
分别表示第
类的真阳性、假阳性和假阴性,N为类别总数。
5. 结果分析
Table 2. Distillation training losses comparison of student models on the training set
表2. 学生模型在训练集上的蒸馏训练损失对比
模型名称 |
任务损失 |
输出蒸馏损失 |
归一化特征蒸馏损失 |
总损失 |
MALUNetSmall |
0.3496 |
0.0007 |
0.5460 |
0.6229 |
MALUNetMedium |
0.3348 |
0.0004 |
0.4287 |
0.5493 |
MALUNetShallow |
0.3455 |
0.0004 |
0.4530 |
0.5722 |
MALUNetNarrow |
0.3395 |
0.0005 |
0.5330 |
0.5632 |
MALUNetEfficient |
0.3605 |
0.0008 |
0.7518 |
0.7368 |
MALUNetAttLight |
0.3403 |
0.0004 |
0.4919 |
0.5865 |
MALUNetGlobalAtt |
0.3399 |
0.0004 |
0.3973 |
0.5388 |
MALUNetMobile |
0.4006 |
0.0020 |
1.3487 |
1.0760 |
Table 3. Distillation training losses comparison of student models on the validation set
表3. 学生模型在验证集上的蒸馏训练损失对比
模型名称 |
任务损失 |
输出蒸馏损失 |
归一化特征蒸馏损失 |
总损失 |
MALUNetSmall |
0.3827 |
0.0008 |
0.5172 |
0.6417 |
MALUNetMedium |
0.3779 |
0.0004 |
0.4351 |
0.5959 |
MALUNetShallow |
0.3827 |
0.0004 |
0.4335 |
0.5996 |
MALUNetNarrow |
0.3915 |
0.0005 |
0.4535 |
0.6014 |
MALUNetEfficient |
0.3911 |
0.0008 |
0.7364 |
0.7598 |
MALUNetAttLight |
0.3814 |
0.0004 |
0.4590 |
0.6111 |
MALUNetGlobalAtt |
0.3815 |
0.0004 |
0.3775 |
0.5705 |
MALUNetMobile |
0.4287 |
0.0018 |
1.2255 |
1.0423 |
表2、表3显示,MALUNetGlobalAtt均获得了在ISIC2017数据集上的归一化蒸馏训练获得了最低的总损失值,接下来在我们将上述八种模型量化部署至RK3588开发板,并进行单样本推理,并与原始MALUNet模型在PC端的测试结果做对比。
经测试,原始MALUNet模型在单张NVIDIA GeForce RTX4060 GPU上的mIoU为0.8847,Dice系数为0.9194,单样本推理时间为0.2503 s,FLOPs为355886。
选取索引为0的样本图片进行单样本推理,结果见表4,同时选取索引为0,52,147的样本图片进行单样本推理,结果见图5。在同一张图片上,MALUNetGlobalAtt模型取得了更高的mIoU和Dice系数,具有更好的模型性能,并且推理时间大幅度降低。
Table 4. Comparison of experimental results of student models
表4. 学生模型实验结果对比表
模型名称 |
mIoU |
Dice系数 |
(s) |
推理时间降低率 |
FLOPs |
MALUNetSmall |
0.8122 |
0.8963 |
0.1854 |
25.9% |
73,530 |
MALUNetMedium |
0.6114 |
0.7005 |
0.0588 |
76.5% |
265,982 |
MALUNetShallow |
0.7154 |
0.8341 |
0.0131 |
94.8% |
428,978 |
MALUNetNarrow |
0.6559 |
0.7922 |
0.0100 |
96.0% |
126,266 |
MALUNetEfficient |
0.7594 |
0.8632 |
0.0140 |
94.4% |
37,606 |
MALUNetAttLight |
0.7640 |
0.8662 |
0.0070 |
97.2% |
252,184 |
MALUNetGlobalAtt |
0.8126 |
0.8966 |
0.0100 |
96.0% |
253,436 |
MALUNetMobile |
0.8012 |
0.8896 |
0.0201 |
92.0% |
236,514 |
Figure 5. Comparison of experimental results of student models
图5. 学生模型实验结果对比图
6. 结论
本文针对医学影像分割模型在边缘设备上难以兼顾精度与实时性的问题,设计了轻量化网络MALUNet,并基于国产瑞芯微RK3588平台完成了全流程量化部署。研究通过引入空洞门控、反转外部注意力及特征桥接模块优化网络结构,并结合归一化知识蒸馏与一次性层剪枝策略,有效解决了复杂模型在受限算力下的运行瓶颈。在ISIC2017数据集上的实验结果表明,本文提出的轻量化MALUNetGlobalAtt综合性能最优,在保持较高分割精度(mIoU 0.8126)的同时,相比原始模型推理时间降低了96%,成功实现了在RK3588 NPU上的毫秒级高效运行。本文提出的软硬协同优化方案验证了深度医学影像算法在国产边缘计算平台落地的工程可行性,为智慧医疗的端侧应用提供了具有实用价值的技术参考。
致 谢
在论文完成后,首先,我们非常想要感谢指导老师的辛勤指导。在我们撰写论文的过程中,从一开始论文的选题,再到论文的构思与撰写方式,我们得到了指导老师悉心细致的教诲和无私的帮助,他那广博的知识、深厚的学术素养、严谨的治学精神使我们这一生都会终身得到受益。在此,我们向指导老师表示真挚的感谢。
在论文的写作过程中,我们也得到了许多同学的宝贵建议,也在网上看到了诠释了开源精神的先辈们,若没有他们将辛苦搜集整理出来的数据开源,我们将寸步难行,在此,对他们表示真诚的谢意。
最后的最后,感谢所有关心、支持、帮助过我们的良师益友。