Iterative Patch Selection for High-Resolution Image Recognition
This addresses a computational bottleneck for applications like autonomous driving and medical imaging by enabling efficient high-resolution image recognition under tight hardware constraints, representing a novel method for a known bottleneck.
The paper tackles the problem of training neural networks on high-resolution images, which is computationally challenging and causes out-of-memory errors, by proposing Iterative Patch Selection (IPS) to decouple memory usage from input size, enabling processing of arbitrarily large images with minimal GPU memory, such as finetuning on whole-slide images with up to 250k patches using only 5 GB of VRAM.
High-resolution images are prevalent in various applications, such as autonomous driving and computer-aided diagnosis. However, training neural networks on such images is computationally challenging and easily leads to out-of-memory errors even on modern GPUs. We propose a simple method, Iterative Patch Selection (IPS), which decouples the memory usage from the input size and thus enables the processing of arbitrarily large images under tight hardware constraints. IPS achieves this by selecting only the most salient patches, which are then aggregated into a global representation for image recognition. For both patch selection and aggregation, a cross-attention based transformer is introduced, which exhibits a close connection to Multiple Instance Learning. Our method demonstrates strong performance and has wide applicability across different domains, training regimes and image sizes while using minimal accelerator memory. For example, we are able to finetune our model on whole-slide images consisting of up to 250k patches (>16 gigapixels) with only 5 GB of GPU VRAM at a batch size of 16.