CVPR 2022 Oral｜节省80%GPU内存！SBP：针对视频任务训练的内存优化策略

↑ 点击蓝字关注极市平台作者丨Lart编辑丨极市平台

极市导读

本文提出了一种名为随机反向传播 (SBP) 的内存高效方法，用于在视频上训练深度神经网络。实验表明，SBP 可以应用于广泛的视频任务模型，最多可节省 80.0% 的 GPU 内存和 10% 的训练加速，而动作识别和时间动作检测的准确率下降不到 1%。>>加入极市CV技术交流群，走在计算机视觉的最前沿

论文：https://arxiv.org/abs/2203.16755

这篇文章是 CVPR 2022 Oral。本文设计了一种针对视频任务训练的内存优化策略。实际上，这一思路从某种角度可以看做是对18年论文Backdrop: Stochastic Backpropagation中所述的随机反向梯度的一种扩展和实用化。接下来会结合我个人的思考仔细介绍一些这一方法的动机和细节。

原始文档链接（内容更近于原始论文）：https://www.yuque.com/lart/papers/xu5t00

问题背景

视频任务中存在多种训练范式，这当然也与目标任务相关。例如我所知道的，视频目标分割任务中，许多方法会对不同的视频随机采样片段（clip）送入模型中。模型每次处理一个片段，而非整个视频连续送入模型中处理。另外也有一些方法会完全按照图像的训练方式，只是利用光流来关联时序信息。。而本文主要所讨论的是另一类，即那些视频理解类的任务，例如动作识别、时间动作定位以及在线动作检测等等。这些任务通常依赖于视频长期上下文信息。此时显存不足的问题尤其突出，这也为端到端训练的设计提供了巨大的阻碍。

这类视频任务中常见的模型可以归类为两种：树模型和图模型。

树模型：对于诸如在线动作检测和离线动作检测等长视频任务而言，大多数现有方法都是 Spatial-then-Temporal (StT) 模型，即他们大多基于一个独立的用于处理空间信息的CNN（2D或3D-CNN）来提取帧/片段的特征，之后通过一个额外的时间模型（RNN，Transformer 等）构建时间上的上下文关联。这类方法主要的特点是空间特征的提取过程是独立的。图模型：最近备受关注的Transformer就是一种典型的例子。在一些短期视频任务，例如行为识别中，他们虽然同样取得了良好的效果，但是其计算成本和内存占用同样也使得其应用收到了限制。这类方法主要的特点是各帧空间特征的提取过程是一个交织的过程。

因此本文主要关注于这两类方法的训练过程的改进，更具体来说，主要针对StT类方法和Video Transformer类方法的优化。

现有策略典型视频方法的显存优化策略

对此，现有方法主要有以下的一些策略：

减少模型参数的梯度：主要基于预先从原始视频中提取的特征，这可以看做是冻结了backbone的形式。DaoTAD 冻结了其 backbone 的前两个阶段。从缩减数据量上入手：压缩输入尺寸：在 AFSD 中使用 96×96，在 DaoTAD 中使用 112×112。

虽然这样的策略确保了方法AFSD[Learning salient boundary feature for anchor-free temporal action localization]和DaoTAD[Rgb stream is enough for temporal action detection]的端到端训练的可行性，但是从最终效果上来看，这也导致了准确性的下降。所以仍然需要寻找更加合适的解决方案。而这种方案又可以不用通过冻结backbone和缩小数据分辨率的方式来压缩显存占用。

以上这些都是从模型训练这一行为本身的视角切入的，但是从另一个角度来看，我们应该反思下，是否所有的视频数据的利用真的是必要的？从视频数据本身来看，其存在着高度的冗余性，这也是众多视频压缩算法的核心动机。如何充分的利用这些冗余性来压缩训练参数是一个值得探索的内容。文中提到了一些典型的工作，主要是通过压缩帧的数量来减少显存占用：

frame dropout: Vatt: Transformers for multimodal self-supervised learning from raw video, audio and textearly pool: No frame left behind: Full video action recognitionframe selections: Smart frame selection for action recognition

这些主动降采样输入帧序列，虽然可以减少由于海量数据带来的显存占用，但是也因此丢失了时域上精细信息的连续性。而且主动降采样输入也限制了他们，只能用于处理基于整体视频的视频理解类的任务，对于那些需要逐帧预测的任务就会面临巨大的压力。同时视频采样的处理方式也要求了视频内容移动速度不可以太快，否则关键内容非常容易被漏掉，影响预测效果。

而基于Transformer的方法则更多关注于对计算成本占比的大头Attention操作的重设计和优化，来强化其对于长序列任务的建模能力。

通用的显存压缩策略时间换空间：梯度检查点技术[Training deep nets with sublinear memory cost]：前向传播不再保存梯度信息，反向传播重新计算关联梯度，更多内容可见PyTorch 之 Checkpoint 机制解析。不足是会拖慢训练速度，基本不影响性能。梯度累加技术：即通过多次前向传播模拟单次大batch的效果，更多内容可见PyTorch中在反向传播前为什么要手动将梯度清零？（https://www.zhihu.com/question/303070254)。会拖慢训练速度，可能由于batch统计量的不同，而与真实大batch有所差异。其他策略：稀疏网络[Sparse networks from scratch: Faster training without losing performance]应用于图像识别模型，但只能在理论上节省内存。Sideways[Sideways: Depth-parallel training of video models, Gradient forward-propagation for large-scale temporal video modelling]中，每当有新激活可用时，就会覆盖原始激活，但这只能应用于因果模型。这也可以降低内存消耗。方案动机实验观察

作者们展示了在Kinetics-400上的视频Swin的激活图和梯度图。作者们假设完整的激活图（前向路径中的）对于提取重要语义信息是必要的，但对于梯度（反向路径）可能是不必要的。图上可以直观的看到模型特定层的激活图和梯度图中都是浅层相似，但是深层差异较大。

实际上如果从数据产生的先后关系来看，激活图是浅层影响深层，而梯度图则是深层影响浅层，因此也可以倒推出作者后续删除浅层结构梯度的一些动机。

相关方法的启发

这一点是被我单独提出来的，我觉得这些工作应该对本文的设计有着重要的启发，或者对于我们理解本文的动机有着重要的参考价值。

Regularizing metalearning via gradient dropout：这篇工作被作者提到了。它通过置零较小比例（10％~20％）的梯度来正则化训练过程。本文与之不同，本文的删除比例更大，而且主要用于压缩显存。Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text：该工作会使用轻量级网络从视频中采样帧，所以会同时删除对应帧的前向和反向路径。而本文仅是随机删除一定比例的反向传播路径，同时保持所有前向路径。完整的视频信息对于时间建模很重要。Backdrop: Stochastic Backpropagation：这篇18年的文章与本文工作同样非常相近，文中虽然没有提到，但是这里我觉得有必要指出一下。

Backdrop这篇工作里提出了对梯度添加drop操作从而提升对于数据利用的效率，减少过拟合。其核心操作和本文非常相似，即不修改前向传播，而仅仅随机删除部分反向传播的路径。同时这篇工作在最终的总结里有如下描述：

In our initial implementation of backdrop, the blocked gradients are still computed but are simply multiplied by zero. A more efficient implementation, where the gradients are not computed for the blocked paths, would therefore lead to a significant decrease in computation time as well as in memory requirements. In a similar fashion, we would expect backdrop to be a natural addition to gradient checkpointing schemes whose aim is to reduce memory requirements [Training deep nets with sublinear memory cost, Memory-efficient backpropagation through time].

Similar to other tools in a machine-learning toolset, backdrop masking should be used judiciously and after considerations of the structure of the problem. However, it can also be used in autoML schemes where a number of masking layers are inserted and the masking probabilities are then fine-tuned as hyperparameters.

可以看到，本文的工作可以看做是对Backdrop中期望的那些的内容在视频任务领域的一种扩展应用。(#^.^#)

具体方案

正如前面的分析，浅层结构的梯度所存在的冗余性，以及与深层特征的无关性，保证了删除浅层梯度从而压缩显存占用的可行性。因此作者们提出了一种名为随机反向传播（SBP）的用于压缩视频模型显存占用的方法。

具体而言，视频模型训练过程中，SBP 保持所有前向传播路径，但在每个训练步骤中随机和独立地删除每个网络层的反向传播路径。虽然这种操作导致模型中部分参数由于梯度链式法则被破坏而无法被更新，但是按照经验，只要保留整体计算图，仍然可以有效的更新网络的参数。

它通过减少对应于反向传播的部分缓存需求，从而降低 GPU 存储成本。而这个量可以通过可调节的保持率（keep-ratio）来控制。最终的节省量仅取决于保持率，并且其较低时，节省内存的程度也更显著。

实验表明，SBP 可以应用于视频任务的各种模型，在动作识别和时间动作检测任务上，可节省高达 80.0％的 GPU 存储，并且提升 10％的训练速度，而准确率下降了不到 1％。

实际应用

由于视频任务中常见的两类模型，即前文中提到的树模型和图模型，具体架构形式不同，这导致了SBP在应用时需要分开考虑。

首先说明一下文中会反复提到的“节点”的定义。

当模型定义时，其反向传播可以看做一个有向无环图（Directed Acyclic Graph (DAG)），其中顶点V对应着为梯度计算缓存的激活，而边E则是用于计算梯度的各个操作。因为模型可以看做不同层的堆叠，所以也可以按照层的方式来组织顶点中的元素，各个元素按照对应层序号来划分成多个。这里的 可以表示层 的输出，而 表示模型输入。每个可以看做是一系列特征向量的集合，例如，如果 是一个 CxDxHxW 的量，就可以把它看做是 DxHxW 个顶点的集合。这里每个顶点（论文中称为“node”），是一个 C 维的特征向量。

作者将将要分析的几类模型分为两个部分：1）空间模型（虚线以下）和 2）时间模型（虚线上方）。

深层结构可以看作是时间模型，因为这些层主要从所有帧中学习全局语义信息。浅层结构通常被视为主要学习空间和局部上下文信息的空间模型。

树模型的一个典型特点底层负责对各帧的空间建模（典型案例如使用固定的空间backbone提取帧特征），而顶层用于联合各帧信息进行时间建模。由于空间模型上的冗余性，实际处理中，SBP仅应用于移除浅层空间结构上的梯度。而且这样做也有好处，就是浅层结构中特征尺寸往往要比深层结构中的大得多，所以移除前者的梯度还可以避免跟踪相应中间节点，从而可以节省大量内存。因此，梯度仅在保留下来的反向路径的节点上传播。而且这类模型的特殊性也导致采样节点的梯度与其他节点是否被采样无关。

而图模型中节点之间存在强关联性，Vision Transformer 是一种典型的全连接图模型。图模型中会在层级别应用SBP，而非树模型中的空间模型级别。反向传播同样仅在被采样的节点之间进行。这些特性导致采样节点梯度计算会随着反向传播过程的进行越来越不精确。但是反向传播从上到下的冗余性却确保了这样的操作不会带来太大的问题。不过跨层维持采样节点不变也可以缓解这一问题。通过为所有层设置一个均匀的保持率可以简化参数的识别。但是可以制定一种系统的策略来设置层特定的保持率。

这里展示了Attention层中的使用形式。提出的方法会对所有层采样相同的节点（灰色节点）。

由于每个query实际上与整个key和value都有关系。所以为了准确计算query节点的梯度，这里会保留对所有key和value节点。这种保留相较于点乘得到的注意力权重的显存占用而言微不足道，后者才是大头。

为了对应于前文中描述的空间模型和时间模型的概念，作者们将Video Transformer中最深层的数层作为时间模型。这些部分有着完整的反向传播路径，而SBP仅作用于剩下的层中。这些层可以看作是空间模型。

具体实现

论文中给出了算法实现的伪代码，这里的任意操作f既包括 StT 模型中的空间模型也包括Video Transformer中的点乘注意力以及多层感知机。

前向传播期间，整个操作过程中自动微分引擎不会跟踪这些参数。通过给定采样索引idx只将需要计算梯度的输入被保存，用于反向传播。反向传播期间，对那些指定采样位置上的回传回来的梯度重新建立梯度路径。即对采样输入重新执行前向传播并回传梯度。

这样的实现可以充分利用自动微分引擎，所以效率很高，几乎没有开销。

空间和时间复杂度

这里引入了一个参数，保持率r。即相较于完整的梯度路径中的节点数量，执行SBP后仍旧维持梯度计算的节点数量的比例。

StT模型中的空间模型和时间模型是相互独立的，并且SBP仅仅作用于空间模型。由于这类模型中空间模型往往占用较大，所以执行SBP后相较于之前的占用可以近似为保持率r。Video Transformer模型中，在假定多头注意力的头的维度和token的数量分别为d和n时，需要计算梯度的激活图的占用，在使用SBP前后的比较如下式，其近似正比于采样率 r。因此r够小的话，显存占用可以显著下降。

关于时间压缩比例，假设前向传播和反向传播的时间是相等的，那么完整的比例就是(1+2r)/2，因为实际上我们执行了 1+r 倍的前向传播和 r 倍的反向传播。

实验效果相似性度量

实验中，作者们引入了Similarity of neural network representations revisited中用于测量两个神经网络之间表征相似性的CKA相似性来度量Video Swin-T的相同层激活的相似度。图中展示了三种训练设定下与完整的端到端训练时（也就是“Video Swin”）的相似性。曲线FD（frame dropout）展示了使用 1/4 数量帧训练的模型与使用所有帧训练的模型（即基线模型“Video Swin”）之间的 CKA 相似性。从另一个角度来看，FD方案相当于在所有层中去掉3/4的梯度。

从图中可以明显看到三种结构对应的浅层结构表征的相似性。据我推测，作者这里展示的图的曲线对应的应该是训练完整的模型的表征相似度。

可见，使用SBP和FD的模型中，浅层都具有较高的相似性，而深层结构的相似度都不是很高。作者认为随机截断反传路径后，浅层特征仍然保持着较高的相似性的现象，意味着丢弃一部分梯度并不会显着降低模型的性能。同时使用SBP的方法在深层结构中的CKA相似性高于使用FD，作者认为这也暗示了在深层结构保持完整的梯度计算的意义。

但是值得我们关注的是，这样的截断操作可以看到对于深层特征有着比较明显的影响，而这似乎并没有对性能造成巨大的差异？又或者说，是否正是因为这样的差异，导致了最终性能仍然稍弱于完整的训练呢？从另一个层面来说，这里似乎也反映出了前向传播与反向传播之间的紧密耦合，虽然在SBP中前向传播中被完整的保留，但是反向传播的调整仍然造成了与原始基线表征上的明显差异。可惜的是，作者们对于这些并未给出更多分析。

行为识别

实验中SBP仅对浅层的时间维度均匀随机采样。这样可以使与帧相对应的节点要么都被采样或被丢弃。在所有层中，采样节点保持一致。

对比中，SBP 在性能和显存占用的权衡下实现了最佳。在重度依赖时间信息的数据集（Epic-55）上FD性能的严重损失也表明了在训练和推理过程中保持所有前向传播路径的重要性。

在于梯度检查点技术的对比中可以看到：

单独使用时，二者内存消耗接近，但是SBP可以快约30%。一起使用时，相较于单独使用SBP，可以额外节省约40%的内存占用。

另外，按照文中所述，作者们也测试了SBP和O1级别的混合精度训练的搭配效果，发现这可以节省 Swin-T 超过约 50%的内存，而没有性能的下降。

实际上三者可以同时使用，但是作者对此并未提供更多信息。

前面的实验都是在时间维度上进行的。一帧的节点要么全部保留要么全部丢弃。作者们补充了在时间和空间上同时执行梯度丢弃的操作，使用了3D棋盘模式的时空梯度丢弃策略。在对比中，时间采样（即帧级别的处理）和时空采样（即 3D 棋盘形式）实现了相近的性能。这也反映了SBP用于空间采样以及被迁移到图像域的潜力。

时间行为检测

为了验证在长时间视频任务中的表现，作者们也在时间行为检测上进行了验证。

在线动作检测任务需要对实时视频流去检测正在发生在每一帧的动作，而不能利用未来信息。现存方法延续 StT 范式，基于提取的视频帧特征在时间模型上执行时序信息的建模。

在实验中可以看到：

端到端训练优于冻结 backbone 的策略，但是显存占用大得多。较低保持率 r=0.125 的 SBP的最终性能甚至略高于端到端训练。作者将其归因于SBP的正则化作用，可以缓解过拟合。LSTR 的改进大于 TRN 的改进，这表明更强的时间模型可以从端到端训练中受益。

针对SBP采样节点的方式，这里对比了三种策略：

Uniform Random Sampler：沿着时间维度均匀划分候选节点到多个 chunk 中。每个 chunk 的大小为 1/r。最终在每个 chunk 中随机采样一个节点。这也是默认的方式。Diverse-Feature Sampler：根据幅值来排序候选节点，然后执行均匀随机采样。Diverse-Grad Sampler：按照幅度排序候选节点梯度，然后执行均匀随机采样。

从对比中看到：

diverse-feature sampler 实现了整体最好的 mAP。三种采样器 mAP 确实相近，也反映出了方法对采样器的鲁棒性。

因为均匀采样实现了不错的 mAP，而且还更加简单，于是默认被用于其他所有实验。

作者还对空间模型的学习率进行了实验。对于使用SBP的StT模型：

对于所有的保持率，最佳的空间模型的学习率大约是时间模型的 0.1 倍左右。使用 SBP，空间模型确实对于学习率更加鲁棒一些，只要它们相对小一点。因此，如果未声明的话，默认设置空间模型的学习率是时间模型的 0.1 倍。时间动作定位

时间动作定位旨在检测未修剪视频中每个动作实例的时间边界。大多数现有方法遵循 StT 范式。

实验中，端到端训练性能优于冻结 backbone 的，但是显存占用也大得多。引入 SBP 后获得了更好的权衡。

局限性首先，SBP 是一个新的原型技术，这可以证明可以通过丢弃反向传播路径来节省内存，但是它也附带着精度的稍微下降。未来的探索可以关注于更加自适应的层保持率和采样方法，以降低准确性的损失，甚至由于梯度丢弃带来的的正则效果而提高准确性。其次，SBP 仅提升了少量的训练速度（~1.1×）。但是，如何将其与其他有效的模型训练策略（例如 Multigrid [A multigrid method for efficiently training video models]）集成，以加快训练的速度仍然有待探索。

公众号后台回复“CVPR 2022”获取论文合集打包下载～

极市干货数据集资源汇总：90+深度学习开源数据集整理｜包括目标检测、工业缺陷、图像分割等多个方向实操教程：Pytorch - 弹性训练极简实现( 附源码)｜PyTorch常用代码段合集CVPR 2022：CVPR"22 最新132篇论文分方向整理｜CVPR"22 最新106篇论文分方向整理｜一文看尽 CVPR 2022 最新 20 篇 Oral 论文

#极市平台签约作者#

Lart

知乎：人民艺术家

CSDN：有为少年

大连理工大学在读博士

研究领域：主要方向为图像分割，但多从事于二值图像分割的研究。也会关注其他领域，例如分类和检测等方向的发展。

作品精选

实践教程 | PyTorch中相对位置编码的理解实操教程 | 使用Docker为无网络环境搭建深度学习环境实践教程 | 一文让你把Docker用起来!投稿方式：添加小编微信Fengcall（微信号：fengcall19），备注：姓名-投稿△长按添加极市平台小编觉得有用麻烦给个在看啦~

关键词： 空间模型 时间模型 传播路径