比SOTA模型更全能！商汤科技和上海人工智能实验室联手打造统一模型架构UniFormer

UniFormer用Transformer将3D卷积和时空自注意力有效地统一，是目前最强最快的视频理解架构。

本文主要向大家介绍UniFormer，包括ICLR2022接受的video backbone，以及为下游密集预测任务设计的拓展版本。本套框架在各种任务上都能取得了比现有SOTA模型更好的性能。

目前UniFormer框架代码、模型、日志以及训练脚本都已开源，欢迎大家试用反馈！

✨UniFormer代码链接：

https://github.com/Sense-X/UniFormer

🪐论文标题：

UniFormer: Unified Transformer for Efficient Spatiotemporal Representation Learning

🪐论文链接：

https://arxiv.org/abs/2201.04676

🪐扩展版本：

UniFormer: Unifying Convolution and Self-attention for Visual Recognition

https://arxiv.org/abs/2201.09450

🪐代码链接：

https://github.com/Sense-X/UniFormer

“

相对SOTA，UniFormer的性能提升

1

图像分类

在Token Labeling 的加持下，仅靠ImageNet-1K训练。

39GFLOPs的UniFormer-L-384在ImageNet上实现了86.3%的top-1精度。

统一模型架构UniFormer

2

视频分类

仅用ImageNet-1K预训练，UniFormer-B在Kinetics-400和Kinetics-600上分别取得了82.9%和84.8% 的top-1精度（比使用JFT-300M预训练，相近性能的ViViT 的GFLOPs少16倍）。

在Something-Something V1和V2上分别取得60.9%和71.2%的top-1精度，为同期模型的SOTA。

统一模型架构UniFormer

3

密集预测

仅用ImageNet-1K预训练，COCO目标检测任务上取得了53.8的box AP与46.4的mask AP；ADE20K语义分割任务上取得了50.8的mIoU；COCO姿态估计任务上取得了77.4的AP。后文将会介绍为下游任务设计的训练和测试时模型适配。

点击可看高清大图

图注：

图像分类与视频分类任务性能比较（上方为ImageNet上224x224与384x384分辨率输入）

统一模型架构UniFormer

“

UniFormer设计灵感

对image和video上的representation learning而言，目前存在两大痛点：

• local redundancy:
  
  视觉数据在局部空间/时间/时空邻域具有相似性，这种局部性质容易引入大量低效的计算。

• global dependency:
  
  要实现准确的识别，需要动态地将不同区域中的目标关联，建模长时依赖。

现有的两大主流模型CNN和ViT，往往只关注解决以上部分问题。

Convolution只在局部小邻域聚合上下文，天然地避免了冗余的全局计算，但受限的感受也难以建模全局依赖。

而self-attention通过比较全局相似度，自然将长距离目标关联，但通过如下可视化我们可以发现，ViT在浅层编码局部特征十分低效。

图注：

DeiT可视化后，可以发现即便经过三层self-attention，输出特征仍保留了较多的局部细节。我们任选一个token作为query，可视化attention矩阵可以发现，被关注的token集中在3x3邻域中（红色越深表示关注越多）

图注：

TimeSformer可视化后，同样可以发现即便是经过三层self-attention，输出的每一帧特征仍保留了较多的局部细节。我们任选一个token作为query，可视化spatial attention和temporal attention矩阵都可以发现，被关注的token都只在局部邻域中（红色越深表示关注越多）。

无论是spatial attention抑或是temporal attention，在ViT的浅层，都仅会倾向于关注query token的邻近token。

但是attention矩阵是通过全局token相似度计算得到的，这无疑带来了大量不必要的计算。相较而言，convolution在提取这些浅层特征时，无论是在效果上还是计算量上都具有显著的优势。

那么为何不针对网络不同层特征的差异，设计不同的特征学习算子，从而将convolution和self-attention有机地结合，物尽其用呢？

本论文中提出的UniFormer (Unified Transformer)，旨在以Transformer的风格，有机地统一convolution和self-attention，发挥二者的优势，同时解决local redundancy和global dependency两大问题，从而实现高效的特征学习。

“

UniFormer模型架构

图注：

模型整体框架，标红维度仅对video输入作用，对image输入都可视作1

模型整体框架如上图所示，借鉴了CNN的层次化设计，每层包含多个Transformer风格的UniFormer block。

🪐UniFormer block：

每个UniFormer block主要由三部分组成，动态位置编码DPE、多头关系聚合器MHRA及Transformer必备的前馈层FFN，其中最关键的为多头关系聚合器：

🪐MHRA：

与多头注意力相似，将关系聚合器设计为多头风格，每个头单独处理一组channel的信息。

每组的channel先通过线性变换生成上下文token ，然后在token affinity 的作用下，对上下文进行有机聚合。

基于前面的可视化观察，可以认为在网络的浅层，token affinity应该仅关注局部邻域上下文，这与convolution的设计不谋而合。因此，将局部关系聚合 设计为可学的参数矩阵：

🪐local MHRA：

其中 为anchor token， 为局部邻域 任一token， 为可学参数矩阵， 为二者相对位置，表明token affinity的值只与相对位置有关。

这样local UniFormer block实际上与MobileNet block 的设计风格相似，都是PWConv-DWConv-PWConv（见原论文解析），不同的是我们引入了额外的位置编码以及前馈层，这种特别的结合形式有效地增强了token的特征表达。

在网络的深层，需要对整个特征空间建立长时关系，这与self-attention的思想一致，因此我们通过比较全局上下文相似度建立token affinity：

🪐global MHRA：

其中 为不同的线性变换。先前的video transformer，往往采用时空分离的注意力机制 ，以减少video输入带来的过量点积运算，但这种分离的操作无疑割裂了token的时空关联。

相反，我们的UniFormer在网络的浅层采用local MHRA节省了冗余计算量，使得网络在深层可以轻松使用联合时空注意力，从而可以得到更具辨别性的video特征表达。

再者，与以往ViT中使用绝对位置编码不同，我们这里采用卷积风格的动态位置编码，使得网络可以克服permutation-invariance的同时，对不同长度的输入更友好。

流行的ViT往往采用绝对或者相对位置编码 ，但绝对位置编码在面对更大分辨率的输入时，需要进行线性插值以及额外的参数微调，而相对位置编码对self-attention的形式进行了修改。

为了适配不同分辨率输入的需要，采用了最近流行的卷积位置编码设计动态位置编码:

🪐DPE：

其中DWConv为零填充的的深度可分离卷积。

一方面，卷积对任何输入形式都很友好，也很容易拓展到空间维度统一编码时空位置信息。

另一方面，深度可分离卷积十分轻量，额外的零填充可以帮助每个token确定自己的绝对位置。

“

UniFormer整体框架

🪐图像分类：

模型细节

我们设计了三种不同规模的模型，每个模型包含4层，前两层使用local MHRA，后两层使用global MHRA。

对于local MHRA，卷积核大小为5x5，归一化使用BN（使用LN性能较差）。

对于global MHRA，每个head的channel数为64，归一化使用LN。动态位置编码卷积核大小为3x3，FFN的拓展倍数为4。

对于特征下采样，我们采用非重叠卷积，其中第一次下采样卷积核大小为4x4、步长为4x4，其余三次下采样卷积核大小为2x2、步长为2x2。

在每次下采样卷积之后，额外增加LN归一化。网络最后接平均池化层与线性分类层，输出最终预测。当使用Token Labeling时，额外加入一个线性分类层以及辅助损失函数。

对于UniFormer-S，我们设计了增强版本，每层block数量为[3, 5, 9, 3]，下采样使用重叠卷积，FLOPs控制为4.2G，保证与其他SOTA模型可比。

🪐视频分类：

对于video使用的3D backbone，通过加载ImageNet-1K预训练的UniFormer-S和UniFormer-B，并进行卷积核展开。

具体来说，动态位置编码和local MHRA分别展开为3x3x3和5x5x5卷积。

对于下采样层，只在第一次下采样同时压缩时间和空间维度，而在其余三次下采样仅压缩空间维度，即第一次下采样卷积核大小为3x4x4、步长为2x4x4，其余卷积核大小为1x2x2、步长为1x2x2。这样就可以在减少计算量的同时，还能保证模型的高性能。

对于global MHRA，直接继承相应参数，将时空token序列化进行统一处理。

🪐密集预测：

对于下游密集预测任务，直接使用原始网络作为backbone并不合适。

因为这些任务往往输入大分辨率图像，比如目标检测中输入1333x800的图像，使用global MHRA会带来过多的计算量。

我们以UniFormer-S为例，统计了不同分辨率输入时，不同操作所需的计算量。

图注：

左边图片为模型整体计算量与第三/四层中MatMul运算所需计算量的结果；右边图片为第三层采用不同风格的MHRA所需的MatMul计算量的结果

从上图中可以看到，第三层中MHRA所需的MatMul运算随着分辨率的增加而急剧上升，在输入分辨率为1008x1008时，甚至占了总运算量50%以上，而第四层仅为第三层的1/28。因此，仅对第三层的MHRA进行改进。

受先前工作的启发，将global MHRA应用在限制的窗口内，这样会把原本的 复杂度降至 ，其中p为窗口大小。

然而直接应用纯window化操作，不可避免地会带来性能下降，为此我们将window和global操作结合。

每个hybrid分组中包含4个block，前3个为window block，最后1个为global block。UniFormer-S和UniFormer-B分别包含2个和5个分组。

不同任务在训练和测试时使用特定改进

如上图所示，在五种任务上，训练和测试采用特定改进。

对目标检测任务，由于训练和测试时输入分辨率都很大（如1333x800），在第三层都采用hybrid block。

对姿态估计任务，输入分辨率相对较小（如384x288），在第三层采用原本的gloabl block。

而对于语义分割任务，由于往往在测试时使用几倍于训练输入的分辨率（如2048x512 vs. 512x512），因此在训练时，对第三层采用global block。

而在测试时采用hybrid block，但需要注意测试时hybrid block中window size需要与训练时global block的感受野一致（如32x32），感受野不一致会导致急剧性能下降。

这种设计可以保证训练高效的同时，提高测试的性能。

“

实验结果

🪐图像分类：

我们用ImageNet-1K进行图像分类实验，采用了DeiT 的代码与训练策略，UniFormer-S/B/L使用droppath比例为0.1/0.3/0.4，对大模型额外加入Layer Scale防止网络训练崩溃 。对于Token Labeling，沿用了其代码训练框架与超参。

ImageNet-1K

结果如上所示，其中带*为加入Token Labaleing进行训练，UniFormer-B\dagger为前述UniFormer-B的增强版本。

可以看到UniFormer在不同规模下都取得了SOTA性能，UniFormer-L-384仅需39G的FLOPs，即可取得86.3%的top-1精度。

🪐视频分类：

我们在Kinetics-400/600以及Something-Something V1/V2上进行了视频分类实验，沿用了MViT的代码和训练策略。

对Kinetics采用dense采样方式，加载了ImageNet-1K的预训练模型加速训练，droppath比例保持与前述一致。

对Sth-Sth采用uniform采样方式，加载了Kinetics的预训练，droppath设置为前述两倍，并不进行水平翻转。

左：Kinetics

 右：Sth-Sth

结果如上图所示，仅使用ImageNet-1K预训练，我们在Kinetics上取得了与使用JFT-300M预训练的ViViT-L、使用ImageNet-21K预训练的Swin-B相近的性能，计算量大幅减小。

而在Sth-Sth上，UniFormer大幅高于先前CNN和ViT的结果，取得了新的SOTA结果，在Sth-Sth V1上为61.2%，V2上为71.4%。

🪐目标检测与实例分割：

我们在COCO2017上进行了目标检测和实例分割实验，沿用了mmdetection 的代码框架，配置了Mask R-CNN与Cascade Mask R-CNN框架，使用Swin Transformer的训练参数，均取得了SOTA性能。

COCO目标检测，Mask R-CNN

COCO目标检测，Cascade Mask R-CNN

🪐语义分割：

我们在ADE20K上进行了语义分割实验，沿用了mmsegmentation 的代码框架，配置了Semantic FPN与UperNet两种框架，分别使用了PVT和Swin的训练参数，均取得了SOTA性能。

ADE20K语义分割

左：Semantic FPN 右：UperNet

🪐人体姿态估计：

我们在COCO2017上进行了姿态估计实验，沿用了mmpose 的代码框架，配置了简单的Top-down框架，使用了HRFormer的训练参数，取得了SOTA性能。

COCO姿态估计，Top-down

🪐消融实验：

我们进行了详尽的消融实验，首先在图像和视频分类任务上验证了backbone的结构设计。其次，对video backbone的预训练、训练以及测试策略进行了探究。最后，验证了下游密集预测任务上改进的有效性。

🤖️

Model designs for image and video backbones

结构设计

FFN：

首先将local block替换为MobileNet block，其中ReLU替换为GeLU，expand ration设置为3保持计算量一致，保留动态编码以公平比较，可以发现local block在图像和视频分类任务上都会明显优于MobileNet block。由此可知transformer风格，以及其特有的FFN确实增强了token的特征表达。

DPE：

将动态位置编码去掉，在图像和视频任务上性能均下降，视频任务上更是掉了1.7%，由此可知位置编码有助于更好的时空特征学习。

Local MHRA size：

将local MHRA的卷积核大小设置为3、5、7、9，性能差异并不大，最终采用大小为5的卷积核，以取得最好的计算量与准确率的权衡。

MHRA configuration：

由纯local MHRA（LLLL）出发，逐层替换使用global MHRA。

结果可以发现，仅使用local MHRA时，计算量很小，但性能下降明显。逐层替换global MHRA后，性能逐渐提升。

但全部替换为global MHRA后，视频分类准确率急剧下降，计算量急剧上升，这主要是因为网络缺失了提取细节特征的能力，冗余的attention在有限的video数据下导致了急剧的过拟合。

🤖️

Pre-training, training and testing for video backbone

迁移性能、卷积核展开方式、训练采样方式

Transfer learning：

如上图所示，表11比较了不同结构的预训练性能以及小数据集上的迁移学习性能，可以发现，联合的时空学习方式，不仅在预训练模型上性能更好，在小数据集上的迁移学习性能提升明显。

而纯local MHRA以及时空分离的学习方式，迁移小数据训练未能带来提升。

Infalting methods：

如上图所示，表12中比较了是否对卷积核进行展开，可以发现，展开为3D卷积核，在场景相关的数据集Kinetics-400上性能接近，但在时序相关的数据集Sth-Sth V1上提升明显，尤其是在强的预训练的加持下，这表明3D卷积核对时空特征的学习能力更强。

Sampling strides of dense sampling：

如上图所示，表13中比较了我们在训练Kinetics时，使用不同间隔采样的结果。可以发现更稀疏的采样，在单clip测试上效果往往更好，但在多clip测试时，间隔4帧采样更好。

Sampling methods of Kinetics pre-trained model：

由于加载Kinetics预训练模型训练Sth-Sth，而Sth-Sth采用uniform采样，所以有必要知道预训练覆盖更多帧是否能带来提升。

如上图所示，表14的结果表明，预训练的不同采样方式差别并不大，16x4采样在大部分条件下性能都较好。

🤖️

左：不同数据集预训练
右：不同测试策略    

 

Pre-trained dataset scales：

如上图所示，图6比较了不同规模数据预训练的结果，可以发现对于小模型，大数据集预训练的提升非常明显，而对于大模型则相差无几。

Testing strategies：

如上图所示，图7比较了不同的测试策略，可以发现对于使用dense采样方式训练的场景相关数据集Kinetics而言，多clip测试方案较好，且1x4综合性能最优。对于使用uniform采样方式训练的时序相关数据集Sth-Sth而言，多crop测试方案较好，且3x1综合性能最好。

🤖️

Adaption designs for downstream tasks

下游改进

object detection：

如上图所示，表15比较了目标检测任务上，第三层采用不同类型block的结果。尽管在1x训练时，hybrid block的性能比纯global block性能略差，但经过3x的充分训练后，hybrid block的性能已经能和纯global block持平。

semantic segmentation：

如上图所示，表16比较了语义分割任务上，第三层采用不同类型block的结果。可以发现更大的窗口，以及global block的使用都能明显提升性能，由于纯global block计算量较大，我们采用性能相近的hybrid block。

pose estimation：

如上图所示，表17分别比较了姿态估计任务上，第三层采用不同类型block的结果。由于图像分辨率较小，zero padding消耗了更多的计算量。

🤖️

Visualizations

图像分类与视频分类

目标检测、语义分割、姿态估计

“

总结与思考

在过去一年多，我们在视频模型设计上尝试了：

CNN（CTNet，ICLR2021）

ViT（UniFormer，ICLR2022）

MLP（MorphMLP，arxiv）

我们发现：Transformer风格的模块+CNN的层次化架构+convolution的局部建模+DeiT强大的训练策略，保证了模型的下限不会太低。

但相比convolution以及linear layer而言，self-attention的性价比仍是最高的，同等计算量尤其是小模型下，self-attention带来的性能提升明显，并且对大数据集预训练的迁移效果更好。

✨CTNet链接：

https://github.com/Andy1621/CT-Net

✨UniFormer链接：

https://github.com/Sense-X/UniFormer

✨MorphMLP链接：

https://arxiv.org/abs/2111.12527

不过传统ViT对不同分辨率输入并不友好，并且对大分辨率输入的计算量难以承受，这在我们的工作UniFormer中都尽可能以简洁的方式解决，DWConv有限制地引入，也并不会导致过多的显存开销与速度延迟，最后在不同的任务上都能取得很好的性能权衡。

UniFormer提供了一个尽可能简单的框架，也希望后面的研究工作能在这个框架的基础上，去考虑视频中的运动信息、时空维度的冗余性、帧间的长时关系建模等等更复杂的问题，实现更大的突破。

商汤研究院基模型组火热招聘中

实习/全职均可

技术方向： 

1. 超大模型设计与优化，大规模神经网络理解； 

2. 基础研究与应用，包括但不限于：目标检测/识别/分割，知识蒸馏，基础模型结构设计等； 

3. 无监督/半监督训练; 

4. 大规模数据训练优化、通用表征学习； 

5. 长尾任务，开集类别检测等； 

6. 轻量化模型设计与优化； 

7. AutoML相关技术研发； 

团队优势： 

1. 1000+独有GPU( V100+A100)，5000+共享GPU；

2. 良好的研究氛围与技术指导，有足够的warm up周期；

3. 团队技术积累丰富，对解决具有挑战性的问题充满激情，获得多项著名竞赛冠军，如ImageNet2017、OpenImage2019、ActivityNet2020、NIST FRVT、MOT2016、MMIT2019, MFR等 。 

岗位要求： 

1. 熟练掌握机器学习（特别是深度学习）和计算机视觉的基本方法；

2. 具备以上一个或多个技术方向的研究经历，对该领域技术理解扎实；

3. 优秀的分析问题和解决问题的能力，对解决具有挑战性的问题充满激情，自我驱动力强；

4. 有较强的研究能力优先，如发表过第一作者CCF A类会议或期刊等论文。 

联系方式（简历投递，请注明实习或全职）： 

songguanglu@sensetime.com

🔍

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

进入知乎首页搜索「PaperWeekly」

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

·

·

·