基于混合知识蒸馏的轻量级胸部疾病分类算法

针对现有胸部疾病分类算法参数量较大、对运行设备的硬件资源要求较高的问题,本文基于混合知识蒸馏(knowledge distillation,KD) 的训练策略提出一种轻量级胸部疾病分类算法RMSNet。首先,该算法将优化后的残差收缩模块加入到基础网络MobileViT中,利用软阈值化的方式过滤X光片中的背景噪声;其次,提出混合知识蒸馏训练策略,利用多层级注意力图和相似性激活矩阵作为监督信号,提升轻量级模型的分类性能;最后,使用焦点损失函数(focal loss,FL) 缓解数据集正负样本不均衡的问题。在ChestX-ray14数据集上展开验证,蒸馏训练后的RMSNet学生模型识别14类胸部疾病的平均AUC值为0.836,而参数量和浮点计算量分别为0.96 M和0.27 G。实验结果表明,本文算法在保持轻量化的同时分类精度更高,能有效降低算法运行时的硬件要求。