OREPA：阿里提出训练也很快的重参数策略，内存减半，速度加倍 | CVPR 2022

最新推荐文章于 2024-09-03 08:41:29 发布

晓飞的算法工程笔记

最新推荐文章于 2024-09-03 08:41:29 发布

阅读量805

点赞数 1

分类专栏：晓飞的算法工程笔记文章标签：深度学习人工智能机器学习神经网络计算机视觉

本文链接：https://blog.csdn.net/lichlee/article/details/125890034

版权

晓飞的算法工程笔记专栏收录该内容

211 篇文章 24 订阅

订阅专栏

论文提出了在线重参数方法OREPA，在训练阶段就能将复杂的结构重参数为单卷积层，从而降低大量训练的耗时。为了实现这一目标，论文用线性缩放层代替了训练时的BN层，保持了优化方向的多样性和特征表达能力。从实验结果来看，OREPA在各种任务上的准确率和效率都很不错

来源：晓飞的算法工程笔记公众号

论文: Online Convolutional Re-parameterization

论文地址：https://arxiv.org/abs/2204.00826
论文代码：https: //github.com/JUGGHM/OREPA_CVPR2022

Introduction

除了准确率外，模型的推理速度也是很重要的。为了获得部署友好且精度高的模型，近期很多研究提出基于结构重参数化来提高模型性能。用于结构重参数化的模型在训练阶段和推理阶段具有不同的结构，训练时使用复杂的结构来获得高精度，而训练后通过等效变换将一个复杂的结构压缩成能够快速推理的线性层。压缩后的模型通常具有简洁的架构，例如类似VGG的或类似ResNet的结构。从这个角度来看，重参数化策略可以在不引入额外推理时间成本的情况下提高模型性能。公众号之前发了RepVGG的论文解读《RepVGG：VGG，永远的神！ | 2021新文》，有兴趣可以去看看。

BN层是重参数模型中的关键组成部分，在每个卷积层之后添加一个BN层，如果图1b所示，移除BN层会导致严重的精度下降。在推理阶段，复杂的结构可以被压缩到单个卷积层中。而在训练阶段，由于BN层需要非线性地将特征图除以其标准差，只能单独计算每个分支。因此，存在大量中间计算操作（大FLOPS）和缓冲特征图（高内存使用），带来巨大的计算开销。更糟糕的是，高额的训练消耗阻碍了探索更复杂和可能更强大的重新参数结构。
为什么BN层对重参数化如此重要？根据实验和分析，论文发现BN层中的缩放因子能够使不同分支的优化方向多样化。基于这个发现，论文提出了在线重参数化方法OREPA，如图1c所示，包含两个步骤：

block linearization：去掉所有非线性泛化层，转而引入线性缩放层。线性缩放层不仅能与BN层一样使不同分支的优化方向多样化，还可以在训练时合并计算。
block squeezing：将复杂的线性结构简化为单个卷积层。

OREPA减少了由中间层引起的计算和存储开销，能够显著降低训练消耗(65%-75%显存节省、加速1.5-2.3倍)且对性能的影响很小，使得探索更复杂的重参数化结果成为可能。为了验证这一点，论文进一步提出了几个重新参数化的组件以获得更好的性能。
论文的贡献包含以下三点：

提出在线重参数化方法OREPA，能够大幅提高重参数化模型的训练效率，使得探索更强的重参数结构成为可能。
根据对重参数模型原理的分析，将BN层替换为线性缩放层，保持优化方向多样化特性和特征表达能力。
通过各种视觉任务实验表明，OREPA在准确率和训练效率方面都优于以前的重参数化模型。

Online Re-Parameterization

OREPA能够将训练期间的复杂结构简化为单一卷积层，维持准确率不变。OREPA的变换流程如图2所示，包含block linearization和block squeezing两个步骤。

Preliminaries: Normalization in Re-param

BN层是重参数中多层和多分支结构的关键结构，是重参数模型性能的基础。以DBB和RepVGG为例，去掉BN层后（改为多分支后统一进行BN操作）性能会有明显的下降，如表1所示。
比较意外的是，BN层的使用会带来过高的训练消耗。在推理阶段，重参数结构中的所有中间操作都是线性的，可以进行合并计算。而在训练阶段，由于BN层是非线性的（需要除以特征图的标准差），无法进行合并计算。无法合并就会导致中间操作需要单独计算，产生巨大的计算消耗和内存成本。而且，过高的成本也阻碍了更复杂的结构的探索。

Block Linearization

虽然BN层阻止了训练期间的合并计算，但由于准确率问题，仍然不能直接将其删除。为了解决这个问题，论文引入了channel-wise的线性缩放作为BN层的线性替换，通过可学习的向量进行特征图的缩放。线性缩放层具有BN层的类似效果，引导多分支向不同方向进行优化，这是重参数化性能的核心。

基于线性缩放层，对重参数化结构进行修改，如图3所示，以下三个步骤：

移除所有非线性层，即重参数化结构中的归一化层。
为了保持优化的多样性，在每个分支的末尾添加了一个缩放层，即BN层的线性替代。
为了稳定训练过程，在所有分支之后添加一个BN层。

经过block linearization操作后，重参数结构中就只存在线性层，这意味着可以在训练阶段合并结构中的所有组件。

Block Squeezing

Block squeezing将计算和内存过多的中间特征图上的操作转换为更快捷的单个卷积核核操作，这意味着在计算和内存方面将重参数的额外训练成本从 $O(H\times W)$ 减少到 $O(KH\times KW )$ ，其中 $(KH, K W)$ 是卷积核的形状。
一般来说，无论线性重参数结构多复杂，以下两个属性都始终成立：

重参数结构中的所有线性层（例如深度卷积、平均池化和建议的线性缩放）都可以用具有相应参数的卷积层来表示，具体证明可以看原文的附录。
重参数结构可表示为一组并行分支，每个分支包含一串卷积层。

有了上述两个属性，就以将多层（即顺序结构）和多分支（即并行结构）压缩为单个卷积，如图4a和图4b所示。原文有部分转换的公式证明，有兴趣的可以去看看原文对应章节，这块不影响对Block Squeezing的思想的理解。

Gradient Analysis on Multi-branch Topology

论文从梯度回传的角度对多分支与block linearization的作用进行了分析，里面包含了部分公式推导，有兴趣的可以去看看原文对应章节。这里总结主要的两个结论：

如果使用分支共享的block linearization，多分支的优化方向和幅度与单分支一样。
如果使用分支独立的block linearization，多分支的优化方向和幅度与单分支不同。

上面的结论表明了block linearization步骤的重要性。当去掉BN层后，缩放层能够保持优化方向的多样化，避免多分支退化为单分支。

Block Design

由于OREPA节省了大量训练消耗，为探索更复杂的训练结构提供了可能性。论文基于DBB设计了全新的重参数模块OREPA-ResNet，加入了以下组件：

Frequency prior filter：Fcanet指出池化层是频域滤波的一个特例，参考此工作加入1x1卷积+频域滤波分支。
Linear depthwise separable convolution：对深度可分离卷积进行少量修改，去掉中间的非线性激活以便在训练期间合并。
Re-parameterization for 1x1 convolution：之前的研究主要关注3×3卷积层的重参数而忽略了1×1卷积，但1x1卷积在bottleneck结构中十分重要。其次，论文添加了一个额外的1x1卷积+1x1卷积分支，对1x1卷积也进行重参数。
Linear deep stem：一般网络采用7x7卷积+3x3卷积作为stem，有的网络将其替换为堆叠的3个3x3卷积取得了不错的准确率。但论文认为这样的堆叠设计在开头的高分辨率特征图上的计算消耗非常高，为此将3个3x3卷积与论文提出的线性层一起压缩为单个7x7卷积层，能够大幅降低计算消耗并保存准确率。