[논문 리뷰] LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation

Wentao Zhu, Can Zhao외 4명|arXiv (Cornell University)|2020. 06. 22.

Advanced Neural Network Applications참고 문헌 29인용 수 6

한 줄 요약

LAMP는 전체 이미지 입력을 사용하여 대규모 3D U-Net 및 SEU-Net 아키텍처를 훈련하기 위한 자동 모델 병렬화를 도입함으로써 분할 정확도와 추론 속도에 있어 뚜렷한 향상을 이끌어내었다. GPipe 기반의 병렬 U-Net 설계를 활용함으로써, 모델 아키텍처 변경 없이도 의료 영상 분할 작업에서 최신 기술 수준의 성능을 달성하였다.

ABSTRACT

Deep Learning (DL) models are becoming larger, because the increase in model size might offer significant accuracy gain. To enable the training of large deep networks, data parallelism and model parallelism are two well-known approaches for parallel training. However, data parallelism does not help reduce memory footprint per device. In this work, we introduce Large deep 3D ConvNets with Automated Model Parallelism (LAMP) and investigate the impact of both input's and deep 3D ConvNets' size on segmentation accuracy. Through automated model parallelism, it is feasible to train large deep 3D ConvNets with a large input patch, even the whole image. Extensive experiments demonstrate that, facilitated by the automated model parallelism, the segmentation accuracy can be improved through increasing model size and input context size, and large input yields significant inference speedup compared with sliding window of small patches in the inference. Code is available\footnote{https://monai.io/research/lamp-automated-model-parallelism}.

연구 동기 및 목표

• 모델 크기와 입력 컨텍스트 크기를 증가시킬 경우 3D 의료 영상 분할 정확도에 미치는 영향을 조사하기 위해.
• 의료 영상에서 대규모 3D ConvNet과 전체 이미지 입력을 훈련할 수 없게 만드는 GPU 메모리 제약을 해결하기 위해.
• 의료 영상 분석에서 널리 사용되는 표준 U-Net 아키텍처와 호환되는 자동 모델 병렬화 프레임워크를 개발하기 위해.
• 큰 수신 필드를 활용한 단일 통과 전체 이미지 추론으로 슬라이딩 윈도우 추론을 대체하여 추론 시간을 단축하기 위해.
• 클래스 불균형 문제를 다루기 위해 커리큘럼 훈련 및 적응형 가중 손실의 효과를 검증하기 위해.

제안 방법

• GPipe의 파이프라인 병렬화 기반으로 대규모 3D ConvNet 레이어를 여러 GPU에 분산하는 병렬 U-Net 아키텍처를 설계하기 위해.
• 수동 레이어 분할이나 아키텍처 변경 없이도 장치 간 네트워크 분할이 가능한 자동 모델 병렬화를 구현하기 위해.
• 큰 모델 훈련을 안정화시키기 위해 작은, 중간, 큰 입력 패치 순서로 단계적으로 훈련하는 커리큘럼 훈련 전략을 사용하기 위해.
• 작은 기관이나 병변을 포함한 의료 영상의 클래스 불균형 문제를 해결하기 위해 dice 손실과 포칼 손실을 조합한 적응형 가중 손실을 적용하기 위해.
• 큰 모델의 훈련 안정성을 확보하기 위해 RMSProp 옵티마이저를 사용하고 학습률 스케줄링 및 배치 정규화를 적용하기 위해.
• 큰 입력 패치(최대 512×512×704)를 사용하여 전체 이미지 추론을 가능하게 하여 슬라이딩 윈도우 추론이 필요 없도록 하기 위해.

실험 결과

연구 질문

• RQ1 3D ConvNet의 크기를 증가시킬 경우 의료 영상 데이터셋에서 분할 정확도에 어떤 영향을 미치는가?
• RQ2 슬라이딩 윈도우 추론과 비교해 전체 이미지 또는 큰 패치를 사용한 큰 입력 컨텍스트가 추론 속도와 정확도에 어떤 영향을 미치는가?
• RQ3 자동 모델 병렬화가 아키텍처 수정 없이도 표준 의료 영상 데이터셋에서 대규모 3D U-Net 및 SEU-Net 모델을 훈련시킬 수 있는가?
• RQ4 입력 크기를 점차 증가시키는 커리큘럼 훈련 전략이 큰 모델의 수렴 및 성능 향상에 어떤 영향을 미치는가?
• RQ5 큰 모델과 큰 입력을 사용하면서 적응형 가중 손실을 조합할 경우 소형 구조 분할에서의 클래스 불균형 문제를 어느 정도 완화할 수 있는가?

주요 결과

• 전체 이미지 입력을 사용한 SEU-Net-128 모델은 AnatomyNet보다 뛰어난 성능을 보였으며, 이는 큰 모델과 큰 컨텍스트가 정확도 향상에 기여함을 시사한다.
• 전체 이미지 입력을 사용함으로써 1장당 추론 시간이 1.52초로 단축되어, 64³ 패치를 사용한 슬라이딩 윈도우 추론 대비 5.8배 빠른 속도 향상을 달성하였다.
• 큰 입력을 사용한 더 큰 모델들(예: U-Net-128)은 간 분할에서 90.99%의 Dice 스코어, 종양 분할에서 56.48%의 Dice 스코어를 기록하여 더 작은 설정에 비해 뚜렷한 성능 향상을 보였다.
• 가장 큰 입력(256³)을 사용한 경우 U-Net-128의 추론 시간은 4.39초로 단축되었으며, 이는 전방 전파 횟수를 줄임으로써 큰 입력이 더 빠른 추론을 가능하게 함을 보여준다.
• 커리큘럼 훈련 전략은 큰 모델의 안정적인 훈련을 가능하게 하였으며, 점차 커지는 패치로 훈련할수록 더 높은 정확도를 달성하였다.
• 자동 모델 병렬화 프레임워크는 블록당 최대 128개 필터와 입력 크기 최대 512×512×704까지의 모델 훈련에 성공하였으며, GPU 메모리 제약을 극복하였다.