开源！ViTAE模型再刷世界第一：COCO人体姿态估计新模型取得最高精度81.1AP

精度不错，模型B就354M，做标注工具感觉可以，工业落地还需要压缩。

人体姿态估计是计算机视觉领域的基本研究问题之一，具有很广泛的实际应用，例如医疗健康领域的行为分析、元宇宙领域的AIGC内容生成等。但是，由于人体姿态丰富，人物衣着表观变化多样，光照和遮挡等因素，人体姿态估计任务非常具有挑战性。之前的深度学习方法主要研究了新的骨干网络设计、多尺度特征融合、更强的解码器结构和损失函数设计等。

最近，探索研究院联合悉尼大学首次提出了一种基于vision transformer结构的简单基线模型ViTPose。该模型仅采用vision transformer结构作为编码器和少量反卷积层作为解码器，即可以在人体姿态估计任务上达到很好的性能。此外，我们进一步分析了模型规模、输入分辨率、特征图尺寸等因素对ViTPose在人体姿态估计任务上的影响，证明了ViTPose具有很好的扩展性和灵活性。采用探索研究院之前提出的具有10亿参数的ViTAE-G模型，ViTPose在人体姿态估计权威榜单MS COCO数据集上获得世界第一，达到了81.1AP的最高精度。

 

【链接：https://competitions.codalab.org/competitions/12061#results】

01

研究背景

近期，基于vision transformer的人体姿态估计算法开始在学术界流行起来。通过设计精巧的递归逐次解码结构，或者是采用多尺度并行的高分辨率特征融合，它们在人体姿态估计任务上取得了很好的结果。然而，简单的vision transformer结构是否就能够解决人体姿态估计任务尚缺乏研究。在本项研究中，我们从模型尺寸，输入分辨率，特征分辨率，预训练等多个方面探讨了简单的vision transformer在人体姿态估计任务上的表现。

02

方法概述

出于尽可能不使用精巧的设计并最大化探索vision transformer本身作用的考虑，我们直接将人体图像输入到vision transformer进行特征提取，并用简单的反卷积层对所得到的特征进行解码，最后对解码特征进行双线性上采样得到人体关键点热图。基于ViT-B骨干网络，这样一个简单的结构在MS COCO验证集上即可达到75.8 AP的检测精度，其性能与目前基于vision transformer的最强的人体姿态估计算法相当。当使用更大的分辨率和更大的模型尺寸，如表1和表2所示，我们的ViTPose模型可以取得更好的人体姿态估计效果。这体现了ViTPose良好的扩展性和灵活性。

表1. 基于ViT-B骨干网络的ViTPose模型在不同输入图像尺寸下的性能（COCO val set）

表2. ViTPose模型使用不同骨干网络ViT-B，ViT-L和ViT-H的性能（COCO val set）

在之前的工作中，ImageNet预训练对于vision transformer在下游任务上的泛化性能具有重要作用。但是，由于人体姿态估计数据集本身包含了大量的人体实例图像，我们怀疑是否有必要继续使用大型ImageNet数据集进行vision transformer的预训练。为此，我们分别测试了采用ImageNet数据（100万张分类图像）和仅仅使用人体姿态估计数据（50万张人体图像）进行MAE自监督预训练对人体姿态估计任务的性能影响。我们发现对于人体姿态估计任务，采用人体姿态估计数据进行预训练即可得到与采用更大规模的ImageNet数据集进行预训练相当的性能。

表3. 不同预训练数据对人体姿态估计任务的影响（COCO val set）

此外，之前的研究表明使用更大的特征分辨率，如1/4尺寸的特征，可以得到更为精细的人体关键点热图和更好的人体姿态估计结果。在这里，我们也探索了特征分辨率对ViTPose性能的影响。需要注意的是，与卷积神经网络不同，采用更大的特征分辨率会对vision transformer带来更多的显存占用和计算开销。因此，我们把vision transformer所有层的全局注意力运算改为窗口注意力运算，从而降低显存消耗和计算量。结果如表4所示，基于全窗口注意力的ViT-B作为骨干网络的ViTPose模型，由于缺乏窗口间的信息交互，使用更大分辨率特征（1/8）的性能（72.0AP），反而不如原始的ViTPose模型使用低分辨率特征（1/16）的性能（75.8AP）。为了解决这个问题，我们采用了探索研究院之前提出的ViTAE transformer作为骨干网络。ViTAE模型采用了并行的卷积分支和注意力分支结构，其卷积分支可以促进窗口间的信息交互。采用了ViTAE-B作为骨干网络，ViTPose在使用更大分辨率特征时取得了相比ViT-B更好的结果，也优于使用低分辨率特征的模型结果。

表4. ViTPose采用不同骨干网络ViT-B和ViTAE-B在更大特征分辨率下性能比较（COCO val set）

表5. ViTPose模型和其他模型的性能比较（COCO test-dev set）。其中，MSPN，RSN，和UDP++为2018，2019，2020年COCO竞赛的冠军方法。上标“+”表示采用了多个模型集成的结果。

03

性能比较

我们在COCO的测试集上对我们所提出的ViTPose方法和之前的代表性方法进行了比较，结果如表5所示。在不采用模型集成的情况下，使用ViTAE-G作为骨干网络的ViTPose已经取得了世界第一的结果，超过了采用了17个模型进行集成并使用了更强的目标检测器进行人体的检测的2020年冠军方案。采用4个ViTPose模型进行集成，我们的方法可以取得81.1AP的最佳性能。该结果充分说明了vision transformer模型在人体姿态估计方面的潜力和适应性。

04

结论

通过本项研究，我们探索并展示了简单的vision transformer模型可以用来解决人体姿态估计任务并取得了很好的性能。得益于更大的模型、输入分辨率、以及特征图尺寸，我们所提出的ViTPose模型在人体姿态估计领域的国际权威榜单MS COCO上取得了最高的81.1AP检测精度，刷新了世界记录。我们希望本项工作可以进一步激发社区力量，推进预训练vision transformer大模型在更多下游任务中的研究。

图像数据均来自公开数据集或公开视频，仅用于研究目的。

论文链接：https://arxiv.org/abs/2204.12484

代码链接：https://github.com/ViTAE-Transformer/ViTPose

参考文献

1. Xu Y, et al. ViTPose: Simple Vision Transformer Baselines for Human Pose Estimation. arXiv preprint arXiv:2204.12484, 2022.

2. Dosovitskiy A, et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. International Conference on Learning Representations, 2020.

3. Xu Y, et al. Vitae: Vision transformer advanced by exploring intrinsic inductive bias. Advances in Neural Information Processing Systems, 2021.

4. Zhang Q, et al. Vitaev2: Vision transformer advanced by exploring inductive bias for image recognition and beyond. arXiv preprint arXiv:2202.10108, 2022.

5. Li Y, et al. Exploring Plain Vision Transformer Backbones for Object Detection. arXiv