K400首次90%准确率！UniFormerV2开源：基于ViT的高效视频识别

©作者 | 黎昆昌 

单位 | 上海AI Lab 

研究方向 | 视频行为识别

TL;DR

论文标题：

UniFormerV2: Spatiotemporal Learning by Arming Image ViTs with Video UniFormer

论文链接：

https://arxiv.org/abs/2211.09552

代码链接：

https://github.com/OpenGVLab/UniFormerV2

简单介绍一下我们最近放出来的工作 UniFormerV2，方法受 UniFormer 的启发，设计了通用高效的时序建模模块，可以无缝插入到各种开源预训练图像 ViT 中，显著增强模型对时序信息的处理能力。为进一步提升主流 benchmark 上的性能，我们将 K400/K600/K700 进行数据清洗，得到更精简的 K710 数据集（0.66M 训练视频），在该数据集上训练，只需简单微调即可得到超过 1% 的性能提升。

基于 CLIP 提供的视觉编码器，我们的 UniFormerV2 最终在 8 个主流 benchmark 都取得了 SOTA 结果，包括场景相关数据集（短时的  K400/K600/K700 和 Moments in Time，以及长时的 ActivityNet 和 HACS），和时序相关数据集（Something-SomethingV1&V2）。 仅单个模型（354M 参数），我们最终在 K400 上首次取得了 90.0% 的准确率。

▲ 8个主流benchmark上的结果

代码、模型和相应的训练配置均已开源，应该 目前开源模型里，在这 8 个流行 benchmark 上性能最强的。 欢迎小伙伴们试用，有用的话可以随身点个 star~

Motivation

▲ motivation

去年在做 UniFormer  [1] 时，我们时常觉得实验的周期太长了。由于 UniFormer 是全新的框架，每次调整结构，我们都需要先经过一轮 ImageNet 图像预训练，之后将模型展开后，在 Kinetics 等视频数据集进行二次微调。根据结果反馈，再进行结构的改进。

更简洁的训练方式是直接在视频数据上从头预训练（如 SlowFast [2] 和 MViT [3]），但这种方式只适合于 Google/Meta 等大的研究机构，并且直接在视频数据预训练，但训练开销也更加大。

尽管这能更仔细地设计对视频任务友好的高效框架，在相同 FLOPs 的确远超同期基于 ViT 的一些列工作（如 TimeSformer [4] 、ViViT [5] 和 Mformer [6] ），但二次训练的代价高昂，也限制了我们去探索更大规模的模型。我们当时最大的模型为 UniFormer-B（50M 参数），输入 32 帧，需要使用 32 张 V100 在 Kinetics 上训练约 2 周。

那么有没有更好的结构设计与训练方案，一来能保证模型的性能足够好，另外又能节省训练开销，让一般的实验室算力，也能进行 2022 年的视频框架探索呢？

回想起 2019-2021 年，彼时 ViT 还未提出，学术界设计了一系列基于 ResNet 的轻量视频模型，如 TSM [7] 和 STM [8] 等。这些方法设计即插即用的时序模块，在 ResNet50 和 ResNet101 进行了一系列充分的验证。由于 ResNet 网络较小，这些方法一直未能在视频数据集上取得突破性的性能提升，此时在 Kinetics 系列上最好的方法仍是 Meta 使用大数据集预训练的大模型 CSN [9] 和 SlowFast 等。

2021 年，ViT 的兴起也促使了 TimeSformer 和 ViViT 等工作的出现，这些工作进一步拓展实际了时间注意力，充分了利用开源的 ImageNet-21K 预训练 ViT，一度刷新 Kinetics 系列的性能上限。

然而注意力机制的开销仍不是一般实验室可以承受的，此时的模型的计算量已经以 TFLOPs 为计量单位。并且基于图像预训练 ViT 的方法，难以真正地提升时序建模能力，在 Something-Something 等强时序相关数据集上效果较差（甚至不如基于 ResNet 的轻量模型）。这也促使我们去探究卷积和注意力的特点，设计高效的视频框架 UniFormer。

时间来到 2022 年，更多样预训练的 ViT 模型开源，如有监督的 Deit III [10] 、对比学习的 CLIP [11] 和 DINO [12] ，以及掩码学习的 MAE [13] 和 BeiT [14] ，规模也日益增长。 我们完全可以充分利用上这些开源的预训练图像大模型，设计轻量的时序建模模块，以较小的训练代价迁移到视频任务中！

UniFormerV2 也便是在这种想法下产生，我们沿用了 UniFormerV1 中的结构设计思想，设计了高效的局部与全局时空学习模块，可无缝插入到强预训练的图像 ViT 中，实现强大的视频建模。

我们在一系列预训练以及不同规模的模型都进行了适配性实验，在流行的数据集上都能取得出众的性能。如上图右下角所示，我们的模型在 K400 上取得了最优的准确率和计算量权衡，并首次达到 90.0% 的 top-1 准确率。

Method

▲ 整体框架

3.1 整体框架

模型的整体框架如上图所示，主要包含四个主要改进，后面消融实验也会进行具体介绍：

• Temporal Downsampling：在 Patch Embedding 层进行时序下采样，该操作将 Patch Embedding 拓展为 3D 卷积，在时序上下采样，变相在相同计算量的前提下，可输入两倍的帧数，可明显提升模型对强时序相关行为的判别能力。

  
  

• Local UniBlock：在保留原始 ViT 对空间的全局建模的前提下，我们额外引入局部时间建模，该模块遵循 UniFormer 局部模块的设计，但 depth-wise 卷积仅在时间维度操作。在较小计算量的前提下提升骨架网络的时序建模能力。

  
  

• Global UniBlock：对于视频任务而言，对全局时空 token 学习 long-term dependencies 尤其关键，为压缩全局模块的计算量，我们引入交叉注意力，使用单个可学的 token 作为 query，不同层的输出 token 作为 keys 和 values，设计。并且我们引入 UniFormer 中的 DPE，增强 token 的时空位置信息。这样，每个全局模块会将对应层的时空信息压缩成单个 informative token。

  
  

• Multi-stage Fusion：为了将不同层的 informative token 融合，得到更复杂的表征，我们探索了包括串行、并行、层次化等几种融合方式，并最终采用最简单的串行设计。即前面层出来的 informative token 会作为下一层的 query token，不断迭代融合多层信息。最后，这些来自多层融合的时空表征会与骨架网络的输入表征进行加权融合，得到最终送入分类层的判别表征。

▲ 几种多层融合方式

3.2 实现细节

对于结构，经过消融实验我们发现：

• 对于 Kinetics 系列等场景相关视频，也即单图即可较好判断行为的前提下，模型只需要在深层额外插入 Global UniBlock，便可以满足这一大类数据的需求。

  
  

• 对于 Something-Something 这类强时序相关视频，对模型的时序建模能力要求极高，除了在 Patch Embedding 引入 temporal downsampling 之外，我们还在每层插入 Local UniBlock，在网络的中层和深层插入 Global UniBlock，方能实现强大的时序判别。

但饶是如此简单的结构，在 训练时需要一些超参的设计保证模型正常收敛 ，我们在论文正文和附录的实现细节中有所说明：

• 模型初始化：为了保证模型初始训练与原图像模型输出一致，我们对插入的部分模块进行了零初始化，包括 Local UniBlock 输出的 linear 层、可学 query、Global UniBlock 中 FFN 的输出 linear 层，以及加权注意力的可学权重。

  
  

• 训练超参数：对强预训练模型，在视频任务上进行 full-tuning 迁移，需要使用较小的学习率，并且在数据规模较小是，需要引入额外的正则化如 droppath 等，具体超参可看原文附录 Table11。

▲ 训练超参

3.3 Kinetics-710

为了进一步增强图像预训练模型的迁移性能，目前 SOTA 的工作 MTV [15] 引入了大规模视频文本数据 WTS-60M，进行二次多模态预训练。考虑到这种方式的训练成本（并且收集数据也太花钱了），我们产生只利用开源数据进行有监督预训练。

我们尝试了多种数据集联合训练的方式，如 COVER [16] 中的多源数据集联合训练，但在我们的实验中 K400+SSV2+MiT 的设置难以 work。于是我们干脆将同源数据联合，也即 K400、K600 和 K700，其中主要涉及两个重要步骤：

• 数据清洗：由于不同数据集之间，训练集和训练集，训练集和测试集，以及测试集和测试集之前存在重叠，包含信息泄露，我们首先根据 YouTube ID 对重叠视频进行删除，得到 0.66M 视频。

  
  

• 标签清洗：我们根据 YouTube ID 和 3 个数据集的类别名，进行标签清洗，最终保留了 710 类标签，论文附录的 Table20 提供了标签 map。

我们将这个精简的数据集称为 Kinetics-710（K710），对于二次有监督预训练，我们直接使用单个数据集上训练超参，模型仅输入 8 帧。预训练完成后，无论输入多少帧（8/16/32/64），该模型都仅需要在单个数据集 K400/K600/K700 上训练 5 个 epoch，即可实现超过 1% 的性能提升，同时极大地减小了训练开销。需要注意的是， 在 K710 预训练模型微调时，我们会对分类层权重进行映射 ，即根据标签 map 保留 400/600/700 类对应权重。

实验

4.1 Kinetics

▲ SOTA对比：Kinetics-400，Kinetics-600和Kinetics-700

除了计算量和准确率以外，我们额外提供了几个指标，包括图像预训练模型是否 ready（即不需要额外预训练），是否包含二次预训练，最终需要在单个数据集上训练多少 epoch。 综合考虑下，我们的 UniFormerV2 不仅在训练上高效，测试性能同样超越了以往的 SOTA 模型。

4.2 Moments in Time

▲ Moments in Time

对 MiT 这种类内类间差异更大的行为，UniFormerV2 同样取得了最佳的性能 ， 充分验证了我们方案的鲁棒性。

4.3 Something-Something

▲ Something-Something V2

在 Something-SomethingV2 这种强时序相关数据集上，UniFormerV2 远好于以往基于 ViT 设计的一系列方法 ，值得注意的是，这些方法往往在 Kinetics 上表现优秀，这也说明了他们实质上并没有真正的对时间进行建模。

与以往的 SOTA 方法相比，受助于强大图像预训练的优势，UniFormerV2 只需要在单个视频数据上训练少量 epoch，而以往方法往往需要先在 K400 上预训练，或者进行数千个 epoch 的自监督训练，才能达到最好的性能。

▲ Something-Something V1 和 K400预训练实验

在 Something-SomethingV1 数据集上，UniFormerV2 刷新了最先进的性能。如右表所示，我们也尝试过额外引入 K400 预训练，但 fine-tuning 效果反倒更差。可能的原因是，我们使用的模型经过 400M 图文数据预训练，而 K400 只包含 0.24M 训练视频，并且视频是场景相关的（即与单张图片效果类似），并且 K400 与 Sth-Sth 存在域差异，额外的 K400 训练可能会破坏 CLIP 预训练的特征分布。

4.4 ActivityNet和HACS

▲ ActivityNet和HACS

在长时行为数据集 ActivityNet 和 HACS 上，UniFormerV2 以显著的优势高于先前的 SOTA 方法。

4.5 不同预训练ViT

▲ 不同预训练

我们对有监督预训练、对比学习、掩码学习的预训练 ViT 都进行了适配验证，可以看到即便使用 ImageNet-1K 预训练的 ViT，我们的方法都能优于 TimeSformer，尤其是在 Sth-Sth 这类强时序相关数据集上。另外一个有趣的观察时， 不同的预训练方式并不能拉开太大的差距，但预训练的数据规模越大，在视频任务上的提升也越明显。

4.6 消融实验

▲ 消融实验

我们在 K400 和 Sth-SthV2 上进行了广泛的消融实验，充分验证了我们方法中提出的各个模块的有效性，以及 K710 二次预训练的高效性，感兴趣的小伙伴可以到正文看看进一步的分析。

4.7 可视化

可视化特征可以发现，UniFormerV2 能更好地在浅层和深层保留特征的细节，并关注前景区域。

结束语

在 UniFormerV2 中，我们探索设计了更高效更通用的时空建模模块，可以无缝适配图像预训练 ViT，显著增强对视频任务的处理效果。开源模型和 K710 的加持，让整个模型的训练谱系的训练十分高效。仅利用开源的 CLIP 预训练和开源的有监督数据，UniFormerV2 可在 8 个流行 benchmark 上超越以往 SOTA。

目前 UniFormerV2 的相关代码、模型和配置都已经开源，在这 8 个数据集上也是目前开源模型最强的。希望能为社区做出一点贡献，欢迎各位试用，帮忙找 bug！

参考文献

[1] UniFormer https://arxiv.org/abs/2201.04676

[2] SlowFasthttps://arxiv.org/abs/1812.03982

[3] MViThttps://arxiv.org/abs/2104.11227

[4] TimeSformerhttps://arxiv.org/abs/2102.05095

[5] ViViThttps://arxiv.org/abs/2103.15691

[6] Mformerhttps://arxiv.org/abs/2106.05392

[7] TSMhttps://arxiv.org/abs/1811.08383

[8] STMhttps://arxiv.org/abs/1908.02486

[9] CSNhttps://arxiv.org/abs/1904.02811

[10] DeiT IIIhttps://arxiv.org/abs/2204.07118

[11] CLIPhttps://openai.com/blog/clip/

[12] DINOhttps://arxiv.org/abs/2104.14294

[13] MAEhttps://arxiv.org/abs/2111.06377

[14] BeiThttps://arxiv.org/abs/2106.08254

[15] MTVhttps://arxiv.org/abs/2201.04288

[16] COVER https://arxiv.org/abs/2112.07175

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：hr@paperweekly.site 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·