paper link: https://arxiv.org/pdf/2204.12484.pdf
code link: https://github.com/ViTAE-Transformer/ViTPose
- 本文探索了纯transformer模型在人体关键点检测任务中的表现,验证了transformer模型的有效性、可拓展性和灵活的训练策略。但是模型参数量和训练代价(包括训练时间和硬件要求)限制了该方法的应用。
- 当训练集较大时(MS-COCO),基于训练集的预训练模型;
- 基于token的知识蒸馏;
- 重encoder轻decoder设计。共享encoder,在多个数据集上训练不同的decoder
目录
The transferability of ViTPose
摘要
- 针对Human Pose Estimation问题,提出一种纯transformer结构:ViTPose;
- ViTPose有四个优点:
- 模型结构简单(simplicity):纯transformer结构,无层级结构;
- 模型大小可拓展性强(scalability):从ViT-B到ViTAE-G;
- 训练策略灵活(flexibility):可基于不同注意力机制,输入分辨率,预训练策略,或finetune策略;
- 便捷的知识蒸馏(transferability):可以很方便的将知识从大ViTPose传递给小ViTPose;
- 在MS COCO上达到80.9 AP
方法
The simplicity of ViTPose
- ViTPose可以分为Encoder和Decoder。
- Encoder:首先将输入图片切块,随后求得tokens = PatchEmbed(image patches),其次,tokens送入L个堆叠的transformer block,每个transformer block包含一个Multi-head self-attention (MHSA)和feed-forward network (FFN)。最终:F_out = Encoder(image),其特征维度为(H/d, W/d, C)。
- Decoder:将Encoder的输出特征,解码为热图。Heatmap = Decoder(F_out),其特征维度为(H/4, W/4, N_k),其中N_k是关键点数量,Decoder()有两种选择,分别是基于Deconv和基于双线性插值,两种解码器效果相近。
The scalability of ViTPose
- ViTPose可以受益于vision transformers的发展,例如:ViT-B,ViT-L,ViT-H,和ViTAE-G。
The flexibility of ViTPose
- pre-training data flexibility:本文提出在训练集上通过MAE (masking 75%) 预训练ViTPose的backbone。
- Resolution flexibility:体现在输入图片和特征图分辨率上:1)可以通过resize提高输入图片分辨率;2)可以通过让图片分割时有重叠,提高特征图分辨率。
- Attention type flexibility:提出一种window-based attetion,有两个特点:1)shift window:在相邻窗口传播信息;2)pooling window:每个窗口对应token经过pooling后会送入其他窗口作为key和value token,求得cross-window特征。该注意力模块帮助减少内存和计算开销,并能提高模型表现。
- Finetuning flexibility:1)在冻结MHSA module参数的情况下finetune;2)对完整模型finetune。ViTPose在两者情况下的表现近似。
- Task flexibility:在共享backbone encoder的情况下,可以在多个数据集上训练,针对每个数据集设置特定的decoder。
The transferability of ViTPose
提出一种token-based蒸馏方法。具体来说,基于一个训练好的老师模型,冻住参数,并训练额外添加的learnable token。训练结束后,冻住该token参数,添加到学生模型中,参与学生模型的训练。
Experiment
- Implementation details:finetune需要210 epochs。
- backbone encoder越大效果越好。
- 训练集较大时,可以在训练集上预训练。
- 输入图片分辨率越大越大
- window attention可以一定程度上缓解训练内存和计算开销。
- 在多个数据集上训练可以带来效果提升
- Token-based知识蒸馏有效,可以和基于输出的知识蒸馏叠加使用。
- 在COCO上sota方法的比较