[논문 리뷰] OSS-Net: Memory Efficient High Resolution Semantic Segmentation of 3D Medical Data

Christoph Reich, Tim Prangemeier외 2명|arXiv (Cornell University)|2021. 10. 20.

Radiomics and Machine Learning in Medical Imaging인용 수 7

한 줄 요약

OSS-Net는 메모리 효율적이고 고해상도 3D 해부학적 분할을 위한 프레임워크로, 기존의 점유도 네트워크(O-Nets)에 局소 관측과 사전 지침 기반 추론 전략을 통합하여 확장한 것이다. 이는 3D U-Net 기준선과 유사한 정확도를 달성하면서도 메모리 사용을 크게 줄이고 추론 속도를 높여 소비자용 GPU에서도 고해상도 3D 분할을 가능하게 한다.

ABSTRACT

Convolutional neural networks (CNNs) are the current state-of-the-art meta-algorithm for volumetric segmentation of medical data, for example, to localize COVID-19 infected tissue on computer tomography scans or the detection of tumour volumes in magnetic resonance imaging. A key limitation of 3D CNNs on voxelised data is that the memory consumption grows cubically with the training data resolution. Occupancy networks (O-Nets) are an alternative for which the data is represented continuously in a function space and 3D shapes are learned as a continuous decision boundary. While O-Nets are significantly more memory efficient than 3D CNNs, they are limited to simple shapes, are relatively slow at inference, and have not yet been adapted for 3D semantic segmentation of medical data. Here, we propose Occupancy Networks for Semantic Segmentation (OSS-Nets) to accurately and memory-efficiently segment 3D medical data. We build upon the original O-Net with modifications for increased expressiveness leading to improved segmentation performance comparable to 3D CNNs, as well as modifications for faster inference. We leverage local observations to represent complex shapes and prior encoder predictions to expedite inference. We showcase OSS-Net's performance on 3D brain tumour and liver segmentation against a function space baseline (O-Net), a performance baseline (3D residual U-Net), and an efficiency baseline (2D residual U-Net). OSS-Net yields segmentation results similar to the performance baseline and superior to the function space and efficiency baselines. In terms of memory efficiency, OSS-Net consumes comparable amounts of memory as the function space baseline, somewhat more memory than the efficiency baseline and significantly less than the performance baseline. As such, OSS-Net enables memory-efficient and accurate 3D semantic segmentation that can scale to high resolutions.

연구 동기 및 목표

• 고해상도 의료 영상 분할에서 3D CNN의 입체적 메모리 증가 문제를 해결하기 위해.
• 기본 점유도 네트워크(O-Nets)의 표현력 부족과 느린 추론 문제를 극복하기 위해.
• O-Nets의 메모리 효율성과 3D CNN의 성능을 결합하여 3D 의미 분할을 수행하기 위해.
• 소비자용 GPU에서 고해상도 분할을 가능하게 하기 위해 메모리 사용과 추론 시간을 줄이기 위해.
• 복잡한 해부학적 구조(예: 뇌 종양, 간 병변)를 포괄하는 확장 가능한 기능 공간 기반 방법을 개발하기 위해.

제안 방법

• OSS-Net는 복잡한 3D 형태의 표현 능력을 향리하기 위해 국소 관측을 통합하여 원래 O-Net을 수정한다.
• 기존 추론 시간을 절반으로 줄이기 위해 사전 인코더 예측을 초기 상태로 사용하여 추론 속도를 향상시킨다.
• 한 번의 순방향 전파에서 다수의 3D 위치에서 점유도를 평가하는 미분 가능 샘플링 전략을 적용한다.
• 다중 척도 특징 추출 메커니즘이 체적 데이터에서 세밀한 해부학적 세부 정보를 포착할 수 있도록 한다.
• 기능 공간 표현을 위해 조정된 3D U-Net 유사 인코더-디코더 백본을 활용한다.
• 사전 예측 기반으로 동적으로 샘플링 위치를 결정하여 추론을 최적화함으로써 정확도를 유지하면서 효율성을 향상시킨다.

실험 결과

연구 질문

• RQ1 점유도 네트워크는 3D 의료 영상에서 3D CNN 수준의 분할 성능을 달성할 수 있는가?
• RQ2 O-Nets의 표현력을 어떻게 향상시켜 종양과 같은 복잡한 해부학적 구조를 모델링할 수 있는가?
• RQ3 O-Nets의 추론 속도를 정확도를 유지하면서 크게 향상시킬 수 있는가?
• RQ4 기능 공간 3D 분할에서 샘플링 위치와 메모리 소비 간의 최적의 트레이드오프는 무엇인가?
• RQ5 OSS-Net은 LiTS 간 분할 데이터셋과 같이 훈련 데이터가 제한된 도전적인 데이터셋에서 어떻게 성능을 내는가?

주요 결과

• BraTS 2020 뇌 종양 데이터셋에서 OSS-Net은 3D 잔차 U-Net(성능 기준선)과 유사한 분할 성능을 달성하였으며, Dice 스코어는 0.878을 기록하였다.
• LiTS 간 분할 데이터셋에서 OSS-Net은 O-Net 기준선과 2D U-Net 효율성 기준선을 모두 초월하여, Dice 스코어 0.852를 달성하였다.
• 기존 O-Net 접근 방식 대비 추론 속도가 두 배 향상되었으며, 2^12개 이상의 샘플링 위치를 초과해도 런타임이 일정하게 유지되었다.
• 추론 중 메모리 사용은 O-Net과 유사한 1.2GB이며, 3D U-Net 기준선(6.5GB)보다 5배 이상 낮아 소비자용 GPU에 배포가 가능했다.
• 기존 추론 방식과 개선된 추론 방식 간 예측 불일치율은 5.9658×10⁻⁴%에 불과하여 성능 저하가 거의 없음을 확인하였다.
• 훈련 중 OSS-Net의 메모리 사용은 2.6GB로, 3D U-Net 기준선(6.5GB)보다 크게 줄었으며, 높은 정확도를 유지하였다.