[Paper Review] Deep learning at the edge enables real-time streaming ptychographic imaging

Anakha V Babu, Tao Zhou+16 more|arXiv (Cornell University)|Sep 20, 2022

Advanced X-ray Imaging Techniques4 citations

TL;DR

This paper presents an AI-powered edge computing workflow that enables real-time, streaming ptychographic imaging at up to 2 kHz by combining online deep learning training on high-performance HPC systems with low-latency inference on embedded GPU devices. The approach achieves real-time reconstruction with 90% structural similarity at 10× lower data rates than conventional methods, eliminating the need for oversampling and enabling low-dose, high-speed nanoscale imaging.

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.

Motivation & Objective

• To overcome the computational bottleneck of real-time ptychographic imaging caused by slow iterative phase retrieval algorithms.
• To reduce data requirements and enable low-dose imaging by eliminating the need for oversampling in ptychography.
• To develop a scalable, end-to-end AI workflow that supports real-time inference at the edge using embedded GPU systems.
• To demonstrate the feasibility of streaming ptychographic reconstruction at detector frame rates up to 2 kHz.
• To enable high-speed, low-latency imaging for beam-sensitive materials and dynamic processes in synchrotron and electron microscopy.

Proposed method

• A three-component workflow is implemented: (1) X-ray diffraction data streaming from a detector at up to 2 kHz, (2) online training of a deep neural network on HPC resources using iterative phase retrieval (e.g., ePIE) as ground truth, and (3) real-time inference on an embedded GPU at the beamline.
• A deep convolutional neural network (PtychoNN and PtychoNN 2.0) is trained to reconstruct complex-valued images directly from diffraction intensities, bypassing iterative optimization.
• The model is deployed on a NVIDIA Jetson AGX Xavier edge device using TensorRT for inference acceleration, achieving sub-millisecond inference times.
• The network is fine-tuned in real time using newly acquired experimental data, enabling continual adaptation to changing sample conditions.
• Data preprocessing includes logarithmic scaling of diffraction intensities to improve robustness under low-signal conditions.
• Model conversion from PyTorch to ONNX format enables efficient deployment and inference via TensorRT on the edge device.

Experimental results

Research questions

• RQ1 Can deep learning models enable real-time ptychographic image reconstruction at detector-limited frame rates of up to 2 kHz?
• RQ2 To what extent can AI-based reconstruction reduce data requirements compared to conventional ptychography, which demands high overlap and oversampling?
• RQ3 How does the performance of lightweight neural networks (e.g., PtychoNN 2.0) compare to larger models in terms of reconstruction accuracy and inference speed on embedded hardware?
• RQ4 Can online, continuous training of the neural network on HPC systems maintain high reconstruction fidelity while adapting to new experimental data?
• RQ5 What is the achievable inference latency and throughput on low-cost embedded GPU platforms like the Jetson AGX Xavier for real-time ptychographic imaging?

Key findings

• The proposed AI-enabled workflow achieves real-time ptychographic reconstruction at up to 2 kHz, matching the maximum detector frame rate.
• The model maintains a structural similarity (SSIM) of over 90% even at 10× lower photon counts per diffraction pattern, demonstrating robustness under low-dose conditions.
• PtychoNN 2.0 achieves inference times of 2.3 ± 0.4 ms on the Jetson AGX Xavier using TensorRT, a 4× speedup over native PyTorch inference.
• The system reduces data requirements by orders of magnitude compared to conventional ptychography by eliminating the need for oversampling and high overlap.
• The workflow enables low-dose imaging with minimal sample damage, crucial for beam-sensitive materials such as biological specimens and quantum devices.
• The edge inference latency is sufficiently low (under 10 ms per frame) to support real-time feedback and dynamic imaging workflows.