[論文レビュー] Hierarchical Amortized Training for Memory-efficient High Resolution 3D GAN

Li Sun, Junxiang Chen+4 more|arXiv (Cornell University)|Aug 5, 2020

Generative Adversarial Networks and Image Synthesis参考文献 59被引用数 4

ひとこと要約

本稿では、共有された潜在空間を用いて階層的な生成器と識別器を訓練することで、高解像度ボリュームデータを生成するメモリ効率の良い3D GANフレームワーク、階層的アモアタイズドトレーニング（HA-GAN）を提案する。部分体積選択とアモアタイズド推論を用いることで、HA-GANは最大80％のメモリ使用量削減を達成しながら、COPDGeneおよびGSPデータセットで最先端のFIDスコアを達成し、計算コストを低減しつつ高精細な3D医療画像合成を可能にする。

ABSTRACT

Generative Adversarial Networks (GAN) have many potential medical imaging applications, including data augmentation, domain adaptation, and model explanation. Due to the limited memory of Graphical Processing Units (GPUs), most current 3D GAN models are trained on low-resolution medical images, these models either cannot scale to high-resolution or are prone to patchy artifacts. In this work, we propose a novel end-to-end GAN architecture that can generate high-resolution 3D images. We achieve this goal by using different configurations between training and inference. During training, we adopt a hierarchical structure that simultaneously generates a low-resolution version of the image and a randomly selected sub-volume of the high-resolution image. The hierarchical design has two advantages: First, the memory demand for training on high-resolution images is amortized among sub-volumes. Furthermore, anchoring the high-resolution sub-volumes to a single low-resolution image ensures anatomical consistency between sub-volumes. During inference, our model can directly generate full high-resolution images. We also incorporate an encoder with a similar hierarchical structure into the model to extract features from the images. Experiments on 3D thorax CT and brain MRI demonstrate that our approach outperforms state of the art in image generation. We also demonstrate clinical applications of the proposed model in data augmentation and clinical-relevant feature extraction.

研究の動機と目的

• 高解像度3D GANのトレーニングに伴う高いメモリ要件を解決すること。
• COPDGeneおよびGSPのような大規模な3D医療データセットにおける3D GANの効率的トレーニングを可能にすること。
• 画像品質を損なわずにメモリ消費量を低減する階層的かつアモアタイズドなトレーニングフレームワークの開発。
• HA-GANの有効性を、教師あり3D医療画像分類のデータ拡張における応用を通じて示すこと。

提案手法

• 生成器ネットワーク $G^A$ はまず、ノイズベクトル $Z$ を $64^3$ の低解像度特徴マップにマップし、その後、GroupNormおよびReLU活性化関数を用いた複数の残差ブロックを通じてアップサンプリングされる。
• $G^L$ および $G^H$ ブランチは、逆3次元畳み込みおよびスキップ接続を用いて中間解像度および高解像度出力を生成し、$G^H$ は2段階のアップサンプリングステージを経て $256^3$ の出力を生成する。
• 識別器 $D^L$ および $D^H$ はスペクトル正規化とLeakyReLUを用い、$D^H$ はメモリ使用量を削減するために $32 \times 256^2$ の部分体積を処理する。
• エンコーダ $E^H$ および $E^G$ は、実画像から潜在コードを抽出し、アモアタイズド推論とクラスラベルを用いた条件付き生成を可能にする。
• 条件付きHA-GANバージョンはワンホットクラスコード $c$ を組み込み、マルチクラス生成のための補助分類器として識別器を用いる。
• トレーニング中に部分体積選択が適用され、一度に $32 \times 128^2$ パッチしか処理されないため、メモリ使用量が最大80％削減される。

実験結果

リサーチクエスチョン

• RQ1 階層的アモアタイズドトレーニングは、画像品質を損なわずに高解像度3D GANにおけるメモリ消費量を削減できるか？
• RQ2 部分体積選択は、3D GANトレーニング中のGPUメモリ使用量削減にどの程度効果的か？
• RQ3 HA-GANは、データ拡張を通じて教師あり3D医療画像分類の性能向上に寄与する高精細な3D医療画像を生成できるか？
• RQ4 $G^A$、$G^L$、$G^H$ といったアーキテクチャ的要素が、生成された3Dボリュームの品質および解像度に与える影響は何か？
• RQ5 補助分類器を備えた条件付きHA-GANは、マルチクラス3D医療画像タスクにおける性能向上にどのように寄与するか？

主な発見

• 部分体積選択により、HA-GANは最大80％のメモリ使用量削減を達成し、乗数係数を1/8に設定した場合、メモリ使用量は13,185 MBから5,961 MBに低下する。
• COPDGeneおよびGSPデータセットの両方で、最先端のFIDスコアを達成しており、トレーニング中の主な検証指標としてFIDが用いられている。
• 条件付きHA-GANバージョンは、クラスラベルに条件づけられた多様で現実的な3D画像を生成することで、教師あり3D分類の性能を向上させた。
• 生成器ネットワーク $G^A$ は階層的潜在空間を効果的に学習し、段階的なアップサンプリングと残差ブロックを通じて高解像度合成を可能にする。
• 識別器 $D^H$ は $32 \times 256^2$ の部分体積を処理し、メモリ使用量を顕著に削減しながらも、識別性能を維持している。
• 教師あり微調整に用いられた3D CNNは5クラス分類タスクで頑健な性能を示しており、ネットワークアーキテクチャは空間階層性と特徴表現を保持するように設計されている。