[论文解读] Iterative Patch Selection for High-Resolution Image Recognition

Benjamin Bergner, Christoph Lippert+1 more|arXiv (Cornell University)|Oct 24, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出迭代块选择（IPS），一种针对高分辨率图像识别的内存高效方法，通过使用交叉注意力Transformer迭代选择显著块，将GPU内存使用量与输入尺寸解耦。IPS在仅使用5 GB显存（批量大小为16）的情况下，实现了对千兆像素图像（例如25万块）的训练，同时保持高精度，与基线方法相比显著降低了内存消耗。

ABSTRACT

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.

研究动机与目标

• 解决在标准GPU上因显存消耗过高而难以训练深度神经网络处理高分辨率图像的挑战。
• 实现显存使用量与输入尺寸的解耦，使在严格硬件限制下处理任意大小图像成为可能。
• 开发一种在不同领域和图像尺寸下均能以最小加速器显存实现高精度的方法。
• 实现在大规模医学和真实世界图像数据集上高效微调和自监督预训练。

提出的方法

• IPS通过将高分辨率图像划分为图像块，并使用共享编码器独立嵌入每个图像块来处理图像。
• 采用交叉注意力Transformer在无梯度模式下对图像块进行评分，以识别最显著的图像块，使用注意力分数作为选择标准。
• 该方法采用迭代式、自回归的缓冲区更新机制，以保持前M个最显著的图像块，通过每次处理I个图像块的块来管理显存。
• 通过第二个交叉注意力Transformer聚合选定的图像块，形成用于分类的全局图像表征。
• 分类头与图像块选择和聚合模块联合训练，且梯度在反向传播时被截断，不再流经选定的图像块和注意力权重。
• 该方法基于多实例学习，支持弱监督多标签分类，并可进行图像块级别的概率估计。

实验结果

研究问题

• RQ1 基于图像块的方法是否能在保持恒定GPU显存使用量（无论输入尺寸如何）的前提下，实现高分辨率图像的高精度识别？
• RQ2 与标准CNN和Transformer相比，通过交叉注意力实现的迭代图像块选择在大输入数据上的内存效率和性能表现如何？
• RQ3 IPS在仅使用极少显存的情况下，能在多大程度上实现对全切片图像（如16GB以上）的微调？
• RQ4 IPS是否能以低显存开销支持大规模数据集（如CAMELYON16）上的多任务学习和自监督学习？
• RQ5 超参数M（选定图像块数量）和I（每次迭代的批量大小）对计算效率和模型性能有何影响？

主要发现

• IPS在仅使用5 GB GPU显存（批量大小为16）的情况下，实现了对包含最多25万个图像块（>16 GB像素）的全切片图像的微调。
• 在将1k×1k的MNIST图像缩放至10k×10k的千兆像素图像时，峰值显存使用量保持恒定在1.7 GB，而同等规模的CNN在2k×2k分辨率下显存消耗高达24.6 GB。
• 该方法在交通标志识别、合成千兆像素MNIST的多任务分类以及CAMELYON16基准的自监督预训练任务中均取得了高精度。
• 在单张A100 GPU上高效完成了三个基准任务的训练与微调，总运行时间分别为45分钟（交通标志）、5小时（MNIST）和2小时（CAMELYON16）。
• 交叉注意力机制生成的注意力分数支持弱监督图像块级别类别概率估计，可视化为彩色编码的注意力图。
• 消融实验确认，迭代图像块选择与交叉注意力聚合的结合是维持低显存使用量和高性能的关键。