通过两阶段蒸馏实现高效全身姿态估计.

0. TL; DR

本文提出了一种名为 DWPose 的两阶段姿态蒸馏方法，用于提升全身姿态估计的效率和精度。通过利用教师模型的中间特征和最终输出对轻量级学生模型进行监督，并采用权重衰减策略，DWPose 显著提升了学生模型的性能。此外，DWPose 还探索了 UBody 数据集以解决数据稀缺问题，并在 COCO-WholeBody 数据集上取得了新的 SOTA 性能。

1. 背景介绍

全身姿态估计任务旨在定位图像中人体、手部、面部和脚部的关键点。这一任务面临诸多挑战，包括人体部件的多尺度性、低分辨率区域的细粒度定位以及数据稀缺性。现有的姿态估计模型在精度和效率上仍有待提升，尤其是在处理复杂场景和多样化姿态时。为了满足实际应用中对高效、准确姿态估计器的需求，本文提出了一种基于知识蒸馏的两阶段方法 DWPose，以提升轻量级模型的性能。

2. 方法介绍

DWPose 的核心是两阶段姿态蒸馏框架，具体如下：

（1）第一阶段蒸馏

⚪ 特征蒸馏

第一阶段蒸馏中，学生模型通过模仿教师模型的中间层特征来学习知识。具体来说，学生模型的特征$F_s$通过一个 1×1 卷积层调整维度后与教师模型的特征$F_t$计算均方误差（MSE）损失，公式如下：

\[L_{fea} = \frac{1}{CHW} \sum_{c=1}^C \sum_{h=1}^H \sum_{w=1}^W \|F^t_{c,h,w} - f(F^s_{c,h,w})\|^2\]

其中 $f$ 是 $1×1$ 卷积层，$H,W,C$ 分别表示教师特征图的高度、宽度和通道数。

⚪ Logit 蒸馏

对于 Logit 蒸馏，学生模型的最终输出 $S_i$ 通过教师模型的输出 $T_i$ 进行监督。与传统方法不同，DWPose 使用教师模型的完整输出（包括可见和不可见关键点）作为监督信号，损失函数如下：

\[L_{logit} = -\frac{1}{N} \sum_{n=1}^N \sum_{k=1}^K \sum_{i=1}^L T_i \log(S_i)\]

其中 $N$ 是批次中的人数，$K$ 是关键点数量，$L$ 是定位 bin 的长度。

⚪ 权重衰减策略

为了逐步减少蒸馏的权重，DWPose 引入了权重衰减策略，通过时间函数 $r(t)=1− \frac{t-1}{t_{max}}$实现，其中 $t$ 是当前迭代次数，$t_{max}$是总迭代次数。最终的第一阶段蒸馏损失为： \(L_{s1} = L_{ori} + r(t)\alpha L_{fea} + r(t)\beta L_{logit}\)

其中 $L_{ori}$ 是原始监督损失，α 和 β 是超参数。

（2）第二阶段蒸馏

第二阶段蒸馏是一种自蒸馏方法，学生模型通过自身的输出对头部进行微调。具体来说，学生模型的骨干网络被冻结，仅头部通过 Logit 蒸馏进行更新。损失函数为： \(L_{s2} = \gamma L_{logit}\)

其中 $γ$ 是超参数。这一阶段仅需 $20\%$ 的训练时间即可显著提升性能。

3. 实验分析

实验在 COCO-WholeBody 和 UBody 数据集上进行。COCO-WholeBody 包含 118K 训练图像和 5K 验证图像，用于全身姿态估计。UBody 数据集包含超过 1M 帧的真实场景数据，涵盖多样化的面部表情和手势。

DWPose 在 COCO-WholeBody 数据集上取得了显著的性能提升。例如，DWPose-l 在 256×192 输入分辨率下达到了 63.1% 的 AP，超越了其教师模型 RTMPose-x（65.3% AP）。DWPose-m 在 2.2 GFLOPs 的计算量下达到了 60.6% 的 AP，比基线模型高出 4.1%。此外，DWPose 还在 384×288 输入分辨率下进一步提升了性能，DWPose-l 达到了 66.5% 的 AP。

两阶段蒸馏均对性能提升有显著贡献。第一阶段蒸馏通过特征和 Logit 蒸馏显著提升了学生模型的性能，而第二阶段蒸馏进一步提升了性能，仅需 20% 的训练时间即可实现。例如，RTMPose-l 在第一阶段蒸馏后 AP 从 62.1% 提升到 62.9%，第二阶段蒸馏后进一步提升到 63.1%。