[NeurIPS‘22] ViTPose: Simple Vision Transformer Baselines for Human Pose Estimation

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 simplicity of ViTPose

The scalability of ViTPose

The flexibility of ViTPose

The transferability of ViTPose

Experiment

---

摘要

• 针对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方法的比较