[论文解读] RevBiFPN: The Fully Reversible Bidirectional Feature Pyramid Network

Vitaliy Chiley, Vithursan Thangarasa+4 more|arXiv (Cornell University)|Jun 28, 2022

Advanced Neural Network Applications被引用 11

一句话总结

本文提出 RevBiFPN，一种完全可逆的双向特征金字塔网络，通过可逆计算消除激活值存储，实现高分辨率、多尺度网络的高效训练。其在图像分类、目标检测和实例分割任务上达到当前最优性能，训练显存消耗相比 EfficientNet 最多减少 19.8 倍，相比 HRNet 减少 2.4 倍，且在 MS COCO 上 mAP 提升最高达 2.5%。

ABSTRACT

This work introduces RevSilo, the first reversible bidirectional multi-scale feature fusion module. Like other reversible methods, RevSilo eliminates the need to store hidden activations by recomputing them. However, existing reversible methods do not apply to multi-scale feature fusion and are, therefore, not applicable to a large class of networks. Bidirectional multi-scale feature fusion promotes local and global coherence and has become a de facto design principle for networks targeting spatially sensitive tasks, e.g., HRNet (Sun et al., 2019a) and EfficientDet (Tan et al., 2020). These networks achieve state-of-the-art results across various computer vision tasks when paired with high-resolution inputs. However, training them requires substantial accelerator memory for saving large, multi-resolution activations. These memory requirements inherently cap the size of neural networks, limiting improvements that come from scale. Operating across resolution scales, RevSilo alleviates these issues. Stacking RevSilos, we create RevBiFPN, a fully reversible bidirectional feature pyramid network. RevBiFPN is competitive with networks such as EfficientNet while using up to 19.8x lesser training memory for image classification. When fine-tuned on MS COCO, RevBiFPN provides up to a 2.5% boost in AP over HRNet using fewer MACs and a 2.4x reduction in training-time memory.

研究动机与目标

• 解决使用多尺度特征融合的深度高分辨率网络训练过程中显存消耗过高的问题。
• 克服现有可逆方法无法支持双向多尺度特征融合的局限。
• 实现面向检测与分割等空间敏感任务的可扩展、显存高效的 SOTA 视觉模型训练。
• 设计一种可逆模块，支持前向与反向计算，无需存储中间激活值。
• 证明可逆多尺度融合在准确率与效率方面可超越标准 BiFPN 与 HRNet 架构。

提出的方法

• 提出 RevSilo，一种新颖的可逆双向多尺度特征融合模块，通过反向传播期间重新计算特征实现无激活值训练。
• 设计 RevSilo 以支持跨多个特征尺度的双向融合，同时保持局部与全局的空间一致性。
• 利用可逆计算将激活值显存复杂度从与网络深度线性增长降低至常数。
• 堆叠 RevSilo 模块构建 RevBiFPN，即 BiFPN 架构的完全可逆版本。
• 将 RevBiFPN 集成至标准视觉模型中，以 MBConv 作为主要构建模块，支持在显存受限硬件上的高效训练。
• 应用快速扩展（Fast Scaling）技术，高效扩展网络深度、宽度与输入分辨率，同时保持显存效率。

实验结果

研究问题

• RQ1 可逆计算能否有效扩展至视觉网络中的双向多尺度特征融合？
• RQ2 可逆多尺度融合在不损失模型准确率的前提下，能在多大程度上减少训练显存？
• RQ3 与 EfficientNet 和 HRNet 等 SOTA 模型相比，RevBiFPN 在准确率、MACs 和显存使用方面表现如何？
• RQ4 在显存受限条件下，RevBiFPN 是否能支持更大的批量大小与更高的输入分辨率，用于检测与分割任务？
• RQ5 多尺度特征融合中的可逆计算是否可泛化至分类任务之外的其他架构与任务？

主要发现

• RevBiFPN 在 ImageNet 上实现 84.3% 的 top-1 准确率，仅需 5.05 GB 训练显存，与 EfficientNet-B7 相同，但当扩展至相似容量时，显存消耗仅为基线的 19.8 倍。
• 在 MS COCO 目标检测任务中，RevBiFPN-S6 的 Bbox AP 比 HRNetV2p-W32 高 2.4%，Mask AP 高 2.0%，且训练显存减少 1.6 GB。
• 即使 HRNetV2p-W48 使用 2 倍训练周期，其性能仍不及使用 1 倍周期的 RevBiFPN-S6，证明 RevBiFPN 具备更优的显存效率。
• RevBiFPN-S2 在 COCO 上达到与 HRNetV2p-W18 相当的性能，但 MACs 减少 1.2 倍，显存减少 2.5 倍，凸显其效率优势。
• RevBiFPN 支持在检测与分割任务中使用更大的批量大小与更高的输入分辨率，提升训练稳定性和性能。
• 该方法可扩展至其他架构（如 ResNet、Transformer）与推理技术（如剪枝、量化），展现出广泛适用性。