NAS-PRNet: Neural Architecture Search generated Phase Retrieval Net for Off-axis Quantitative Phase Imaging

2022-10-25

arXiv.org

AI Summary

To address the trade-off between reconstruction accuracy and inference speed in off-axis quantitative phase imaging (QPI), this work introduces differentiable neural architecture search (NAS) to QPI for the first time, proposing a lightweight encoder–decoder network. Methodologically, we extend SparseMask-based differentiable NAS to automatically optimize sparse skip connections between encoder and decoder modules; further, we integrate a MobileNet-v2 backbone and a composite loss function comprising a phase reconstruction term and a structural sparsity regularization. Experiments on biological cell datasets demonstrate a PSNR of 36.1 dB—significantly surpassing U-Net—while achieving single-frame inference in only 31 ms, a 12× speedup. Crucially, the learned architecture exhibits strong generalization across diverse interferometric systems, confirming its robustness and practical applicability in real-world QPI scenarios.

Abstract

Single neural networks have achieved simultaneous phase retrieval with aberration compensation and phase unwrapping in off-axis Quantitative Phase Imaging (QPI). However, when designing the phase retrieval neural network architecture, the trade-off between computation latency and accuracy has been largely neglected. Here, we propose Neural Architecture Search (NAS) generated Phase Retrieval Net (NAS-PRNet), which is an encoder-decoder style neural network, automatically found from a large neural network architecture search space. The NAS scheme in NAS-PRNet is modified from SparseMask, in which the learning of skip connections between the encoder and the decoder is formulated as a differentiable NAS problem, and the gradient decent is applied to efficiently search the optimal skip connections. Using MobileNet-v2 as the encoder and a synthesized loss that incorporates phase reconstruction and network sparsity losses, NAS-PRNet has realized fast and accurate phase retrieval of biological cells. When tested on a cell dataset, NAS-PRNet has achieved a Peak Signal-to-Noise Ratio (PSNR) of 36.1 dB, outperforming the widely used U-Net and original SparseMask-generated neural network. Notably, the computation latency of NAS-PRNet is only 31 ms which is 12 times less than U-Net. Moreover, the connectivity scheme in NAS-PRNet, identified from one off-axis QPI system, can be well fitted to another with different fringe patterns.

Problem

• Optimize neural networks for real-time phase retrieval
• Improve inference speed in quantitative phase imaging
• Enhance accuracy and efficiency in aberration compensation

Innovation

• Neural Architecture Search for optimized phase retrieval
• MobileNet-v2 encoder with synthesized loss function
• Gradient descent optimized skip connections