[論文レビュー] 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に基づくデータ拡張手法を提案する。正常と肺炎のクラス間を移行する合成画像を生成することで、トレーニングデータのバランスをとるとともに、分類器の性能を向上させ、検証データにおいてAUC 0.9929（ROC）、AUC 0.9865（PR）を達成した。

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とサイクル整合性損失を用い、正常X線を肺炎風に変換するG_{0→1}と、肺炎X線を正常風に変換するG_{1→0}の2つの生成器を訓練する。
• 生成器を元のデータセット上で訓練し、解剖学的構造と病変特徴を保持したままドメイン間変換を学習する。
• 元の画像と生成画像を組み合わせ、各クラスに等しいサンプル数がなるようにバランスの取れたトレーニングセットを構築する。
• バイナリクロスエントロピー損失を用いて、拡張されたデータセット上でDenseNet-121分類器を微調整する。
• 注意領域の可視化のため、クラスアクティベーションマップを適用し、モデルの解釈可能性を検証する。
• CheXNetと同一の検証スプリットを用いることで、拡張あり・なしのモデル間の公平な比較を実現する。

実験結果

リサーチクエスチョン

• RQ1 CycleGANに基づくデータ拡張は、肺炎分類のための極めて不均衡なX線データセットを効果的にバランス化できるか？
• RQ2 正常と肺炎のクラス間を移行する合成画像を生成することで、分類器の性能が向上するか？
• RQ3 元のデータセットと比較して、拡張データセットで学習した肺炎分類器の性能はどのように変化するか？
• RQ4 1枚のGPUといった限られた計算リソースでも、この手法は効率的に適用可能か？
• RQ5 生成画像は臨床的に妥当で、モデルの頑健性を損なわない程度に現実的か？

主な発見

• 提案手法により、検証セットにおけるAUC（ROC）がベースラインの0.9745から0.9929に上昇し、顕著な性能向上が確認された。
• AUC（PR）は0.9580から0.9865に改善され、少数クラス（肺炎）における性能強化が示された。
• 追加の事前学習データセットを用いることで、AUC（ROC）は0.9939、AUC（PR）は0.9883にさらなる向上が得られた。
• クラスアクティベーションマップにより、拡張モデルが臨床的に関連する領域に注目していることが確認され、解釈可能性が向上した。
• 単一のNVIDIA 1070 GPUでも高い性能が達成されたことから、計算コストが低く抑えられることを示した。
• 生成画像は視覚的に妥当で、区別可能であり、臨床的評価への応用可能性が示唆された。