大尺寸图像分类检测分割统一模型：Resource Efficient Perception for Vision Systems

论文题目：Resource Efficient Perception for Vision Systems

论文链接：http://arxiv.org/abs/2405.07166

代码链接：https://github.com/Visual-Conception-Group/Localized-Perception-Constrained-Vision-Systems

作者设计了一个统一的模型，用于图像分类、目标检测和分割，训练还可用Jetson Nano进行。

一、摘要

文中提出了一种框架，旨在通过利用内存高效的基于patch的处理方法来缓解这些挑战。它结合了全局上下文表示和局部patch信息，实现了对图像内容的全面理解。与受内存限制的传统训练方法不同，所提方法能够处理超高清图像。通过在7个不同的基准测试（包括分类、目标检测和分割）上表现出色，证明了所提方法的有效性。
在这里插入图片描述

分类主要是一个全局任务，需要理解整个图像。相比之下，目标检测和分割是密集型任务，需要全局和空间理解的结合。

文章目录

- 一、摘要
- 二、关键问题
- 三、原理
- - 1、分类
  - 2、检测
  - 3、分割
- 四、实验
- - 1、图像分类
  - 2、目标检测
  - 3、分割
- 五、总结

二、关键问题

对于如 $4096 \times 4096$ 这样的大图像，它可能不足以充分捕捉全局上下文。此外，全局语义的推断高度依赖于采样率。如果增加采样率，内存大小会增加，这在资源受限的设备上训练变得困难。

三、原理

1、分类

本文引入全局上下文。这通过使用全图像的下采样版本作为与其他块一起的块来实现，称其为全局块。全局块的维度为 $M/m \times N/n \times c$ 。通过 $\theta_{1}$ 获取的全局块特征向量 $g$ 与其它块的特征融合，得到联合表示 $Z_{g} = Z + g$ 。然后，将 $Z_{g}$ 传递给 $\theta_{2}$ 以获取分类标签。训练目标为：
$L_{cls} = L_{CE} (X, p_{j}, \theta_{1}, \theta_{2} ),$
其中 $\theta_{1}$ 和 $\theta_{2}$ 表示模型参数， $p_{j}$ 表示在第 $j$ 个内循环中从图像 $X$ 中采样的块。由于仅 $k$ 个块在 $Z$ 中更新，计算梯度仅针对这些更新的样本，这有助于高效利用 GPU 内存。值得注意的是，还有其他方法将 $Z$ 和 $g$ 融合，如连接。实验发现对于分类任务，加法效果更好。所提方法的块图示如图 2(a) 所示。
在这里插入图片描述

2、检测

以经典的物体检测器 $\theta_{1}$ 作为基础模型（例如，无锚点的FCOS算法 [32] 可以作为基础）。有 $mn$ 个局部区域和一个全局区域。将这些区域分别输入到 $h$ 个头中。通过 $h$ 个检测器头和 $mn + 1$ 个区域，单个头获取特征表示为 $\in R^{(mn+1) \times H \times W \times d}$ 。 $Z$ 的更新与分类中描述的分类更新类似。 $\theta_{1}$ 包含backbone、pyramid 和 heads。对于每个头的输出，将维度重塑为 $\times W \times d(mn+ 1)$ 。重塑后的特征输入到 $\theta_{2}$ ，它包含一个单层卷积神经网络（CNN）。经过这一步，得到 $h$ 个输出，维度为 $\times W \times d$ 。然后，这些输出传递到分类、中心化和回归的最终层，以获取实际的边界框和标签，这与 [32] 中的方法类似。使用FCOS [32] 的目标函数进行网络训练，其中包括focal loss [33] 和IOU loss [34]。所提方法的框图如图 2(b) 所示。

3、分割

与分类任务类似，文中也将一个经典的分割模型（如 U-Net）作为基础模型 $\theta_{1}$ 。每个小块作为输入送入模型 $\theta_{1}$ ，输出形式为 $M/m \times N/n \times C$ 。所有 $mn$ 个小块的特征被拼接，形成图像的联合表示，维度为 $\in R^{M \times N \times C}$ 。 $Z$ 的更新方式类似于上上节中的分类任务。

当资源受限时，这种表示过于庞大，无法在 GPU内存中处理。为了减小潜在表示的大小，用一个 $\times 1$ 卷积层替换 $\theta_{1}$ 的分类层，选择输出通道数为8，以便在内存有限的情况下进行处理。然后，将联合表示传递给 $\theta_{2}$ ，以学习全局上下文。为了进一步提供全局信息，使用全局小块并通过 $\theta_{1}$ 获取其表示，** 将其上采样至 $\times N \times C$ 大小，然后与小块表示拼接，得到 $Z_{g} \in R^{M \times N \times 2C}$ ** 。在这里， $\theta_{2}$ 包含一个 $\times 1$ 卷积层以生成最终的掩码。训练目标采用Dice 损失和 BCE（二元交叉熵）。由于物体相对于背景可能非常小，Dice损失有助于很好地处理类别不平衡问题。

四、实验

这些算法的实现使用了PyTorch。采用AdamW作为优化器。对于分类和检测任务，设置初始学习率为1e-3。训练模型时，将批次大小限制为4。还应用了 [35]中提出的学习率warmup和线性衰减策略。在进行分割任务时，使用的学习率是1e-4。实验结果显示，当在GPU内存限制下训练模型时，所提方法显著优于基线。这表明，所提方法能够在使用小分辨率切片的同时，通过在高分辨率图像上保持模型性能，实现在内存较小的设备上训练。为了进一步缓解内存限制，采用了梯度累积技术。在每次内部迭代中，计算梯度并保存。对于分类任务，设置的最大梯度累积步数为3，检测任务为10，而分割任务为2。

1、图像分类

针对提出的分类算法进行实验，使用的数据集是前列腺癌分级评估（Prostate cANcer graDeAssessment, PANDA）[36]和UltraMNIST [37]。训练轮数设置为100。
在这里插入图片描述

PANDA: 在PANDA数据集上（内存配置为24GB和16GB）训练所提算法。使用8,616个样本进行训练，1,000个样本用于验证，1,000个样本用于推理，并在表I中报告了推理结果。对于 $\theta_{1}$ 和 $\theta_{2}$ ，均采用了ResNet50 [2]。考虑了 $512 \times 512$ 、 $2048 \times 2048$ 和 $4096 \times 4096$ 三种不同的图像尺寸，以 $128 \times 128$ 或 $256 \times 256$ 的patch大小。在 $512 \times 512$ 的分辨率下，所提方法达到了59.3%的准确率，而PatchGD为53.4% （*表示作者重新实现的结果）。在极高分辨率 $4096 \times 4096$ 和 $256 \times 256$ 的patch大小下，所提方法比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$ ，inner iterations设置为3，采样率为25%。使用ResNet50（RN50）和MobileNetV2（MNV2）[38]作为两个网络，并在表II中报告结果。需要注意的是，Jetson Nano的GPU内存有限，使用 $512 \times 512$ 分辨率进行训练可能会不稳定，尤其是在使用SGD时。使用ResNet50的SGD训练得到约52.9%的准确率，而使用MobileNetV2则达到67.1%。相比之下，PatchGD在使用ResNet50时达到62.1%。在MobileNetV2上观察到1.2%的提升。同样，作者注意到峰值内存使用量有轻微增加，与PANDA中的观察结果类似。

2、目标检测

针对提出的检测算法，在COCO数据集[42]、驾驶场景下的小型物体检测数据集(SODA-D)[23]以及Pang等人 [43]的TJU-Ped-traffic上进行了实验。实验选择FCOS [32]作为 $\theta_{1}$ 。模型训练了300个epoch。使用两块Nvidia 2080ti GPU（每块11GB内存）进行训练。

在这里插入图片描述

COCO：遵循 [32]的协议。由于图像大小为 $640 \times 480$ ，创建了一个 $4096 \times 4096$ 的空白图像，并添加36张图像，用于在这些大型合成图像上训练和测试模型。还将新合成的图像调整为 $800 \times 1024$ ，与 [32]中的图像大小相同，然后训练FCOS算法。还在新合成的图像降采样到表III中提到的分辨率后，训练CenterNet、CornerNet和YOLOX。实验结果见表III。在AP上，所提模型达到25.4%，而FCOS为21.5%。在参数数量上，所提模型增加了约3.1%的轻微开销。由于基于patch的处理，延迟也有所增加。YOLOX的得分是23.3%。其他指标也显示出类似的趋势。结果表明，在给定的GPU内存下，所提训练方法更有效，因为它能够捕捉细粒度和全局细节。

在峰值内存使用上，所提方法使用20GB，YOLOX使用17GB，而FCOS使用15GB。尽管比FCOS多消耗5GB，但可以看到，所提方法允许在处理大型图像的同时提高性能。

在这里插入图片描述

Predicted Bounding Boxes：使用提出的物体检测算法在新合成图像上预测的边界框如图3(a)所示。我们还给出了使用FCOS（图3(b)）和YOLOX（图3©）预测的边界框。可以推断，当图像大小较大时，FCOS未能检测到许多小到中等大小的物体，而YOLOX表现良好，但未能超越所提模型。

在这里插入图片描述

SODA-D: 遵循 [23]中的协议。平均图像大小为 $3407 \times 2470$ 。将图像调整为 $4096 \times 4096$ 像素时，所提模型在AP上达到28.6%，相较于FCOS的23.9%和YOLOX的26.7%有所提升。在AP50和AP75上，所提方法也优于其他方法。值得注意的是，[23]在 $1200 \times 1200$ 分辨率的图像块上运行检测器，然后将检测结果映射回原始图像，再应用非极大值抑制消除重叠和冗余预测。这是一种在使用GPU内存和处理高分辨率图像之间取得有效平衡的方法。所提方法使用 $4096 \times 4096$ 的全尺寸图像，采用 $512 \times 512$ 的patch大小，峰值内存使用量为18GB。而[23]使用4块NVIDIA GeForce RTX 3090 GPU进行模型训练。从表IV中可以推断，所提方法在内存限制下处理高分辨率图像时，取得了更好的平衡。预测的边界框结果如图4所示。
在这里插入图片描述

TJU-Ped-traffic：遵循 [43]的协议。使用13,858个样本进行训练，2,136个样本进行验证，4,344个样本进行推理。该数据集有52%的图像分辨率为 $1624 \times 1200$ ，48%的图像分辨率为至少 $2560 \times 1440$ 。在表V中展示了结果。所提方法获得AP为29.45，YOLOX为27.2%，而FCOS为23.3%。其他指标也显示出类似的趋势。实验还发现，所提方法的峰值内存为20GB，相比之下，YOLOX为18GB，FCOS为14GB。这表明提出的这种方法可以在内存限制下处理全分辨率图像，并且性能更优。

3、分割

实验使用U-net [44] 和 DeepLab v3 [45] 作为 $\theta_{1}$ 。模型在3090Ti GPU （24GB内存）和Jetson Nano上训练100个epoch。

DRIVE - 视网膜血管：DRIVE [46] 数据集共有100张图像，每张分辨率为 $512 \times 512$ 。按照 [47]的协议，将其分为80张用于训练，20张用于测试。在表VI 中，U-net和DeepLab v3（Type Full）在 $512 \times 512$ 全尺寸图像上进行训练。在推理时，也使用全尺寸输入。
在这里插入图片描述

U-net和DeepLab v3 （Type Downsampled，简称Down.） 使用 $128 \times 128$ 的图像进行训练，推理时输入调整为 $128 \times 128$ ，输出掩码上采样至 $512 \times 512$ 。Ours‡表示在提出的框架下，使用 $128 \times 128$ 的裁剪进行训练，但不包含全局区域。Ours表示使用 $128 \times 128$ 的裁剪以及全局区域进行训练。

Type Full版本使用全尺寸图像，因此占用更多GPU内存，这里作为参考。特别关注Type Down-sampled的比较，因为它仅占用1.39 GB的GPU内存，可以在4GB限制内进行训练。在F1、IoU、平衡精度、PLR和NLR方面，所提方法相对于U-net Down有显著提升。F1从61.0%提高到79.8%，IoU从43.9%提升到66.5%。同时，Ours的结果与Full相当或更好。实验还发现，全局区域的加入略微增加了GPU峰值使用，但整体性能有所提高。使用DeepLab v3作为 $\theta_{1}$ 时，得出类似结论。将在图5中展示定性结果。如图 5(f)所示，所提方法能很好地保持连续性和细节。
在这里插入图片描述

航空影像：该数据集包含72张图像，每张图像分辨率为 $1024 \times 1024$ 。使用52个样本进行训练，20个样本进行测试。观察到所提方法（Ours‡）和原始版本的Ours在准确性上优于下采样方法，尽管在其他指标上稍有落后。然而，在DeepLab v3中，所提方法在准确性、IoU、平衡精度、PLR和NLR上都优于下采样。实验还在Jetson Nano上进行训练，结果见表VIII。结果显示，对于DeepLab v3和Unet，只能对 $128 \times 128$ 的下采样图像进行训练，而所提方法（Ours‡）可以处理 $1024 \times 1024$ 分辨率的图像。与下采样版本相比，所有指标都有显著提升。
在这里插入图片描述