CardioCaps：用于类别不平衡超声心动图分类的基于注意力的胶囊网络
计算机视觉与模式识别
2024-05-21 v2
摘要
胶囊神经网络 (CapsNets) 是一种利用由多个神经元形成的向量表示的新颖架构。具体而言，动态路由胶囊网络 (DR-CapsNets) 采用仿射矩阵和动态路由机制来训练胶囊并获取平移等变性，与传统的卷积神经网络 (CNNs) 相比增强了其鲁棒性。超声心动图捕捉心脏的运动图像，对传统的图像分类方法提出了独特的挑战。在本文中，我们探索了 DR-CapsNets 的潜力，并提出了 CardioCaps，一种用于类别不平衡超声心动图分类的基于注意力的 DR-CapsNet 新颖架构。CardioCaps 包含两个关键组件：结合回归辅助损失的加权边缘损失和注意力机制。首先，加权边缘损失优先考虑阳性案例，并由基于射血分数 (EF) 回归任务（心脏功能的关键指标）的辅助损失函数补充。这种方法增强了模型在面对类别不平衡时的韧性。其次，认识到动态路由的二次复杂性导致训练效率低下，我们采用注意力机制作为计算上更高效的替代方案。我们的结果表明，CardioCaps 超越了传统的机器学习基线方法，包括带有采样方法和类别权重矩阵的逻辑回归、随机森林和 XGBoost。此外，CardioCaps 优于其他深度学习基线方法，如 CNNs、ResNets、U-Nets 和 ViTs，以及先进的 CapsNets 方法，如 EM-CapsNets 和 Efficient-CapsNets。值得注意的是，我们的模型展示了对类别不平衡的鲁棒性，即使在阴性案例占很大比例的数据集中也能实现高精度。
引用
@article{arxiv.2403.09108,
title = {CardioCaps: Attention-based Capsule Network for Class-Imbalanced Echocardiogram Classification},
author = {Hyunkyung Han and Jihyeon Seong and Jaesik Choi},
journal= {arXiv preprint arXiv:2403.09108},
year = {2024}
}
备注
8 pages, 6 figures