Image Gradient Decomposition for Parallel and Memory-Efficient Ptychographic Reconstruction

TL;DR

This paper introduces a novel gradient decomposition method for ptychographic reconstruction that significantly reduces memory usage and improves scalability and speed on GPU clusters, enabling high-resolution imaging with less resource consumption.

Contribution

The paper presents a new image gradient decomposition technique that tessellates data for memory efficiency and a parallel method enabling asynchronous communication, scaling efficiently on large GPU clusters.

Findings

Reduces memory footprint by 51 times.

Achieves 2.2-minute reconstruction time on 4158 GPUs.

Outperforms state-of-the-art algorithms in speed and scalability.

Abstract

Ptychography is a popular microscopic imaging modality for many scientific discoveries and sets the record for highest image resolution. Unfortunately, the high image resolution for ptychographic reconstruction requires significant amount of memory and computations, forcing many applications to compromise their image resolution in exchange for a smaller memory footprint and a shorter reconstruction time. In this paper, we propose a novel image gradient decomposition method that significantly reduces the memory footprint for ptychographic reconstruction by tessellating image gradients and diffraction measurements into tiles. In addition, we propose a parallel image gradient decomposition method that enables asynchronous point-to-point communications and parallel pipelining with minimal overhead on a large number of GPUs. Our experiments on a Titanate material dataset (PbTiO3) with 16632 probe locations show that our Gradient Decomposition algorithm reduces memory footprint by 51 times. In addition, it achieves time-to-solution within 2.2 minutes by scaling to 4158 GPUs with a super-linear strong scaling efficiency at 364% compared to runtimes at 6 GPUs. This performance is 2.7 times more memory efficient, 9 times more scalable and 86 times faster than the state-of-the-art algorithm.