[논문 리뷰] Deep learning at the edge enables real-time streaming ptychographic imaging

Anakha V Babu, Tao Zhou외 16명|arXiv (Cornell University)|2022. 09. 20.

Advanced X-ray Imaging Techniques인용 수 4

한 줄 요약

이 논문은 고성능 HPC 시스템에서의 온라인 딥 러닝 훈련과 임베디드 GPU 장치에서의 저지연 추론을 조합하여 최대 2 kHz에서 실시간 스트리밍 페르티오그래픽 영상 촬영을 가능하게 하는 AI 기반 엣지 컴퓨팅 워크플로우를 제시한다. 이 방법은 전통적인 방법보다 10배 낮은 데이터 전송률에서도 실시간 복원을 달성하며, 90% 이상의 구조적 유사도(SIM)를 유지함으로써 과도한 샘플링이 필요 없게 되어 저선량, 고속 나노스케일 영상 촬영이 가능해진다.

ABSTRACT

Coherent microscopy techniques provide an unparalleled multi-scale view of materials across scientific and technological fields, from structural materials to quantum devices, from integrated circuits to biological cells. Driven by the construction of brighter sources and high-rate detectors, coherent X-ray microscopy methods like ptychography are poised to revolutionize nanoscale materials characterization. However, associated significant increases in data and compute needs mean that conventional approaches no longer suffice for recovering sample images in real-time from high-speed coherent imaging experiments. Here, we demonstrate a workflow that leverages artificial intelligence at the edge and high-performance computing to enable real-time inversion on X-ray ptychography data streamed directly from a detector at up to 2 kHz. The proposed AI-enabled workflow eliminates the sampling constraints imposed by traditional ptychography, allowing low dose imaging using orders of magnitude less data than required by traditional methods.

연구 동기 및 목표

• 느린 반복적 위상 복원 알고리즘으로 인해 발생하는 실시간 페르티오그래픽 영상 촬영의 계산적 병목 현상을 해결하기 위해.
• 페르티오그래피에서 과도한 샘플링이 필요 없도록 하여 데이터 요구량을 줄이고 저선량 영상 촬영을 가능하게 하기 위해.
• 임베디드 GPU 시스템을 사용하여 엣지에서 실시간 추론을 지원하는 확장 가능한 종단 간 AI 워크플로우를 개발하기 위해.
• 최대 2 kHz의 검출기 프레임 속도에서 스트리밍 페르티오그래픽 복원의 가능성을 입증하기 위해.
• 빔에 민감한 물질과 동적 과정에 대한 고속, 저지연 영상 촬영을 가능하게 하기 위해, 싱크로트론 및 전자 현미경에서의 응용을 위해.

제안 방법

• 세 구성 요소로 이루어진 워크플로우가 구현됨: (1) 최대 2 kHz 속도로 검출기에서 유입되는 X선 衍射 데이터 스트리밍, (2) 반복적 위상 복원(e.g., ePIE)을 기준값으로 사용하여 HPC 자원에서 딥 뉴럴 네트워크의 온라인 훈련, (3) 범선에서의 임베디드 GPU에서의 실시간 추론.
• 복소수 영상 복원을 위해 딥 컨volution 신경망(PtychoNN 및 PtychoNN 2.0)을 훈련시켜 반복 최적화 과정을 생략함.
• NVIDIA Jetson AGX Xavier 엣지 장치에 TensorRT를 사용하여 추론 가속 기능을 구현하여 밀리초 이내의 추론 시간을 달성함.
• 최신 실험 데이터를 실시간으로 활용하여 네트워크를 미세 조정함으로써 변화하는 샘플 조건에 지속적으로 적응 가능하게 함.
• 저신호 조건에서의 강인성을 향상시키기 위해 행렬 강도에 로그 스케일링을 적용한 데이터 전처리 수행.
• PyTorch에서 ONNX 포맷으로의 모델 변환을 통해 엣지 장치에서의 효율적 배포 및 추론 가능.

실험 결과

연구 질문

• RQ1 딥 러닝 모델이 최대 2 kHz의 검출기 제한 프레임 속도에서 실시간 페르티오그래픽 영상 복원을 가능하게 할 수 있는가?
• RQ2 AI 기반 복원 방식은 기존 페르티오그래피가 요구하는 고중첩 및 과도한 샘플링이 필요 없기 때문에, 전통적 방법 대비 데이터 요구량을 얼마나 줄일 수 있는가?
• RQ3 경량 신경망(예: PtychoNN 2.0)의 성능은 임베디드 하드웨어에서 더 큰 모델 대비 복원 정확도와 추론 속도 측면에서 어떻게 비교되는가?
• RQ4 HPC 시스템에서의 온라인 연속 훈련이 새로운 실험 데이터에 적응하면서도 높은 복원 정밀도를 유지할 수 있는가?
• RQ5 저비용 임베디드 GPU 플랫폼인 Jetson AGX Xavier에서 실시간 페르티오그래픽 영상 촬영을 위한 추론 지연 시간과 처리량은 얼마인가?

주요 결과

• 제안된 AI 기반 워크플로우는 최대 2 kHz에서 검출기 최대 프레임 속도에 맞추어 실시간 페르티오그래픽 복원을 달성함.
• 모델은 10배 낮은 광자 수(1회 행렬 패턴당)에서도 여전히 90% 이상의 구조적 유사도(SSIM)를 유지하여 저선량 조건에서의 강인성을 입증함.
• Jetson AGX Xavier에서 TensorRT를 사용할 경우 PtychoNN 2.0의 추론 시간은 2.3 ± 0.4 ms이며, 원본 PyTorch 추론 대비 4배 빠름.
• 과도한 샘플링과 고중첩이 필요 없어지면서 기존 페르티오그래피 대비 데이터 요구량이 수개의 주기로 감소함.
• 최소한의 샘플 손상으로 저선량 영상 촬영이 가능해져 생물학적 샘플 및 양자 장치와 같은 빔에 민감한 물질에 매우 중요함.
• 엣지 추론 지연 시간이 1 프레임당 10ms 이내로 매우 낮아 실시간 피드백 및 동적 영상 촬영 워크플로우를 지원함.