Enabling Next Generation Machine Learning for Large Scale Image Analysis

Project Summary/Abstract Deep learning has transformed medical image analysis by delivering clinically meaningful results on challenging problems like prostate cancer detection and lung cancer screening. FDA approval of whole-slide digital pathology imaging (WSIs) for primary diagnosis is further increasing interest, adoption, and investment in artificial intelli- gence (AI) technology for pathology. Learning from large medical images using patient-level labels (PLLs) has become an active computational pathology research area. PLLs such as pathology diagnosis or clinical outcomes are generated through healthcare operations and are often readily available. In contrast to learning paradigms that depend on the expert annotation of images (e.g., delineating tumor regions) and are therefore time-intensive and limited to smaller cohorts, training directly from WSIs using PLLs will allow the development of realistic training datasets containing tens-of-thousands of subjects that can produce models with clinically-meaningful ac- curacy. GPU accelerators have played a significant role in advancing deep learning methods for computational pathology tools. Machine Learning Frameworks (MLFs), e.g., Pytorch and TensorFlow, provide researchers with abstractions to quickly develop models that utilize GPUs. The evolution of GPUs and MLFs has been driven by the analysis of small images, and so applying these tools directly to WSIs or other large medical images like volumetric magnetic resonance or computed tomography is challenging. Adapting medical imaging problems to the small image paradigm leads to many compromises resulting in suboptimal performance, increased imple- mentation effort, and increased software/design complexity (e.g., patch based techniques or multiple instance learning). As a result, the development of scalable ML models from PLLs by directly processing WSI images through a deep learning pipeline is infeasible today on GPUs. Recent efforts that use unified GPU memory or streaming approaches to overcome GPU memory limits and attempt to perform end-to-end training at WSI scale have demonstrated superior performance to annotation or MIL. However, these approaches are either slow (due to suboptimal data movement strategies), complex to adapt/use, or highly specific to a given network architecture (limiting the ability to develop and explore new architectures). More general-purpose, efficient, and user-friendly frameworks are needed to allow the development of WSI scale deep learning. This project will develop a robust software framework to facilitate seamless development and use of scalable ML models, without the imposition of any limits on the sizes of handled images, unhindered by the limited memory capacity in GPUs. The proposed SSTEP (Seamless Scalable Tensor-Expression Execution via Partitioning) soft- ware framework will allow scalable and portable neural network models that directly process full high-resolution images of arbitrary size for training or inference, on any (multi) GPU platform. SSTEP will allow the development of novel deep learning paradigms that are purpose-built for medical applications, and will enable developers to rapidly create and evaluate these tools using familiar MLFs - PyTorch or TensorFlow.