Scalable and accurate multi-GPU-based image reconstruction of large-scale ptychography data

Abstract

While the advances in synchrotron light sources, together with the development of focusing optics and detectors, allow nanoscale ptychographic imaging of materials and biological specimens, the corresponding experiments can yield terabyte-scale volumes of data that can impose a heavy burden on the computing platform. Although graphics processing units (GPUs) provide high performance for such large-scale ptychography datasets, a single GPU is typically insufficient for analysis and reconstruction. Several works have considered leveraging multiple GPUs to accelerate the ptychographic reconstruction. However, most of these works utilize only the Message Passing Interface to handle the communications between GPUs. This approach poses inefficiency for a hardware configuration that has multiple GPUs in a single node, especially while reconstructing a single large projection, since it provides no optimizations to handle the heterogeneous GPU interconnections containing both low-speed (e.g., PCIe) and high-speed links (e.g., NVLink). In this paper, we provide an optimized intranode multi-GPU implementation that can efficiently solve large-scale ptychographic reconstruction problems. We focus on the maximum likelihood reconstruction problem using a conjugate gradient (CG) method for the solution and propose a novel hybrid parallelization model to address the performance bottlenecks in the CG solver. Accordingly, we have developed a tool, called PtyGer (Ptychographic GPU(multiple)-based reconstruction), implementing our hybrid parallelization model design. A comprehensive evaluation verifies that PtyGer can fully preserve the original algorithm’s accuracy while achieving outstanding intranode GPU scalability.

Image Gradient Decomposition for Parallel and Memory-Efficient Ptychographic Reconstruction

“…To explain how the state-of-the-art Halo Voxel Exchange method works [8], [9], [7], we reuse the 9 overlap probe locations example from Fig. 1(b) in Fig.…”

“…Despite its capability in delivering unprecedented image resolution, ptychographic imaging requires enormous memory to store diffraction measurements and 3D image pixels, also called "Voxels" [7], [8], [9]. This memory requirement, in turn, constrains the achievable ptychographic image resolution and stops many scientists from discovering important features that are present in the samples.…”

“…To address the memory constraints and improve reconstruction speed, the state-of-the-art Halo Voxel Exchange algorithm takes advantage of the fact that ptychography diffraction measurements are often organized into a sequence of matrices indexed in time order [7], [8], [9]. Each time step of the data acquisition is called a "Probe Location" and the diffraction measurements at a probe location corresponds to a group of spatially contiguous voxels in the shape of a circle in the image domain.…”

Ultrafast Error-bounded Lossy Compression for Scientific Datasets

“…Over the past decade, GPUs have become prevalent because of their massive parallelism and computational power [25]. Various applications have been successfully accelerated on GPU-based platforms [36,37,[39][40][41][42]. Because of the high demand for ultrafast error-bounded lossy compressors, a few specific error-controlled lossy compression algorithms have been developed for GPU accelerators; cuSZ [31] and cuZFP [10] are two leading ones.…”

SOLAR: A Highly Optimized Data Loading Framework for Distributed Training of CNN-based Scientific Surrogates

“…They can provide accurate results, however, are quite time-consuming, even with the accelerations on high-performance computing (HPC) platforms. For instance, the traditional iterative phase retrieval algorithm for ptychographic image reconstruction takes hours or even days to converge on multiple GPUs with only a medium-scale dataset [47]. To accelerate the reconstruction, physicists develop a CNN-based autoencoder (named PtychoNN [1]) to approximate the reconstructed images from raw X-ray data and achieve 300× speedup compared to the traditional physical approach [6].…”