Network architecture search of X-ray based scientific applications

Adarsha Balaji,Ramyad Hadidi,Gregory Kollmer,Mohammed E. Fouda,Prasanna Balaprakash
2024-04-17
Abstract:X-ray and electron diffraction-based microscopy use bragg peak detection and ptychography to perform 3-D imaging at an atomic resolution. Typically, these techniques are implemented using computationally complex tasks such as a Psuedo-Voigt function or solving a complex inverse problem. Recently, the use of deep neural networks has improved the existing state-of-the-art approaches. However, the design and development of the neural network models depends on time and labor intensive tuning of the model by application experts. To that end, we propose a hyperparameter (HPS) and neural architecture search (NAS) approach to automate the design and optimization of the neural network models for model size, energy consumption and throughput. We demonstrate the improved performance of the auto-tuned models when compared to the manually tuned BraggNN and PtychoNN benchmark. We study and demonstrate the importance of the exploring the search space of tunable hyperparameters in enhancing the performance of bragg peak detection and ptychographic reconstruction. Our NAS and HPS of (1) BraggNN achieves a 31.03\% improvement in bragg peak detection accuracy with a 87.57\% reduction in model size, and (2) PtychoNN achieves a 16.77\% improvement in model accuracy and a 12.82\% reduction in model size when compared to the baseline PtychoNN model. When inferred on the Orin-AGX platform, the optimized Braggnn and Ptychonn models demonstrate a 10.51\% and 9.47\% reduction in inference latency and a 44.18\% and 15.34\% reduction in energy consumption when compared to their respective baselines, when inferred in the Orin-AGX edge platform.
Machine Learning,Signal Processing
What problem does this paper attempt to address?
### Problems the Paper Aims to Solve The paper aims to address two key issues in X-ray and electron diffraction microscopy techniques: 1. **Bragg Peak Detection**: In high-energy diffraction microscopy (such as HEDM), it is necessary to accurately detect the positions of Bragg peaks. Traditional detection methods (such as pseudo-Voigt functions) are computationally complex and time-consuming, usually performed offline. 2. **Tomographic Imaging Reconstruction**: In tomographic imaging, it is necessary to recover the 3D structure of a sample from coherent diffraction images, which involves solving a complex inverse problem. Traditional algorithms are very computationally intensive, making real-time processing difficult. To improve the efficiency and accuracy of these tasks, the authors propose using deep neural networks (DNNs) and further automate the design and optimization of these neural network models through neural architecture search (NAS) and hyperparameter optimization (HPS). Specific goals include: - **Model Size**: Reduce the number of model parameters to lower storage and transmission costs. - **Energy Consumption**: Optimize the model to reduce energy consumption on edge devices. - **Throughput**: Increase the inference speed of the model to process large amounts of data in real-time. ### Main Contributions 1. **Architecture Optimization**: Using the Deephyper framework to optimize the architecture for two scientific applications (BraggNN and PtychoNN), demonstrating the Pareto optimality of model size and accuracy. Results show that the trainable parameters of BraggNN and PtychoNN models were reduced by 87.57% and 12.67%, respectively, while improving accuracy. 2. **Performance Improvement**: The Deephyper framework outperforms existing HPS and NAS tools (such as Optuna) in terms of convergence speed. 3. **Hardware Evaluation**: The optimized models were evaluated on the Nvidia Jetson AGX Orin platform, showing that the optimized BraggNN and PtychoNN models reduced inference latency and energy consumption by 10.51% and 9.47%, and 44.18% and 15.34%, respectively, compared to the baseline models. ### Method Overview 1. **Neural Architecture Search (NAS)**: - Defined the search space for multi-layer perceptrons (MLP) and deep convolutional networks (CNN). - Used the Deephyper framework for NAS and HPS to explore combinations of different architectures and hyperparameters. 2. **Hyperparameter Optimization (HPS)**: - Optimized hyperparameters through repeated sampling, evaluation, and updating of target values. - For multi-objective optimization problems, used the Chebyshev scalarization function to combine multiple objectives into a single scalar value. 3. **Model Optimization**: - Improved model accuracy and efficiency by optimizing hyperparameters such as learning rate, batch size, and training epochs. - Evaluated the optimized models on the Nvidia Jetson AGX Orin platform, comparing their inference latency and energy consumption. ### Experimental Setup - **Computational Resources**: Used the high-performance computing cluster Swing at Argonne National Laboratory, equipped with 6 nodes, each containing 2 AMD EPYC 7742 64-core processors and 8 NVIDIA A100 GPUs. - **Inference Platform**: Used the Nvidia Jetson AGX Orin 32GB module for inference, evaluating the performance of the optimized models in practical applications. ### Results Discussion - **Single-Objective Optimization**: - BraggNN: The optimized mlpBragg model outperformed the baseline model in terms of accuracy and model size, with an average MSE loss reduction of 31.3% and a model size reduction of 87.57%. - PtychoNN: The optimized cnnPtycho model outperformed the baseline model in terms of accuracy and model size, with an average MAE loss reduction of 16.77% and a model size reduction of 12.67%. - **Multi-Objective Optimization**: - BraggNN: The optimized mlpBragg model outperformed the baseline model in terms of accuracy and model size, with an average MSE loss reduction of 51.72% and a model size reduction of 67.37%. - PtychoNN: The optimized cnnPtycho model in