[논문 리뷰] High Resolution Medical Image Analysis with Spatial Partitioning

Le Hou, Youlong Cheng외 7명|arXiv (Cornell University)|2019. 09. 06.

Medical Image Segmentation Techniques참고 문헌 18인용 수 29

한 줄 요약

이 논문은 메시 텐서플로우에서 공간 분할 프레임워크를 제안하여 512×512×512 밴델스까지의 고해상도 의료 영상에서 3D U-Net 모델을 엔드 투 엔드로 훈련할 수 있도록 한다. 할로 인터페이스를 통해 분할된 패치 간의 컨volution 연결성을 유지한다. 우리가 아는 바로는, 이는 내림샘 또는 데이터 자르기 없이 전체 해상도 CT 스캔에서 직접 이러한 모델을 훈련하는 데 있어 첫 번째 방법이며, 훈련 시간에 5%의 오버헤드만으로 LiTS 벤치마크에서 최고 성능을 기록한다.

ABSTRACT

Medical images such as 3D computerized tomography (CT) scans and pathology images, have hundreds of millions or billions of voxels/pixels. It is infeasible to train CNN models directly on such high resolution images, because neural activations of a single image do not fit in the memory of a single GPU/TPU, and naive data and model parallelism approaches do not work. Existing image analysis approaches alleviate this problem by cropping or down-sampling input images, which leads to complicated implementation and sub-optimal performance due to information loss. In this paper, we implement spatial partitioning, which internally distributes the input and output of convolutional layers across GPUs/TPUs. Our implementation is based on the Mesh-TensorFlow framework and the computation distribution is transparent to end users. With this technique, we train a 3D Unet on up to 512 by 512 by 512 resolution data. To the best of our knowledge, this is the first work for handling such high resolution images end-to-end.

연구 동기 및 목표

• GPU/TPU 메모리 제약으로 인해 고해상도 의료 영상(예: 10⁸–10⁹ 밴델스를 가진 3D CT 스캔)에서 CNN을 직접 훈련하는 것이 불가능한 문제를 해결하기 위해.
• 자르기, 내림샘 또는 굵은 해상도에서 세밀한 해상도로의 방법 등 기존 접근 방식의 한계를 극복하기 위해, 정보 손실과 복잡한 구현을 초래함.
• 전체 해상도 3D 의료 영상 분석을 위한 효율적이고 투명한 모델 및 데이터 병렬 처리를 가능하게 하기 위해.
• TPU와 GPU 모두에서 고해상도 데이터를 대상으로 하는 확장 가능하고 오픈소스 프레임워크를 개발하기 위해.
• 합성 기반 데이터 증강 방법을 도입하여 간 종양 세그멘테이션의 일반화 성능을 향상시키기 위해.

제안 방법

• 공간 분할은 고해상도 3D 영상을 겹치지 않는 패치로 나누고, 이를 여러 GPU/TPU에 분산 배포한다.
• 각 컨볼루션 레이어 전에 할로 교환을 수행하며, 이는 인접한 장치들이 패치의 가장자리(커널 크기의 절반)를 교환하여 공간적 맥락을 유지한다.
• 프레임워크는 메시 텐서플로우 기반으로 구축되어 최소한의 사용자 간섭으로 자동 데이터 및 모델 병렬 처리를 가능하게 한다.
• 합성 기반 데이터 증강 방법은 현실적인 강도와 형태 변화를 가진 합성 종양을 생성하여 모델의 강건성을 향상시킨다.
• 내림샘 또는 패치 기반 추론 없이 512×512×512 해상도 데이터에서 3D U-Net 모델의 엔드 투 엔드 훈련을 지원한다.
• 혼합 정밀도(반정밀도) 부동소수점 연산과 조합된 Dice 및 교차 엔트로피 손실을 사용한 Adafactor 옵timizer를 사용하여 훈련한다.

실험 결과

연구 질문

• RQ1 내림샘 또는 자르기 없이 전체 해상도 512×512×512 CT 스캔에서 3D U-Net 모델을 엔드 투 엔드로 훈련할 수 있는가?
• RQ2 할로 교환을 동반한 공간 분할이 고해상도 의료 영상에서 최소한의 계산 오버헤드로 효율적이고 확장 가능한 훈련을 가능하게 하는가?
• RQ3 합성 기반 데이터 증강 방법이 3D CT 스캔에서 간 종양 세그멘테이션의 일반화 및 성능 향상에 기여하는가?
• RQ4 제안된 프레임워크를 사용할 때 입력 해상도가 증가함에 따라 LiTS 벤치마크에서의 Dice 스코어는 어떻게 변화하는가?
• RQ5 분산 TPU 환경에서 공간 분할과 할로 교환으로 인해 실제로 발생하는 훈련 시간 오버헤드는 얼마인가?

주요 결과

• 제안된 프레임워크는 512×512×512 해상도 CT 스캔에서 3D U-Net 모델을 성공적으로 훈련하여 LiTS 검증 세트에서 평균 케이스 당 Dice 스코어 0.4547 ± 0.0475를 달성한다.
• 512×512×512 해상도에서 전역 Dice 스코어는 0.7180 ± 0.0446에 도달하여 대용량 세그멘테이션에서 뛰어난 성능을 보인다.
• 분할, 재형성 및 할로 교환 연산으로 인해 약 5%의 추가 훈련 시간 오버헤드만 발생한다.
• 제안된 데이터 증강이 없을 경우, 케이스 당 Dice 스코어가 최소 10% 감소하여 이 방법이 모델 일반화에 핵심적인 역할을 한다는 것을 입증한다.
• 더 높은 해상도 입력(예: 512³)은 더 낮은 해상도(예: 64³, 128³)보다 일관되게 더 높은 Dice 스코어를 기록하며, 전체 해상도 학습의 이점이 확인된다.
• 프레임워크는 TPU와 GPU 훈련을 모두 지원하며, 네트워크 아키텍처를 수정하지 않고도 장치 간 계산을 투명하게 분산시킬 수 있다.