训练高分辨率图像任务，突破 GPU 内存限制，Jetson Nano 上也能起飞！

集智书童公众号

发布于 2024-05-28 21:20:35

2100

发布于 2024-05-28 21:20:35

尽管在图像识别领域取得了快速进展，但处理高分辨率图像仍然是一个计算挑战。然而，这种处理对于从自动驾驶导航到医学成像分析等领域提取详细的目标洞察至关重要。作者的研究引入了一个框架，旨在通过利用基于内存高效 Patch 处理来缓解这些挑战，用于高分辨率图像。它结合了全局上下文表示和局部 Patch 信息，使得能够全面理解图像内容。与受内存限制的传统训练方法不同，作者的方法能够训练超高分辨率图像。作者通过在分类、目标检测和分割的7个不同基准测试中表现出卓越的性能来证明作者方法的有效性。值得注意的是，即使在资源受限的设备如Jetson Nano上，提出的方法也能取得强大的性能。代码： https://github.com/Visual-Conception-Group/Localized-Perception-Constrained-Vision-Systems。

I Introduction

图像识别领域近年来取得了巨大进步，这得益于像AlexNet ，ResNet ，ViT 和Swin家族这样的突破性架构。这些架构擅长处理像ImageNet [5]这样的数据集，这些数据集主要包含分辨率低于一百万像素的自然图像。然而，许多关键应用依赖于对高分辨率图像的分析，以提取有关感兴趣目标的复杂细节。

例如，自动驾驶汽车需要能够检测远处的交通标志。同样，在医学成像中，准确的诊断往往依赖于检查能够揭示关于细胞和组织的关键信息的吉像素显微镜切片[6]。

虽然高分辨率图像具有无可否认的优势，但将深度学习模型训练在这样的数据上却带来了重大挑战。将模型架构扩展以适应这些巨大的输入会导致计算和内存需求很高。

输入图像大小的看似线性增加会导致计算复杂性和内存使用的二次增长[7]。这些要求很容易造成硬件瓶颈，阻碍训练。

现有解决方案，如流式处理[8]，梯度预训练权重[9]，遍历网络[7]，迭代 Patch 选择[10]和Zoom-In [11]，主要解决了分类任务。

然而，高分辨率图像在诸如目标检测和分割等其他视觉任务中提出了挑战。上述模型不能直接扩展用于这些任务。像[12]，AdaZoom [13]和PRDet [14]这样的目标检测方法解决了小目标检测或训练大分辨率图像的问题。分割方法[15, 16]提出了GPU内存高效的技术。然而，这些检测和分割方法没有深入研究在严格的GPU内存限制下训练非常大分辨率图像的能力。

研究在如此限制下训练模型的方法是必要的，因为这将帮助模型即使在部署到资源受限的系统后也能进行训练。

在本文中，作者提出了一种框架，利用基于块的处理[17]与全局上下文理解，在GPU内存限制下进行训练。

作者的贡献如下：

首先，作者提出了一种用于图像分类、目标检测和分割的统一模型。
作者使用规范的主干网络作为从局部和全局块中提取特征的提取器。然后，聚合网络从这些块中提取所需的细粒度和全局语义。
其次，在作者的实验中，作者证明了基于块的处理有助于在GPU内存预算下训练大图像。此外，这种局部和全局处理能够捕捉到进行密集任务（如目标检测和分割）所需的细粒度细节，以及用于分类的全局语义。作者的模型概述在图1中给出。
作者选择图像分类、目标检测和分割任务，因为它们是不同的，而且架构差异很大。因此，一个针对一个任务的直接扩展方法不会轻易适用于其他任务。此外，分类主要是一项全局任务，需要理解整个图像。另一方面，目标检测和分割是密集任务，需要全局和空间理解。

II 相关工作

分类：卷积神经网络（CNNs）和 Transformer （Transformers）在图像识别任务中表现出色。尽管作者的模型可以扩展到基于 Transformer 的方法，但为了简洁起见，作者专注于CNNs。由于内存需求巨大，在高分辨率图像上训练这些模型仍然是一个挑战[10]。最近的工作通过各种方法着重解决这一挑战。一种常用的方法是先降低图像分辨率，然后再将其输入网络。这种技术实施起来很简单，但可能导致丢失关键信息，尤其是在处理小型目标或细微细节时。当处理4096×4096尺寸的超高分辨率图像时，问题变得更加严重。

另一项工作涉及内存占用低的网络架构[18, 19]。这些架构通常通过使用深度可分离卷积技术或减少网络中的通道数来实现内存效率。然而，与标准CNN架构相比，这些方法可能会牺牲准确性。PatchGD[17]使用基本网络从非重叠图像块中提取特征，然后使用另一个轻量级网络进行池化。注意力机制最近已成为提高神经网络效率的有力工具[20]。这些机制允许网络关注输入图像最相关的部分。然而，注意力机制本身可能计算成本很高，特别是对于高分辨率图像。

检测：以前的文献针对非常大图像上的小目标检测[21]。在一项相关工作中，为了推动小目标检测的进展，还提出了如DOTA[22]，SODA-A和SODA-D[23]等数据集。[12]解决了在有限内存约束下高分辨率图像中小目标检测的问题。虽然这和 Proposal 的工作在精神上相近，但有两个主要区别。首先，[12]仅针对行人和车辆设计，而作者的方法没有这样的限制。其次，[12]独立处理每个图块，并且无法获得全局视图，而作者的方法通过学习以及使用完整图像的调整大小版本来提取全局理解。AdaZoom[13]和PRDet[14]网络在大型和无人机视角场景中的目标检测方面进行了创新，重点关注克服由小目标和尺度变化带来的挑战。QueryDet[24]和SDPDet[25]也解决了类似的挑战。

分割：超高分辨率图像分割在如无人机（UAVs）等应用中由于计算资源有限而构成重大挑战。在诸如制图和决策制定等各项任务中，对内存效率模型的的需求变得至关重要。由于全局分割需要大量的内存需求，一些研究专门使用局部图像块来解决这个问题[26, 27, 28, 29, 30]。还引入了如URUR这样的超高分辨率图像（5120×5120）[31]数据集来评估分割性能。在通用图像分割的情况下，[15]沈等人提出了用于超高分辨率图像分割优化的连续细化模型（CRM）。[16]提出了一种针对超高分辨率图像分割的GPU内存效率框架。

预备知识：作者的方法受到基于图像块处理的启发，这里作者简要回顾一下在PatchGD [17]算法中提出的基于图像块的分类技术。与一次性更新整个图像不同，PatchGD 对小图像块进行操作，假设这种局部化的方法可以在迭代过程中有效地覆盖整个图像，同时产生有效的模型更新。此外，这种图像块处理也有助于对非常大的图像进行训练。PatchGD 的核心是构建一个深度潜在表示，记为

Z\in\mathbb{R}^{m\times n\times d}

块，其中

是块的总数，

表示特征维度。

在之前的更新步骤中从图像的不同部分获取信息。

块作为完整输入图像的编码，是通过使用参数为

\theta_{1}

的基础模型独立处理图像块生成的。

为了池化特征，使用了一个小型子网络

\theta_{2}

。重要的是，集成这个子网络的计算成本非常小。在训练期间，两个模型组件

\theta_{1}

和

\theta_{2}

都会被更新。PatchGD 方法克服了传统梯度下降方法在处理大图像时的局限性，通过避免对整个图像样本进行更新，而是在局部块上计算梯度。这种方法提高了内存效率和计算的可扩展性，使其非常适合在分类中处理高分辨率图像。

III Proposed Methodology

在以下内容中，作者描述了作者的分类、检测和分割方法。作者的方法受到了PatchGD [17]的启发。然而，PatchGD只解决了分类任务。与分类相反，目标检测和分割还需要更好的空间理解。因此，作者的方法解决了这一挑战。

Image Classification

令

X\in\mathbb{R}^{M\times N\times c}

表示一个尺寸为

M\times N

的

通道输入图像。作者首先将

分割成大小为

M/m\times N/n\times c

的非重叠图像块，并使用

\theta_{1}

提取特征。在受限内存的情况下，使用所有图像块是不可行的，从计算资源的角度来看，这相当于使用全分辨率图像。为了通过图像块高效处理图像，随机抽取

个图像块的一个子集。这被称为内迭代中的采样。然后组合这些特征以获得潜在表示

Z\in\mathbb{R}^{m\times n\times d}

，其中

是来自每个图像块的特征的维度。在每次内迭代中，只有

个图像块在

中被更新。因此，在

次内迭代结束时，模型只看到了相当于

个给定图像的图像块。

尽管在任何给定的迭代中只更新图像块的部分子集，但

块累积并保留了所有采样图像块的特征。这种累积表示使模型能够保持全局视角，并确保图像不同部分之间的连贯性。它使得可以从大规模数据中有效学习，而无需一次将整个图像保存在内存中。PatchGD从不同的图像块聚合特征以推理全局语义。然而，作者发现对于非常大的图像，如4096

\times

4096，它并不足以捕捉全局上下文。此外，全局语义推理也严重依赖于采样率。如果作者增加采样率，内存大小会增加，这使得在资源受限的设备上进行训练变得困难。

为了解决上述限制，作者使用了一个全局上下文。这是通过将全图像的下采样版本作为一个图像块与其他图像块一起使用来实现的。作者称这个块为全局块。全局块的维度为

M/m\times N/n\times c

。通过

\theta_{1}

获得的全球块的特征向量

与其他图像块的特征融合，以获得联合表示

Z_{g}=Z+g

。然后作者将

Z_{g}

传递给

\theta_{2}

以获得分类标签。训练目标由

\mathcal{L}_{cls}=\mathcal{L}_{CE}(X,p_{j},\theta_{1},\theta_{2})

给出。

\theta_{1}

和

\theta_{2}

表示模型参数，

p_{j}

表示在第

次内迭代中从图像

中采样的图像块。由于只在

中更新

个图像块，因此只针对那些更新样本计算梯度。这有助于高效使用GPU内存。值得注意的是，还有其他融合

和

的方法，比如拼接。作者发现对于分类任务，加法可以得到更好的结果。

作者的方法的框图在图2(a)中提供。

Object Detection

作者考虑将经典的目标检测器

\theta_{1}

作为基础模型。例如，可以使用 Anchor-Free 点目标检测算法FCOS [32] 作为作者的基础模型。这里有

个局部 Patch 和一个全局 Patch 。作者将这些 Patch 分别输入到

个 Head 。利用检测器的

个 Head 和

mn+1

个 Patch ，作者得到单个 Head 的特征表示为

Z\in\mathbb{R}^{(mn+1)\times H\times W\times d}

。

的更新类似于第三节A部分描述的分类更新。

\theta_{1}

包含了 Backbone 网络、金字塔和 Head 。对于每个 Head 输出，作者将维度重构成

H\times W\times d(mn+1)

。重构后的特征接着输入到

\theta_{2}

，它包含一个单一的CNN层。这一步骤之后，有

个输出，维度为

H\times W\times d

。然后它被传递到分类、中心度和回归的最后一层，以得到实际的边界框和标签，类似于 [32] 中采用的方法。作者使用FCOS [32] 目标训练网络，该目标包括Focal Loss [33] 和IOU损失 [34]。作者方法的框图在图2(b)中提供。

Segmentation

Patch 采样器采样

个 Patch 。这些 Patch 的特征，由

\theta_{1}

获得，然后用来替换潜在特征表示

的特征。全局 Patch 的特征与 Patch 特征融合。聚合特征

随后使用模型

\theta_{2}

进行池化。更新的

用于计算损失，并且仅针对更新后的特征计算梯度。每次内循环后保存梯度，并且梯度累积进一步用于高效利用GPU内存。

类似于分类，作者将一个经典的分割模型作为基础模型

\theta_{1}

。例如，U-net可以用作作者的基础模型。每个 Patch 被输入到模型

\theta_{1}

中。输出形式为

M/m\times N/n\times C

。所有

个 Patch 的相应特征随后被平铺，以创建图像的联合表示，其维度为

Z\in\mathbb{R}^{M\times N\times C}

。

的更新类似于第III-A节中分类的更新。

当资源受限时，这种表示对于GPU内存来说过于庞大而难以处理。为了减少潜在表示的大小，作者将

\theta_{1}

的分类层替换为1

\times

1卷积。在内存有限的情况下，作者选择输出通道为8以便进行处理。联合表示随后通过

\theta_{2}

来学习全局上下文。为了进一步提供全局上下文，作者使用全局 Patch 并通过

\theta_{1}

获得其表示。作者进一步将其上采样到

M\times N\times C

，并与 Patch 表示进行拼接，使得

Z_{g}\in\mathbb{R}^{M\times N\times 2C}

。这里，

\theta_{2}

由1

\times

1卷积组成以获得最终的 Mask 。Dice损失和BCE（二元交叉熵）被用作训练目标。由于目标大小与背景相比可能非常小，Dice损失能很好地处理类别不平衡。

IV Experiments

算法的实现是在PyTorch中完成的。作者使用AdamW作为优化器。在分类和检测的情况下，作者使用1e-3作为初始学习率。作者将批量大小限制为4以训练模型。作者还采用了在[35]中提出的方法论，加入学习率预热和线性衰减。在分割的情况下，作者使用1e-4的学习率。在作者的实验中，作者证明了当模型在GPU内存的限制下训练时，作者的方法比 Baseline 有显著优势。此外，这一观察表明，通过使用小分辨率 Patch ，作者的方法可以在内存较少的设备上进行训练，同时还能在较高分辨率图像上匹配模型的性能。为了进一步解决内存限制问题，作者使用了梯度累积。在每次内部迭代中作者计算梯度并保存。作者将梯度累积步骤设置为：分类最多3步，检测为10步，分割为2步。

Image Classification

作者在Prostate cANcer graDe Assessment (PANDA) 和 UltraMNIST 数据集上进行了作者提出的分类算法的实验。训练的轮数（epochs）为100。

PANDA: 在PANDA数据集上，作者使用24GB和16GB内存训练作者提出的算法。作者用8,616个样本进行训练，1,000个样本进行验证，1,000个样本进行推理，并在表1中报告推理结果。ResNet50 [2] 同时用于

\theta_{1}

和

\theta_{2}

。作者考虑了不同图像大小：512

\times

512，2048

\times

2048和4096

\times

4096。 Patch 大小为128

\times

128或256

\times

256。在512

\times

512的情况下，作者的方法准确率达到59.3%，而PatchGD达到53.4%。注意，[17]中没有提到确切的分割，因此作者重现了结果（用*表示）。在非常高分辨率的4096

\times

4096，且 Patch 大小为256

\times

256的情况下，作者的方法比PatchGD高出6.9%。这是因为作者使用全局 Patch 提供全局上下文。作者还注意到，在这种情况下，作者方法的峰值内存（GB(peak)）消耗为21GB，而PatchGD为20GB。因此，在作者的情况下内存使用略有增加，但在准确性方面有大幅提升。对于512

\times

512的图像，如果 Patch 大小为128

\times

128，则有16个 Patch 。以20%的采样率（表中的S），每个内部迭代（表示为Iter.）随机选择3个 Patch ，总共4次迭代选择12个 Patch 。

UltraMNIST: 在这里，作者在Jetson Nano (4GB内存)上训练作者的算法。作者用22,400个样本进行训练，2,800个样本进行验证，2,801个样本进行推理。图像分辨率为512

\times

512， Patch 大小为256

\times

256，内部迭代为3，采样率为25%。作者使用ResNet50 (RN50) 以及 MobileNetV2 (MNV2) [38] 作为两个网络，并在表2中报告结果。作者要强调的是，Jetson Nano的GPU内存有限，因此使用512

\times

512进行训练可能不稳定。当应用SGD时，这种情况就很明显。使用ResNet50的SGD大约得到52.9%，而使用MobileNetV2得到67.1%。另一方面，PatchGD在使用ResNet50时达到62.1%。相比之下，作者的情况下提升了1.2%。作者在MobileNetV2上也观察到了类似的情形。此外，与在PANDA中的观察类似，作者看到了峰值内存使用的小幅增加。

Object Detection

提出的检测算法实验在COCO数据集[42]，驾驶场景下的小目标检测数据集（SODA-D）[23]和TJU-Ped-traffic[43]上进行。作者使用FCOS[32]作为

\theta_{1}

。作者训练模型300个周期。使用两个Nvidia 2080ti GPU（每个11GB内存）进行训练。

COCO：作者遵循[32]中给出的协议。由于图像大小为640

\times

480，作者创建了一个大小为4096

\times

4096的空白图像，并向其中添加36张图像，以在这些大合成图像上训练和测试模型。作者将新的合成图像调整为[32]中图像大小相当的800

\times

1024，并训练FCOS算法。作者还将在表3中提到的分辨率下缩放的合成图像上训练CenterNet、CornerNet和YOLOX。作者在表3中呈现作者的结果。在

方面，作者的模型获得了25.4%，而FCOS为21.5%。在参数数量方面，作者的模型增加了大约3.1%的轻微开销。由于基于块的处理，作者还看到延迟增加了。YOLOX获得23.3%的得分。作者看到其他指标也有类似的趋势。结果表明，在给定的GPU内存下，作者的训练方法更有效，因为它可以捕捉到细粒度以及全局细节。

在峰值内存使用方面，作者的方法使用了20GB，YOLOX使用了17GB，FCOS使用了15GB。尽管与FCOS相比作者多出了5GB，但作者可以看到，提出的方法允许在大型图像上进行训练，同时提高了性能。

预测边界框 使用作者提出的目标检测算法对新的合成图像的预测边界框在图3(a)中给出。作者还给出了使用FCOS和YOLOX的预测边界框，分别在图3(b)和3(c)中。可以推理出，当图像尺寸较大时，FCOS未能检测到许多小到中等大小的目标，而YOLOX表现良好，但未能超越提出的模型。

SODA-D：作者遵循[23]中的协议。平均图像大小为3407

\times

2470。作者将图像调整为4096

\times

4096。作者在表4中呈现作者的结果。

{}^{\dagger}

表示结果来自[23]。作者的模型达到了28.6%的

，而FCOS为23.9%，YOLOX为26.7%。在

、

AP_{50}

和

AP_{75}

方面，作者的方法优于其他方法。值得注意的是[23]在1200

\times

1200分辨率的块上运行检测器，块状的检测结果映射回原始图像，然后应用非最大抑制来消除重叠和冗余的预测。这是在利用GPU内存和处理非常高分辨率图像之间的有效折中。作者的方法使用了全图像分辨率的4096

\times

4096和512

\times

512的块大小。作者的方法的峰值内存使用量为18GB。[23]使用了4个NVIDIA GeForce RTX 3090 GPU来训练模型。从表4可以推理，在内存限制下，提出的方法处理高分辨率图像的折中方案更好。

预测边界框的结果在图4中给出。

TJU-Ped-traffic：作者遵循[43]中给出的协议。作者使用13,858个样本进行训练，2,136个样本进行验证，以及4,344个样本进行推理。其中52%的图像分辨率为1624

\times

1200，48%的图像分辨率至少为2560

\times

1440。作者在表5中呈现作者的结果。作者的方法获得

为29.45，YOLOX达到27.2%，FCOS获得23.3%。作者看到其他指标也有类似的趋势。作者还注意到，作者方法的峰值内存为20GB，而YOLOX为18GB，FCOS为14GB。这表明，在内存限制下，提出的方法可以处理全分辨率图像，并给出更好的性能。

Segmentation

作者使用U-net [44] 和 DeepLab v3 [45] 作为

\theta_{1}

。这些模型在3090Ti GPU（24GB内存）以及Jetson Nano上进行100个周期的训练。

DRIVE - 视网膜血管数据集：DRIVE [46] 数据集共有100张图像，每张图像的分辨率为512

\times

512。按照文献[47]中的协议，作者将数据集分为80个用于训练，20个用于测试。在表6中，U-net和DeepLab v3的Type Full版本是在完整图像（512

\times

512分辨率）上训练的。在推理时也使用完整的分辨率。U-net和DeepLab v3的Type Downsampled（下采样）版本是在图像下采样到128

\times

128时进行训练的。在推理时，输入被调整到128

\times

128，输出 Mask 被上采样到512

\times

512。Ours

{}^{\ddagger}

表示模型在 Proposal 的框架下使用128

\times

128大小的裁剪进行训练，不包括全局块。最后，Ours表示使用带有全局块的128

\times

128大小的裁剪进行训练。

Type Full版本使用完整图像，因此占用更多GPU内存。作者将其作为参考给出。作者特别关注与Type Downsampled的比较，因为它只使用1.39 GB的GPU内存，因此可以在4GB的限制内进行训练。在F1、IoU、平衡准确率、PLR和NLR方面，作者方法与U-net Down相比有显著提升。F1从Type Downsampled的61.0%增加到Ours的79.8%。IoU从43.9%增加到66.5%。作者还可以看到，Ours的结果与Full相似或更好。作者还发现全局块在略微增加峰值GPU使用量的情况下，提高了整体性能。当使用DeepLab v3作为

\theta_{1}

时，作者也得出了类似的结论。作者在图5中给出了定性结果。作者可以看到，图5(f)中的连续性和细节都得到了很好的保留。

航拍图像数据集：此数据集包含72张图像，每张图像的分辨率为1024

\times

1024。作者使用52个样本进行训练，20个样本进行测试。作者观察到，Ours

{}^{\ddagger}

和Ours在准确度方面都优于Type Downsampled，尽管在其他指标上略有落后。然而，在DeepLab v3中，作者看到Ours

{}^{\ddagger}

和Ours在准确度、IoU、平衡准确率、PLR和NLR方面都优于Downsampled。作者还使用Jetson Nano进行训练。结果在表8中报告。作者发现对于DeepLab v3和Unet，只有下采样图像（128

\times

128大小）可以进行训练，而Ours

{}^{\ddagger}

可以在1024

\times

1024上进行训练。与下采样版本相比，所有分数都有大幅提升。与分类不同，分割是一项资源密集型任务，为了能够在Jetson Nano上进行训练，作者不得不进行很多优化步骤。作者进行了附加优化

例如增加交换内存、从OpenCV切换到PIL并禁用增强，这使得训练成功。然后利用作者的方法，作者将训练扩展到512和1024分辨率的图像。

V Conclusion

在这项工作中，作者解决了在严格内存限制下高效处理高分辨率图像的难题，用于分类、目标检测和分割等任务。作者创新的框架整合了基于局部 Patch 的处理与全局上下文理解，使得在内存限制下能够进行全面图像分析。这种方法不仅保留了准确目标检测和分割所需的细粒度细节，还融入了对稳健分类性能至关重要的全局语义。在七个不同基准测试上的实验表明，作者的方法取得了具有竞争力的性能。作者还展示了在资源受限的设备上，如Jetson Nano上进行训练。