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」
点击「关注」订阅我们的专栏吧
·
·
·