【论文精读】ViTPose

摘要

       验证了ViT在姿态估计方向的潜力。

框架

编码器

       给定person图像作为输入$X\in {\mathbb{R}}^{H\times W\times 3}$，首先通过一个patch embedding将图像嵌入到token $F\in {\mathbb{R}}^{\frac{H}{d}\times \frac{W}{d}\times c}$，其中$d$ (默认16)是块嵌入层的下采样比率，$C$是通道维度。嵌入的token再由几个transformer层处理，每个transformer层由一个多头自注意力(MHSA)层和一个前馈网络(FFN)组成：
$$\begin{gathered} F_{i+1}^{\prime }=F_i+MHSA(LN(F_i)) \\ {F}_{i+1}=F_{i+1}^{\prime }+FFN(LN(F_{i+1}^{\prime })) \end{gathered}$$

       其中$i$表示$i_{th}$ transformer层的输出，初始特征$F_0=PatchEmbed(X)$表示patch embedding层之后的特征。每个transformer层的空间和通道尺寸恒定，故骨干网的输出特征表示为$F_{out}\in {\mathbb{R}}^{\frac{H}{d}\times \frac{W}{d}\times c}$。

解码器

       采用两种轻量级解码器对从骨干网中提取的特征进行处理和关键点定位。第一种是经典的解码器，由两个反卷积块组成，每个反卷积块包含一个反卷积层，然后进行批量归一化和ReLU，每个块对特征图上采样2次，使用一个核大小为$1\times 1$的卷积层来获得关键点的定位热力图。
$$K=Con{v}_{1\times 1}(Deconv(Deconv(F_{out})))$$

       其中$K\in {\mathbb{R}}^{\frac{H}{4}\times \frac{W}{4}\times N_k}$表示估计的热力图(每个关键点一个)，$N_k$为要估计的关键点数量，对于MS COCO数据集，被设置为17。

       第二种直接使用双线性插值对编码器的特征图进行4次上采样，然后用一个ReLU和一个核大小为$3\times 3$的卷积层，获得热力图。
$$K=Con{v}_{3\times 3}(Bilinear(ReLU(F_{out})))$$

实验

实验配置

       ViTPose遵循自顶向下的人体姿态估计设置，检测器用于检测人体实例，ViTPose用于估计检测实例的关键点。使用ViT-B、ViT-L和ViT-H作为骨架，将相应的模型表示为ViTPose-B、ViTPose-L和ViTPose-H。骨架用MAE预训练权重进行初始化。使用$256\times 192$输入分辨率和AdamW优化器，学习率为5e-4。使用Udp进行后处理。模型训练了210次，在第170次和第200次时学习率下降了10次。最佳设置上图所示。

消融实验

       用经典解码器和简单解码器来训练ViTPose，以ResNet 作为骨干训练SimpleBaseline。可以观察到，对于ResNet-50和ResNet-152，使用简单解码器会导致大约18 AP下降。以视transformer为骨干的ViTPose与简单解码器无此想象，只有ViT-B、ViT-L和ViT-H的性能略有下降，对于指标AP50和AR50, ViTPose在使用两个解码器中的任何一个时获得了类似的性能，这表明普通视觉transformer具有很强的表示能力，不需要复杂的解码器。

       从表中还可以得出结论ViTPose的性能随着模型大小的增加而不断提高，说明ViTPose具有良好的可扩展性。

灵活性

预训练数据的灵活性

       上图为采用VitPose-B结构，在MS COCO val set上，四种预训练方案的结果。四种预训练方案具体为：

• 采用ImagNet预训练分类任务
• 采用MS COCO中的单个人体crop预训练MAE任务，将75%的patch随机的mask掉，然后让网络学习恢复这些patch
• 任务框架同方案2，不过数据集采用MS COCO + AI Challenger的单个人体crop
• 任务框架同方案2，不过数据集采用MS COCO + AI Challenger的全部数据

       这些模型分别在三个数据集上进行了1600个epoch的预训练，然后在MS COCO数据集上进行了210个epoch的姿态注释微调。可以看到使用MS COCO + AI Challenger(croping)，在只有一半数据量的情况下，可以达到比ImageNet更好的效果。

分辨率灵活性

       上图为不同模型输入分辨率的结果。采用不同的输入图像大小和下采样率$d$训练ViTPose，可以看到ViTPose-B的性能随着输入分辨率的增加而提高。

注意型灵活性

       上图为VitPose采用不同注意力模块组合实验。直接使用1/8特征大小的全注意力在MS COCO val集上获得了最好的77.4 AP(行一)，但即使在混合精度训练模式下也会存在较大的内存占用。故采用Shift和Pool分别表示Shift窗口和池窗口机制，窗口注意力可以缓解内存问题，但代价是由于缺乏全局上下文建模，从77.4 AP到66.4 AP(行二)，性能下降。移动窗口和池化窗口机制都促进了全局上下文建模的跨窗口信息交换，从而以不到10％的内存增加显著提高了10 AP的性能(行三、行四)。当两种机制同时使用时，性能进一步提升到76.8 AP(行五)。与ViTDet 中提出的联合使用全注意力和窗口注意力的策略(行六)相当，但内存占用更少。
       通过比较表中的第五行和最后一行，通过将窗口大小从$8\times 8$扩大到$16\times 12$，性能可以进一步从76.8 AP提高到77.1 AP，这也优于全部注意力和窗口注意力联合设置。

微调灵活性

       上图为三种微调ViTPose-B模型：完全微调、冻结MHSA模块和冻结FFN模块。冻结MHSA模块后(行二)，与完全微调(行一)相比，性能下降至75.1 AP。冻结FFN模块和仅微调MHSA模块时(行三)，有3.0 AP的显著下降，意味着视觉transformer的FFN模块更负责特定任务的建模。相比之下，MHSA模块更与任务无关。

任务灵活性

       上图为ViTPose整体模型使用MS COCO, AI Challenger 和MPII数据集进行多数据集训练测试，直接使用多数据集训练后的模型进行评估， decoder而没有在MS COCO上进一步微调。通过使用所有三个数据集进行训练，可以观察到ViTPose的性能从75.8 AP持续提高到77.1 AP。尽管与MS COCO和AI Challenger的组合相比，MPII的体积要小得多，但使用MPII进行训练仍然带来了0.1 AP的增加，这表明ViTPose可以很好地利用不同数据集中的多样化数据。

可转移性

       采用两种蒸馏方法测试算法的可转移性。方法一，给定一个教师网络T和学生网络S，采用简单的蒸馏方法，添加一个输出蒸馏损失$L_{t\to s}^{od}$，以迫使学生网络的输出模仿教师网络的输出。
$$L_{t\to s}^{od}=MSE(K_s,K_t)$$

       其中$K_s$和$K_t$是给定相同输入的学生和教师网络的输出。

       方法二为基于token的蒸馏方法来桥接教师和学生模型，是对方法一的补充。随机初始化一个额外的可学习知识标记$t$，并在教师模型的块嵌入层之后将其附加到视觉标记，冻结经过训练的教师模型，调整几个epoch的知识token获得知识：
$$t^{\ast }=\underset{t}{\mathrm{argmin}}(MSE(T(t;X),K_{gt})$$

       其中$K_{gt}$是真实热图，$X$是输入图像，$T(t;X)$表示教师的预测，$t^{\ast }$表示最小化损失的最佳标记。之后，在训练期间，冻结知识token $t^{\ast }$并与学生网络中的视觉token连接，以将知识从教师传递到学生网络。因此，学生网络的损失是：
$$\begin{gathered} L_{t\to s}^{td}=MSE(S(t^{\ast };X),K_{gt}) \\ {L}_{t\to s}^{tod}=MSE(S(t^{\ast };X),K_t)+MSE(S(t^{\ast };X),K_{gt}) \end{gathered}$$

       其中$L_{t\to s}^{td}$和$L_{t\to s}^{tod}$分别表示token蒸馏损失，输出蒸馏损失和token蒸馏损失的组合。

       上图为使用经典的输出蒸馏和知识token蒸馏将知识从ViTPose-l迁移到ViTPose-b。可以看出，基于token的蒸馏为ViTPose-B带来0.2 AP增益(行一)，而输出蒸馏带来0.5 AP(行二)增加。两种蒸馏方法具有互补性，综合使用得到76.6 AP(行三)，验证了ViTPose模型具有良好的可迁移性。

对比实验

       上图为对多种姿态估计算法在MS COCO val的test-dev集的结果对比。使用$256\times 192$的输入分辨率、单个A100 GPU，batch大小为64。

       可以观察到，ViTPose的模型尺寸较大，但它在吞吐量和精度之间获得了更好的权衡。ViTPose在更大的backbone上表现也很好。例如，ViTPose-L在val集上获得了比ViTPose-B更好的性能。ViTPose-L的表现超过了之前的SOTA CNN和transformer模型，包括UPD和TokenPose，并具有类似的推理速度。对比ViTPose-H(第15行)和HRFormer-B(第9行)，ViTPose-H在仅使用MS COCO数据进行训练就获得了更好的性能和更快的推理速度，分别有79.1 AP和75.6 AP以及241fps和158fps。

       使用ViTAE-G骨干构建了一个更强大的模型ViTPose-g，该模型具有1B参数，更大的输入分辨率($576\times 432$)，以及MS COCO和AI Challenger数据进行训练，探索了ViTPose的性能极限。具有ViTAE-G backbone的单个ViTPose模型在MS COCO测试开发集上的性能达到了80.9 AP，优于之前的所有SOTA方法，其中之前的最佳方法UDP++集成了17个模型，并在稍好的检测器下达到了80.8 AP。在集成三个模型后，ViTPose进一步取得了最好的81.1 AP。

       上图为在MS COCO数据集的测试效果。



       上图为在AI Challenger、OCHuman、MPII等具有遮挡、截断、尺度和人体外观的变化的数据集的测试效果。

reference

Xu, Y. , Zhang, J. , Zhang, Q. , & Tao, D. . (2022). Vitpose: simple vision transformer baselines for human pose estimation.