Towards Scalable & GPU Accelerated Earth Science Imagery Processing: An AI/ML Case Study

The expansion of Artificial intelligence and machine learning (AI/ML) in the last few decades is directly associated with the advancements in high-performance computing (HPC) hardware. Many super-computing centers around the country have been enhancing their computing capabilities by integrating new systems, software stacks, and hardware resources. One of the most prolific and popular computing technologies of these times are Graphical Processing Units (GPUs). Designed for parallel processing, GPUs are used in many graphics applications, and have opened the paradigm of General-purpose computing on GPUs (GPGPU). GPGPU allows the GPU to perform computations in applications traditionally handled by the central processing unit (CPU) to dramatically accelerate computational workloads. Earth science imagery data holdings have become a big data analytics problem, and the use of GPGPUs techniques to scale computations is a must. In this work, we enable the acceleration of common Earth science imagery processing pipelines by means of parallel computing and GPU-accelerated data structures. A Python module was developed to fully process the imagery on the GPUs and to provide gains of more than 2x compared to running on serial CPUs. NVIDIA RAPIDS framework is utilized as the backend, where CuPy and cuDF are leveraged for both array and dataframe-like objects, while Dask is used for multi-GPU scalability. In addition, we developed a set of Docker and Singularity containers to make the module and all its dependencies portable and reproducible across HPC environments. All benchmarks were run at the NASA Center for Climate Simulation PRISM GPU cluster and by leveraging NVIDIA v100 GPUs. Here we present several use cases and examples of how this module and its underlying ecosystem can be leveraged to speed up the development and deployment of AI/ML software, including the processing of Landsat, Sentinel-2 and WorldView imagery.