[论文解读] GANs 'N Lungs: improving pneumonia prediction

Tatiana Malygina, Elena Ericheva+1 more|arXiv (Cornell University)|Aug 1, 2019

COVID-19 diagnosis using AI参考文献 7被引用 8

一句话总结

本文提出一种基于CycleGAN的数据增强方法，以提升在高度不平衡X光影像数据集中的肺炎预测性能。通过生成在正常与肺炎类别之间迁移的合成图像，该方法平衡了训练数据并提升了分类器性能，在验证数据上达到0.9929的ROC AUC和0.9865的PR AUC。

ABSTRACT

We propose a novel method to improve deep learning model performance on highly-imbalanced tasks. The proposed method is based on CycleGAN to achieve balanced dataset. We show that data augmentation with GAN helps to improve accuracy of pneumonia binary classification task even if the generative network was trained on the same training dataset.

研究动机与目标

• 解决医学影像数据集高度不平衡的问题，其中正常病例数量远超肺炎病例。
• 开发一种数据增强技术，生成逼真且实现类别间迁移的X光图像，以平衡训练数据。
• 在不需大量计算资源的情况下，提升深度学习模型在二分类肺炎预测中的性能。
• 实现在计算资源有限的临床机器学习流水线中，数据增强技术的实际部署。

提出的方法

• 使用带有循环一致性损失的CycleGAN训练两个生成器：G_{0→1}将正常X光片转换为类似肺炎的图像，G_{1→0}将肺炎X光片转换为类似正常的图像。
• 在原始数据集上训练生成器，以学习域到域的图像翻译，同时保留解剖结构和病理特征。
• 将原始图像与生成图像结合，构建样本数量相等的平衡训练集。
• 在增强后的数据集上，使用二元交叉熵损失对DenseNet-121分类器进行微调，以实现肺炎与正常分类。
• 应用类别激活图以可视化注意力区域，验证模型的可解释性。
• 使用CheXNet的相同验证划分，确保在有无数据增强的模型之间进行公平比较。

实验结果

研究问题

• RQ1 基于CycleGAN的数据增强是否能有效平衡高度不平衡的X光影像数据集以用于肺炎分类？
• RQ2 生成在正常与肺炎类别之间迁移的合成图像是否能提升分类器性能？
• RQ3 在增强数据上训练的肺炎分类器性能与在原始数据上训练相比有何变化？
• RQ4 该方法是否能在计算资源有限（如单张GPU）的条件下高效应用？
• RQ5 在临床合理性与模型鲁棒性方面，生成图像与真实图像相比如何？

主要发现

• 所提方法将验证集上的ROC AUC从基线的0.9745提升至0.9929，表明性能显著提升。
• PR AUC从0.9580提升至0.9865，表明在少数类（肺炎）上的性能更强。
• 使用额外的预训练数据集后，ROC AUC进一步提升至0.9939，PR AUC提升至0.9883。
• 类别激活图显示，增强后的模型能够聚焦于临床相关区域，提升了可解释性。
• 即使仅使用单张NVIDIA 1070 GPU，该方法也实现了高性能，表明计算开销较低。
• 生成图像在视觉上合理且可区分，表明其具有临床评估的潜力。