[论文解读] A Comparative Study of CNN, ResNet, and Vision Transformers for Multi-Classification of Chest Diseases

Ananya Jain, Aviral Bhardwaj+2 more|arXiv (Cornell University)|May 31, 2024

AI in cancer detection被引用 4

一句话总结

本研究比较了卷积神经网络（CNN）、ResNet 和视觉Transformer（ViT）模型在NIH胸部X光数据集上对14种胸部疾病进行多分类的表现。在ImageNet-21k上微调后，预训练的ViT模型表现优于CNN和ResNet，测试准确率达到94%，表明在大规模数据预训练下，基于注意力机制的特征学习在医学影像中具有显著优势。

ABSTRACT

Large language models, notably utilizing Transformer architectures, have emerged as powerful tools due to their scalability and ability to process large amounts of data. Dosovitskiy et al. expanded this architecture to introduce Vision Transformers (ViT), extending its applicability to image processing tasks. Motivated by this advancement, we fine-tuned two variants of ViT models, one pre-trained on ImageNet and another trained from scratch, using the NIH Chest X-ray dataset containing over 100,000 frontal-view X-ray images. Our study evaluates the performance of these models in the multi-label classification of 14 distinct diseases, while using Convolutional Neural Networks (CNNs) and ResNet architectures as baseline models for comparison. Through rigorous assessment based on accuracy metrics, we identify that the pre-trained ViT model surpasses CNNs and ResNet in this multilabel classification task, highlighting its potential for accurate diagnosis of various lung conditions from chest X-ray images.

研究动机与目标

• 评估并比较卷积神经网络（CNN）、ResNet与视觉Transformer（ViT）架构在X光图像中对14种胸部疾病进行多标签分类的性能。
• 评估在ImageNet-21k等大规模数据集上预训练对ViT在医学影像中性能的影响。
• 研究ViT与传统基于CNN的模型在低资源医学影像场景下的泛化能力与特征学习效率。
• 分析ViT中的注意力图及其可解释性，以评估其临床相关性与决策透明度。
• 识别多类别胸部X光图像分类中数据不平衡与模型优化缺陷等局限性。

提出的方法

• 微调了两种ViT变体：一种在ImageNet上预训练，另一种使用NIH胸部X光数据集（112,120张前后位X光图像）从头训练。
• 使用相同的数据集与训练协议，训练并评估了标准CNN和ResNet-50作为基线模型。
• 采用迁移学习方法，使用预训练的InceptionResNetV2模型，通过提取ViT编码器最后一层多头注意力层的注意力权重，生成注意力图。
• 以224×224像素的RGB图像作为输入，应用数据增强技术，并使用带Sigmoid激活函数的交叉熵损失函数进行多标签分类训练。
• 通过在训练集与验证集上评估准确率、AUC-ROC曲线及损失曲线，比较各模型的性能与泛化能力。
• 通过在头与补丁之间平均注意力权重，再上采样并叠加到原始图像上，可视化注意力图以解释模型预测结果。

实验结果

研究问题

• RQ1 预训练的视觉Transformer是否在胸部X光图像的多标签分类任务中优于CNN与ResNet？
• RQ2 在ImageNet-21k等大规模数据集上进行预训练，如何影响ViT在医学图像分类中的性能？
• RQ3 注意力机制在识别胸部X光图像中临床相关区域方面，相较于卷积特征图的作用如何？
• RQ4 数据不平衡与罕见疾病样本不足如何影响不同架构模型的性能？
• RQ5 ViT模型能否通过注意力可视化实现与传统CNN相当或更优的可解释性？

主要发现

• 在NIH胸部X光数据集上微调后，预训练的视觉Transformer（ViT-v2）测试准确率达到94%，显著优于CNN与ResNet模型。
• 在ImageNet-21k上预训练的ViT模型表现优于从头训练的ViT及所有基线模型，凸显大规模预训练的重要性。
• AUC-ROC分析显示，ViT模型在“无异常发现”类别上表现较差（AUC = 0.50），表明其难以区分正常与异常病例，而在肺不张（AUC = 0.51）和胸膜增厚（AUC = 0.53）等疾病上表现中等。
• ViT模型在罕见疾病如心脏扩大（AUC = 0.48）和疝气（AUC = 0.43）上表现不佳，主要由于数据集中存在数据不平衡问题。
• 从ViT模型生成的注意力图突出了X光图像中临床相关的区域，如心脏与肺野，证实了模型能够聚焦于具有诊断意义的区域。
• 尽管平均性能优异，ViT模型在“无异常发现”类别上表现欠佳，提示未来工作需改进数据采样或类别平衡技术。