CVPR 2022 |节省70%的显存，训练速度提高2倍

本篇分享 CVPR 2022 论文『Online Convolutional Re-parameterization』，浙大&阿里提出在线卷积重新参数化OREPA，节省70%的显存！训练速度提高2倍！代码已开源！

详细信息如下：

• 论文地址：https://arxiv.org/abs/2204.00826

• 代码地址：https://github.com/JUGGHM/OREPA_CVPR2022

      01      

摘要

结构重新参数化在各种计算机视觉任务中引起了越来越多的关注。它旨在在不引入任何推理时间成本的情况下提高深度模型的性能。尽管在推理过程中效率很高，但此类模型严重依赖复杂的训练时的block来实现高精度，从而导致大量额外的训练成本。

在本文中，作者提出了在线卷积重新参数化（OREPA），这是一个两阶段的pipeline，旨在通过将复杂的训练时的block压缩到单个卷积中来减少巨大的训练开销。为了实现这个目标，作者引入了一个线性缩放层来更好地优化在线块。在降低训练成本的帮助下，作者还探索了一些更有效的重新参数组件。

与最先进的重新参数模型相比，OREPA 能够节省约 70% 的训练时的显存成本，并将训练速度提高约 2 倍。同时，配备了 OREPA，这些模型在 ImageNet 上的性能比之前的方法高出 +0.6%。作者还对目标检测和语义分割进行了实验，并在下游任务上表现出一致的改进。

      01      

Motivation

卷积神经网络 (CNN) 已经见证了许多计算机视觉任务的成功，包括分类、目标检测、分割等。精度和模型效率之间的差异已被广泛讨论。一般来说，精度更高的模型通常需要更复杂的块，更宽或更深的结构。但是，这样的模型总是太重而无法部署，尤其是在硬件性能不佳的场景中，考虑到效率，更小、更紧凑和更快的模型是首选。

为了获得部署友好的模型并保持高精度，提出了基于结构重新参数化的方法，以“免费”提高性能。在这种方法中，模型在训练阶段和推理阶段具有不同的结构。具体来说，他们使用复杂的训练阶段拓扑（重新参数化块）来提高性能。训练后，他们通过等价变换将一个复杂的块压缩成一个单一的线性层。压缩模型通常具有简洁的架构，例如，通常是 VGG-like或 ResNet-like结构。从这个角度来看，重新参数化策略可以在不引入额外推理时成本的情况下提高模型性能。

归一化层被认为是重新参数模型中的关键组成部分。在重新参数块（上图（b））中，总是在每个计算层之后添加一个归一化层。据观察，去除此类归一化层将导致严重的性能下降。然而，在考虑效率时，这种归一化层的使用出乎意料地在训练阶段带来了巨大的计算开销。在推理阶段，复杂的块可以被压缩到单个卷积层中。但是，在训练期间，归一化层是非线性的，即它们将特征图除以其标准差，这阻止了合并整个块。结果，存在大量中间计算操作（大 FLOPS）和缓冲特征图（高显存使用）。

更糟糕的是，高额的训练预算使得探索更复杂和可能更强大的重新参数块变得困难。自然而然会产生下面的问题：为什么归一化层在重新参数中很重要？根据分析和实验，作者认为归一化层中的缩放因子最重要，因为它们能够使不同分支的优化方向多样化。

基于观察，作者提出了在线重新参数化（OREPA）（上图（c）），这是一个两阶段的pipeline，使模型能够简化复杂的训练时重新参数块。在第一阶段，块线性化，作者移除所有非线性归一化层并引入线性缩放层。这些层与归一化层具有相似的属性，它们使不同分支的优化多样化。此外，这些层是线性的，可以在训练时合并成卷积层。

第二阶段，称为块压缩，将复杂的线性块简化为单个卷积层。OREPA 通过减少由中间计算层引起的计算和存储开销，显着降低了训练成本，而对性能的影响很小。此外，高效率使得探索更复杂的重新参数化拓扑成为可能。为了验证这一点，作者进一步提出了几个重新参数化的组件以获得更好的性能。

作者在 ImageNet分类任务上评估提议的 OREPA。与最先进的重参数模型相比，OREPA 将额外的训练时 GPU 显存成本降低了 65% 至 75%，并将训练过程加快了 1.5 倍至 2.3 倍。同时，本文的 OREPA-ResNet 和 OREPA-VGG 始终优于以前的方法+0.2%∼+0.6%。作者在下游任务上评估 OREPA，即目标检测和语义分割。作者发现 OREPA 可以持续为这些任务带来性能提升。

本文的贡献可以总结如下：

1. 提出了在线卷积重参数化（OREPA）策略，极大地提高了重参数化模型的训练效率，并使得探索更强的重参数块成为可能。
2. 根据对重新参数模型工作机制的分析，作者用引入的线性缩放层替换了归一化层，这仍然提供了多种优化方向并保留了表示能力。
3. 各种视觉任务的实验表明，OREPA 在准确性和训练效率方面都优于以前的重新参数模型。

      03      

方法

上图展示了本文提出OREPA的pipeline。

3.1. Preliminaries: Normalization in Re-param

中间归一化层是重新参数化中多层和多分支结构的关键组成部分。以 SoTA 模型，即 DBB和 RepVGG为例，去除这些层会导致严重的性能下降，如上表所示。因此，作者认为中间归一化层对于重新参数化模型的性能至关重要。

然而，中间归一化层的使用出人意料地带来了更高的训练预算。在推理阶段，重新参数化块中的所有中间操作都是线性的，因此可以合并到一个卷积层中，从而得到一个简单的结构。但是在训练期间，归一化层是非线性的，即它们将特征图除以其标准差。结果，中间操作应该单独计算，这导致更高的计算和显存成本。更糟糕的是，如此高的成本会阻止探索更强大的训练块。

3.2. Block Linearization

中间归一化层阻止模型在训练期间合并单独的层。但是，由于性能问题，直接删除它们并非易事。为了解决这个难题，作者引入了通道线性缩放操作作为归一化的线性替代方案。缩放层包含一个可学习的向量，它在通道维度上缩放特征图。线性缩放层具有与归一化层类似的效果，它们都鼓励多分支向不同方向进行优化，这是重新参数化性能提升的关键。除了对性能的影响之外，线性缩放层可以在训练期间合并，使在线重新参数化成为可能。

基于线性缩放层，作者修改了重新参数化块，如上图所示。具体来说，块线性化阶段包括以下三个步骤。首先，移除所有非线性层，即重新参数化块中的归一化层。其次，为了保持优化的多样性，作者在每个分支的末尾添加了一个缩放层，即归一化的线性替代。最后，为了稳定训练过程，在添加所有分支之后立即添加一个后添加归一化层。

一旦完成线性化阶段，re-param 块中只存在线性层，这意味着可以在训练阶段合并块中的所有组件。

3.3. Block Squeezing

受益于块线性化（第 3.2 节），可以获得了一个线性块。在本节中，作者描述将训练时线性块压缩到单个卷积核中的标准过程。块挤压步骤转换中间特征图上的操作到更经济的核上的操作。这意味着在计算和显存方面将重新参数的额外训练成本从减少到，其中是空间特征图和卷积核的形状。
通常，无论线性重新参数块多么复杂，以下两个属性始终成立：

1. 块中的所有线性层，例如深度卷积、平均池化和线性缩放，都可以用具有相应参数集的退化卷积层来表示。
2. 该块可以由一系列并行分支表示，每个分支由一系列卷积层组成。

有了上述两个属性，如果可以将多层（即顺序结构）和多分支（即并行结构）都简化为单个卷积，就可以压缩一个块。在接下来的部分中，作者将展示如何简化顺序结构（上图（a））和并行结构（上图（b））。
首先定义卷积的符号。令表示大小的 2d 卷积核的输入和输出通道数。表示输入和输出张量。作为一种常见的做法，在这里省略了偏差，卷积过程表示为：

Simplify a sequential structure

考虑一堆卷积层，表示为：

其中 满足。根据结合律，这些层可以通过以下方程对卷积核进行压缩：

其中是第 j 层的权重。表示端到端映射矩阵。

Simplify a parallel structure

并行结构的简化是微不足道的。根据卷积的线性，可以将多个分支合并为一个:

其中是第m个分支的权重，是统一的权重。值得注意的是，在合并不同大小的内核时，需要对齐不同内核的空间中心，例如，一个 1×1 的内核应该与一个 3×3 的内核的中心对齐。

Training overhead: from features to kernels

无论块多么复杂，它都必须由不超过多分支和多层的子拓扑构成。因此，可以根据上面的两个简化规则将其简化为一个。最后，可以获得一体化的端到端映射权重，并且在训练期间只卷积一次。根据上式，作者实际上将中间特征图上的操作（卷积、加法）转换为卷积核上的操作。结果，将重新参数块的额外训练成本从降低到。

3.4. Gradient Analysis on Multi-branch Topology

为了理解为什么块线性化步骤是可行的，即为什么缩放层很重要，作者对重新参数化的统一权重的优化进行了分析。结论是，对于去除了归一化层的分支，使用缩放层可以使它们的优化方向多样化，并防止它们退化。
为了简化符号，这里只取输出 Y 的一维。考虑一个 conv-scaling 序列（conv-norm 序列的简化版本）：

其中 ，是滑动窗口内的矢量化像素，，W 是对应于某个输出通道的卷积核，γ是缩放因子。假设所有参数通过随机梯度下降更新，映射更新为：

其中L是整个模型的损失函数，η是学习率。对于具有共享 γ 的多分支拓扑，即：

权重的优化和上式相同，即：

具有分支 γ 的多分支拓扑提供了这样的变化，例如：

权重的更新与上式的不同，即：
条件1： 所有分支中至少有两个处于活动状态，即，使得

。

条件2：每个活动分支的初始状态彼此不同

同时，当满足条件 2 时，多分支结构不会退化为单分支结构，既可以用于forwarding也可以用于backwarding。这揭示了以下命题，解释了为什么比例因子很重要。请注意，当每个分支的权重随机初始化，并且缩放因子初始化为 1 时，总是满足条件 1 和 2。
命题1： 单分支线性映射，当通过两层以上的多分支拓扑重新参数化部分或全部时，整个端到端权重矩阵将得到不同的优化。如果将一层映射重新参数化为最多一层多分支拓扑，则优化将保持不变。
到目前为止，已经扩展了关于重新参数化如何影响优化的讨论，从仅多层到包括多分支。实际上，当前所有有效的重新参数化拓扑都可以通过之前的工作或命题 1 进行验证。

3.5. Block Design

由于提出的 OREPA 大大节省了训练成本，因此能够探索更复杂的训练块。因此，作者设计了一种新的重新参数化模型，即OREPA-ResNet，通过线性化最先进的模型 DBB，并插入以下组件（如上图）。

Frequency prior filter

在之前的工作中，在块中使用了池化层。之前的工作说明池化层是频率滤波器的一个特例。为此，作者添加了一个 Conv1×1 - 频率滤波器分支。

Linear depthwise separable convolution

作者通过移除中间非线性激活层稍微修改深度可分离卷积，使其在训练期间合并成为可能。

Re-parameterization for 1×1 convolution

以前的工作主要关注 3×3 卷积层的重参数，而忽略了 1×1 的。作者提出重新参数化 1×1 层，因为它们在瓶颈结构中起着重要作用。具体来说，作者添加了一个额外的 Conv1×1-Conv1×1分支。

Linear deep stem

大型卷积核通常放置在最开始的层，例如 7×7 stem层，旨在实现更大的感受野。之前也有工作通过堆叠的3×3卷积替换7×7 stem层。然而，由于高分辨率，最开始层的堆叠卷积需要更大的计算开销。可以使用本文提出的线性缩放层将堆叠的深层stem压缩到 7×7 的卷积层，这可以在保持高精度的同时大大降低训练成本。

对于 OREPA-ResNet（上图）中的每个块，1）作者添加一个频率先验滤波器和一个线性深度可分离卷积。2）作者用提出的线性深度stem替换所有stem层（即最初的 7×7 卷积）。3）在bottleneck块中，作者进一步用提出的Rep-1×1块替换原来的1×1卷积分支。

      04      

实验

作者在上图中对所有分支的分支相似度可视化，发现使用缩放层可以显着增加不同分支的多样性。

上表展示了本文方法中不同模块的消融结果，可以看出每个模块对于模型的性能提升都是有益的。

上表展示了与 ResNet18 模型上的线性化变体比较。

作者比较了 OREPA-ResNet-18 与其离线对应物（即 DBB-18）的训练成本。上图了消耗的显存（a）和训练时间（b）。随着组件数量的增加，离线重新参数化模型受到显存利用率快速增加和训练时间变长的影响。

上表展示了与其他重新参数化模型的比较，可以看出本文的方法相比于baseline有持续的性能提升。

如上表所示，OREPA 不断提高目标检测和语义分割任务的性能。

      05      

总结

在本文中，作者提出了在线卷积重新参数化（OREPA），这是一种两阶段pipeline，旨在通过将复杂的训练时的模块压缩到单个卷积中来减少巨大的训练开销。为了实现这一目标，作者用线性缩放层替换了训练时的非线性归一化，这保持了优化的多样性和表示能力的增强。
因此，本文显着减少了重新参数化模型的训练预算。这对于训练具有复杂拓扑的大规模神经网络至关重要，并且它进一步允许以更经济和有效的方式重新参数化模型。各种任务的结果证明了 OREPA 在准确性和效率方面的有效性。

参考资料

[1]https://arxiv.org/abs/2204.00826
[2]https://github.com/JUGGHM/OREPA_CVPR2022
▊ 作者简介 研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。 知乎/公众号：FightingCV