[논문 리뷰] AirLab: Autograd Image Registration Laboratory

Robin Sandkühler, Christoph Jud외 2명|arXiv (Cornell University)|2018. 06. 26.

Medical Image Segmentation Techniques참고 문헌 41인용 수 49

한 줄 요약

AIRLab은 자동 미분 및 GPU 가속 실행을 통해 의료 영상 정합의 빠른 프로토타이핑과 재현성을 강화하는 Python 기반의 PyTorch 지원 환경을 제공합니다.

ABSTRACT

Medical image registration is an active research topic and forms a basis for many medical image analysis tasks. Although image registration is a rather general concept specialized methods are usually required to target a specific registration problem. The development and implementation of such methods has been tough so far as the gradient of the objective has to be computed. Also, its evaluation has to be performed preferably on a GPU for larger images and for more complex transformation models and regularization terms. This hinders researchers from rapid prototyping and poses hurdles to reproduce research results. There is a clear need for an environment which hides this complexity to put the modeling and the experimental exploration of registration methods into the foreground. With the "Autograd Image Registration Laboratory" (AIRLab), we introduce an open laboratory for image registration tasks, where the analytic gradients of the objective function are computed automatically and the device where the computations are performed, on a CPU or a GPU, is transparent. It is meant as a laboratory for researchers and developers enabling them to rapidly try out new ideas for registering images and to reproduce registration results which have already been published. AIRLab is implemented in Python using PyTorch as tensor and optimization library and SimpleITK for basic image IO. Therefore, it profits from recent advances made by the machine learning community concerning optimization and deep neural network models. The presented draft of this paper outlines AIRLab with first code snippets and performance analyses. A more exhaustive introduction will follow as a final version soon.

연구 동기 및 목표

• 의료 영상 정합에서 복잡한 목표와 대용량 3D 데이터를 다루기 위한 빠른 프로토타이핑 환경의 필요성을 제시합니다.
• 등록 목적에 대한 자동 그래디언트 계산을 가능하게 하는 오픈형의 차별화 가능한 프레임워크로 AIRLab을 소개합니다.
• 빠른 실험을 용이하게 하는 빌딩 블록(유사도 측정, 변환, 규제화)과 예제 방법을 제공합니다.
• 초기 실험 및 코드 예제를 통해 성능 및 재현성의 이점을 시연합니다.]
• method:[

제안 방법

• PyTorch autograd를 통한 등록 목적의 자동 미분.
• f 및 그 매개변수를 위한 모듈식 유사도 측정( MSE, NCC, LCC, SSIM, MI, NGF ) 및 규제화(확산, 비등방성/등방성 TV, 희소성, Demons).
• 선형/치밀, 비선형/보간(예: B-스플라인 및 Wendland 커널), 조밀한/하이브리드 변환 모델에 대한 지원.
• 고정 이미지 도메인에서의 잔차를 계산하기 위한 백워드 이미지 워핑과 GPU 가속 그리드 샘플링.
• 지수 매핑을 갖는 정상 속도장으로의 고정 도메인 변형의 선택적 미분가능 변환.
• POPi와 같은 데이터셋에 대한 자동 등록 평가 파이프라인.
• 매개변수가 큰 문제에 대한 PyTorch의 공통 옵티마이저(예: LBFGS, Adam) 이용 가능.
• 다양한 이미지 크기에 대한 CPU vs. GPU 실행 성능 벤치마크.

실험 결과

연구 질문

• RQ1 자동 미분이 등록 목적의 빠른 프로토타이핑을 단순화하고 수동 그래디언트 도출 없이 재현 가능한 실험을 가능하게 할 수 있는가?
• RQ2 다양한 유사도 측정 및 규제화가 2D/3D 데이터에서 등록 성능과 안정성에 어떤 영향을 주는가?
• RQ3 다양한 변환 모델에 대해 CPU vs. GPU에서 AIRLab의 실제 성능 특성(속도, 확장성)은 어떠한가?
• RQ4 AIRLab이 Rigid, Free Form Deformation, Diffeomorphic Demons와 같은 표준 등록 방법들을 하나의 미분 가능한 프레임워크 내에서 얼마나 재현할 수 있는가?

주요 결과

실험	하드웨어	64	128	256	512	1024	2048	4096
Dense + R_{diff} + S_{MSE}	CPU	2.29	3.29	5.57	14.11	46.44	187.68	832.77
Dense & R_{diff} + S_{MSE} with Diffeomorph	CPU	4.52	9.75	25.81	81.55	309.06	1367.49	5834.58
Dense + R_{diff} + S_{MSE}	GPU GTX 1080	4.41	4.39	4.39	4.29	4.63	9.21	30.91
Dense & R_{diff} + S_{MSE} with Diffeomorph	GPU GTX 1080	7.36	7.32	6.59	6.61	9.14	24.34	89.85

• AIRLab은 등록 목적을 위한 자동 그래디언트 계산을 가능하게 하여 수동 구현 노력을 줄입니다.
• 다양한 유사도 측정 및 규제화와 여러 변환 모델을 제공하여 등록 방법의 빠른 프로토타이핑을 가능하게 합니다.
• 초기 실험에서 AIRLab 내의 Rigid, FFD, Diffeomorphic Demons에 대한 등록 워크플로우와 해당 시각적 결과를 시연합니다.
• 대규모 2D 입력의 경우 CPU 대비 GPU 기반 실행이 계산 시간을 크게 단축시키고 이미지 크기 확장에 따른 스케일링을 보여줍니다.
• POPI와 같은 데이터셋에서 재현성을 촉진하기 위한 자동 평가 파이프라인이 포함되어 있습니다.
• 실행 가능한 코드 스니펫과 B-스플라인 보간 및 다채널 데이터와 같은 추가 기능 통합 계획이 제시됩니다.