[論文レビュー] 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のリアルタイムストリーミングプロビクティック画像化を実現するAI駆動のエッジコンピューティングワークフローを提示する。このアプローチにより、従来手法と比較して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.

研究の動機と目的

• 反復的位相再構成アルゴリズムの遅さに起因するリアルタイムプロビクティック画像化の計算的ボトル neck を克服すること。
• プロビクティックにおけるオーバーサンプリングの必要性を排除することで、データ要件を低減し、低線量画像化を可能にすること。
• 埋め込みGPUシステムを用いたエッジでのリアルタイム推論をサポートするスケーラブルでエンドツーエンドのAIワークフローを構築すること。
• 最大2 kHzの検出器フレームレートまでストリーミングプロビクティック再構成の実現可能性を実証すること。
• 放射線感受性材料やシンクロトロンおよび電子顕微鏡における動的プロセスのための高速・低遅延画像化を可能にすること。

提案手法

• 3段階のワークフローを実装：(1) 最大2 kHzの検出器からのX線回折データストリーミング、(2) 逐次位相再構成（例：ePIE）を教師データとしてHPCリソース上でディープニューラルネットワークのオンライン学習、(3) ビームラインにおける埋め込みGPUでのリアルタイム推論。
• 複素数値画像を回折強度から直接再構成するための深層畳み込みニューラルネットワーク（PtychoNNおよびPtychoNN 2.0）を訓練。反復最適化を回避する。
• NVIDIA Jetson AGX XavierエッジデバイスにTensorRTを用いてモデルをデプロイし、ミリ秒未満の推論時間を達成。
• 新規に取得した実験データを用いたリアルタイムでのファインチューニングにより、変化するサンプル状態への継続的適応が可能になる。
• データ前処理には、低信号条件下でのロバストネスを向上させるために回折強度に対数スケーリングを適用。
• PyTorchからONNXフォーマットへのモデル変換により、エッジデバイスにおけるTensorRTを介した効率的デプロイと推論が可能になる。

実験結果

リサーチクエスチョン

• RQ1 ディープラーニングモデルは、最大2 kHzの検出器制限フレームレートまでリアルタイムでプロビクティック画像再構成を可能にするか？
• RQ2 AIベースの再構成は、高いオーバーラップとオーバーサンプリングを要する従来のプロビクティックと比較して、どの程度データ要件を低減できるか？
• RQ3 軽量ニューラルネットワーク（例：PtychoNN 2.0）は、埋め込みハードウェア上で大規模モデルと比較して再構成精度と推論速度の面でどの程度の性能を示すか？
• RQ4 HPCシステム上で実行されるオンライン連続学習により、新しい実験データに適応しながらも高い再構成忠実度を維持できるか？
• RQ5 低コストの埋め込みGPUプラットフォーム（例：Jetson AGX Xavier）では、リアルタイムプロビクティック画像化に必要な推論遅延とスループットはどの程度達成可能か？

主な発見

• 提示されたAI駆動のワークフローにより、最大2 kHzまで検出器フレームレートに一致するリアルタイムプロビクティック再構成が達成された。
• 10倍の光子数低減（1フレームあたり）でも構造的類似度（SSIM）が90%以上を維持し、低線量条件下でのロバストネスが示された。
• PtychoNN 2.0は、TensorRTを用いたJetson AGX Xavier上で2.3 ± 0.4 msの推論時間を達成し、ネイティブPyTorch推論と比較して4倍の高速化が実現された。
• オーバーサンプリングおよび高オーバーラップの必要性を排除することで、従来のプロビクティックと比較してデータ要件が桁違いに低減された。
• 最小限のサンプル損傷で低線量画像化が可能となり、生物学的標本や量子デバイスなどの放射線感受性材料にとって極めて重要である。
• エッジ推論遅延が1フレームあたり10 ms未満に抑えられており、リアルタイムフィードバックおよび動的画像ワークフローのサポートが可能になった。