[论文解读] 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

一句话总结

该论文提出了一种分层渐近训练（HA-GAN）方法，这是一种内存高效的3D GAN框架，通过在共享潜在空间中训练分层生成器和判别器，实现高分辨率体素数据的生成。通过采用子体积选择和渐近推理，HA-GAN将内存使用量最多降低80%，同时在COPDGene和GSP数据集上实现了最先进（SOTA）的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$ 分支分别生成中间分辨率和高分辨率输出，采用转置3D卷积和跳跃连接，其中 $G^H$ 通过两个上采样阶段生成 $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医学图像，从而通过数据增强提升下游监督学习的性能？
• 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数据集上均实现了最先进（SOTA）的FID分数，FID被用作训练期间的主要验证指标。
• 条件HA-GAN变体成功通过生成多样化且逼真的3D图像（基于类别标签条件）提升了监督3D分类的性能。
• 生成器网络 $G^A$ 有效学习了分层潜在空间，通过渐进式上采样和残差块实现高分辨率合成。
• 判别器 $D^H$ 处理 $32 \times 256^2$ 的子体积，显著降低了内存使用量，同时保持了良好的判别性能。
• 用于监督微调的3D CNN在5分类任务上表现出稳健性能，其网络架构设计旨在保留空间层次结构和特征表示。