[논문 리뷰] Fast and Accurate 3D Medical Image Segmentation with Data-swapping Method

Haruki Imai, Samuel Matzek외 3명|arXiv (Cornell University)|2018. 12. 19.

Advanced Neural Network Applications인용 수 8

한 줄 요약

이 논문은 중간 특징을 GPU 및 CPU 메모리 간에 동적으로 이동시켜, 패치 없이 전체 크기의 192×192×192 의료 영상에서 3D U-Net 모델을 훈련할 수 있도록 하는 데이터 스위핑 기법을 제안한다. 이 방법은 통신 오버헤드를 17.1% 감소시키며, 전체 종양 및 종양 핵에 대해 각각 4.48%와 5.32% 향상된 Dice 점수를 달성하고, 훈련 속도를 3.53배(164시간에서 47시간)로 가속화한다.

ABSTRACT

Deep neural network models used for medical image segmentation are large because they are trained with high-resolution three-dimensional (3D) images. Graphics processing units (GPUs) are widely used to accelerate the trainings. However, the memory on a GPU is not large enough to train the models. A popular approach to tackling this problem is patch-based method, which divides a large image into small patches and trains the models with these small patches. However, this method would degrade the segmentation quality if a target object spans multiple patches. In this paper, we propose a novel approach for 3D medical image segmentation that utilizes the data-swapping, which swaps out intermediate data from GPU memory to CPU memory to enlarge the effective GPU memory size, for training high-resolution 3D medical images without patching. We carefully tuned parameters in the data-swapping method to obtain the best training performance for 3D U-Net, a widely used deep neural network model for medical image segmentation. We applied our tuning to train 3D U-Net with full-size images of 192 x 192 x 192 voxels in brain tumor dataset. As a result, communication overhead, which is the most important issue, was reduced by 17.1%. Compared with the patch-based method for patches of 128 x 128 x 128 voxels, our training for full-size images achieved improvement on the mean Dice score by 4.48% and 5.32 % for detecting whole tumor sub-region and tumor core sub-region, respectively. The total training time was reduced from 164 hours to 47 hours, resulting in 3.53 times of acceleration.

연구 동기 및 목표

• 고해상도 의료 영상에서 깊이 있는 3D U-Net 모델을 훈련할 때 GPU의 메모리 제약을 해결하기 위해.
• 패치 기반 훈련 방법으로 인해 객체가 패치 간에 분할되어 발생하는 품질 저하를 제거하기 위해.
• 효율적으로 GPU 및 CPU 메모리 간 중간 활성화를 관리하는 데이터 스위핑 기법을 개발하여 더 큰 효과적 GPU 메모리처럼 작동하도록 하기 위해.
• 공간 연속성의 손실 없이 더 빠르고 정확한 3D 의료 영상 분할을 달성하기 위해.

제안 방법

• 훈련 중에 GPU 메모리에서 중간 특징 맵을 CPU 메모리로 오프로드하여 GPU 메모리 압박을 줄이는 데이터 스위핑 방법을 사용한다.
• 메모리 사용량과 계산적 의존성에 기반해 GPU 및 CPU 간 데이터 전송을 동적으로 스케줄링한다.
• 3D U-Net에 특화되어 있어 통신 오버헤드를 최소화하기 위해 데이터 전송 빈도와 크기를 최적화한다.
• 패치 없이 전체 192×192×192 볼륨 영상에서 엔드 투 엔드 훈련을 가능하게 하여 공간적 맥락을 유지한다.
• 자주 액세스되는 특징을 우선순위로 정하고 중복 전송을 최소화하는 맞춤형 메모리 관리 전략을 사용한다.
• 표준 평가 지표인 Dice 점수와 훈련 시간을 사용하여 뇌 종양 분할 데이터셋에서 방법을 평가한다.

실험 결과

연구 질문

• RQ1 데이터 스위핑은 패치 없이 전체 크기의 3D 의료 영상에서 3D U-Net을 훈련시킬 수 있으며, 분할 정확도를 유지하거나 향상시킬 수 있는가?
• RQ2 훈련 효율성 측면에서 데이터 스위핑의 통신 오버헤드는 표준 패치 기반 훈련과 비교해 어떻게 되는가?
• RQ3 패치를 피할 경우 소형 또는 분할된 종양 영역에서 분할 성능 향상은 어느 정도 이루어지는가?
• RQ4 훈련 시간을 최소화하고 모델 정확도를 최대화하기 위해 데이터 스위핑 파라미터의 최적 설정은 무엇인가?
• RQ5 패치 기반 대안과 비교해 데이터 스위핑 방법은 더 빠른 수렴과 더 나은 일반화 성능을 달성할 수 있는가?

주요 결과

• 데이터 스위핑 방법은 기준 대비 17.1% 통신 오버헤드를 감소시켜 훈련 효율을 크게 향상시켰다.
• 전체 종양 하위 영역에 대해 패치 기반 훈련 대비 평균 Dice 점수가 4.48% 향상되었다.
• 종양 핵 하위 영역에 대해 Dice 점수는 5.32% 향상되어 중요한 영역의 정확한 국소화를 보여주었다.
• 총 훈련 시간은 164시간에서 47시간으로 감소하여 훈련 속도가 3.53배 가속화되었다.
• 정확도와 속도 모두에서 패치 기반 훈련을 능가했으며, 전체 볼륨 맥락을 유지함으로써 이점이 있음을 입증했다.
• 결과적으로 데이터 스위핑은 패치와 관련된 잡음과 성능 손실 없이 고해상도 3D 분할을 가능하게 한다는 것이 확인되었다.