[论文解读] 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 Techniques被引用 4

一句话总结

本文提出了一种基于人工智能的边缘计算工作流程，通过结合高性能HPC系统上的在线深度学习训练与嵌入式GPU设备上的低延迟推理，实现了高达2 kHz的实时流式波前衍射成像。该方法在数据速率比传统方法降低10倍的情况下，实现了90%的结构相似性，消除了对过采样的需求，从而实现了低剂量、高速的纳米级成像。

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.

研究动机与目标

• 解决由缓慢的迭代相位恢复算法引起的实时波前衍射成像计算瓶颈。
• 通过消除波前衍射中对过采样的需求，降低数据需求，实现低剂量成像。
• 开发一种可扩展的端到端AI工作流程，利用嵌入式GPU系统实现实时边缘推理。
• 展示在高达2 kHz的探测器帧率下流式波前衍射重建的可行性。
• 为同步辐射和电子显微镜中的束敏感材料及动态过程提供高速、低延迟成像能力。

提出的方法

• 实施了三部分工作流程：(1) 从探测器以高达2 kHz的速率流式传输X射线衍射数据，(2) 在HPC资源上使用迭代相位恢复（如ePIE）作为真实值，对深度神经网络进行在线训练，(3) 在光束线的嵌入式GPU上实现实时推理。
• 训练了一个深度卷积神经网络（PtychoNN和PtychoNN 2.0），直接从衍射强度重建复数值图像，绕过迭代优化过程。
• 使用TensorRT将模型部署在NVIDIA Jetson AGX Xavier边缘设备上，实现毫秒级以下的推理时间。
• 通过新获取的实验数据实时微调网络，实现对样本条件变化的持续适应。
• 数据预处理包括对衍射强度进行对数缩放，以提高在低信号条件下的鲁棒性。
• 将模型从PyTorch转换为ONNX格式，实现通过TensorRT在边缘设备上高效部署和推理。

实验结果

研究问题

• RQ1 深度学习模型能否在高达2 kHz的探测器帧率下实现实时波前衍射图像重建？
• RQ2 与传统波前衍射（需高重叠度和过采样）相比，基于AI的重建在多大程度上可减少数据需求？
• RQ3 在嵌入式硬件上，轻量级神经网络（如PtychoNN 2.0）在重建精度和推理速度方面与大模型相比表现如何？
• RQ4 在HPC系统上对神经网络进行在线持续训练，能否在适应新实验数据的同时保持高重建保真度？
• RQ5 在低功耗嵌入式GPU平台（如Jetson AGX Xavier）上，实时波前衍射成像的可实现推理延迟和吞吐量是多少？

主要发现

• 所提出的AI增强工作流程实现了高达2 kHz的实时波前衍射重建，与最大探测器帧率一致。
• 即使每张衍射图案的光子数降低10倍，该模型仍保持超过90%的结构相似性（SSIM），证明了在低剂量条件下的鲁棒性。
• PtychoNN 2.0在使用TensorRT的Jetson AGX Xavier上实现了2.3 ± 0.4 ms的推理时间，相比原生PyTorch推理速度提升4倍。
• 通过消除对过采样和高重叠度的需求，该系统相比传统波前衍射将数据需求降低了数量级。
• 该工作流程实现了极低剂量成像，最大限度减少样品损伤，对生物样本和量子器件等束敏感材料至关重要。
• 边缘推理延迟足够低（每帧低于10 ms），可支持实时反馈和动态成像工作流。