[논문 리뷰] Hierarchical Amortized Training for Memory-efficient High Resolution 3D GAN

Li Sun, Junxiang Chen외 4명|arXiv (Cornell University)|2020. 08. 05.

Generative Adversarial Networks and Image Synthesis참고 문헌 59인용 수 4

한 줄 요약

이 논문은 고해상도 부피 데이터를 생성하기 위해 공유된 잠재 공간을 가진 계층적 생성자와 판별자로 훈련하는 메모리 효율적인 3D GAN 프레임워크인 계층적 약어 훈련(HA-GAN)을 제안한다. 부분 부피 선택과 약어 추론을 사용함으로써 HA-GAN은 메모리 사용량을 최대 80% 감소시키면서도 COPDGene 및 GSP 데이터셋에서 최신 기준(FID 점수)을 달성하여 계산 비용을 줄이고 고해상도 의료 영상 합성의 정밀도를 높인다.

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$ 크기의 저해상도 특징 맵으로 매핑한 후, 그룹 BatchNorm과 ReLU 활성화 함수를 사용한 잔차 블록을 통해 업샘플링된다.
• $G^L$ 및 $G^H$ 브랜치는 전치 3D 컨볼루션과 스kip 연결을 사용하여 중간 및 고해상도 출력을 생성하며, $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 의료 영상 분류 작업을 위한 데이터 증강을 통해 고해상도의 정밀한 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개 클래스 분류 작업에서 우수한 성능을 보였으며, 네트워크 아키텍처는 공간 계층성과 특징 표현을 유지하도록 설계되었다.