ICCVW 2023｜DWPose：两阶段蒸馏让 RTMPose 圆了 SOTA 梦

对于关注 RTMPose 和 ControlNet 的读者而言，最近两周或许都有了解到 DWPose 这篇工作，通过在 RTMPose 模型上进行二阶段蒸馏，成功取得 COCO-WholeBody 数据集 SOTA 的同时，推理速度快如闪电。并且，由于 DWPose 实现了无 mmcv 依赖的推理代码接入到 ControlNet 中，极大提升了人像生成的质量，尤其是手部的重建质量，成功两周涨了 800+ Star。

在 DWPose 代码开源的第一时间，我们也将模型参数加入到了 RTMPose 的项目页面方便下载，并且在 RTMPose 的在线 Demo 中进行了支持，还没有体验过的读者可以到下面的链接进行快速试玩：

https://openxlab.org.cn/apps/detail/mmpose/RTMPose

DWPose 开源链接：https://github.com/IDEA-Research/DWPose/tree/main

论文：https://arxiv.org/abs/2307.15880

结果一览

之前 RTMPose-x 以 0.1 AP 之差没能取得 SOTA 还是有一点点遗憾的，但因为考虑到本身 RTMPose 强调的也不是大参数量堆精度，所以我后来也没有再花时间进一步调参了，这次 DWPose 算是圆了 RTMPose 的 SOTA 梦~（笑）

RTMPose 模型经过两阶段蒸馏后，DWPose-l 模型以 66.5 AP 取得了 SOTA，超越了此前的 SOTA 方案 ZoomNAS (65.4 AP)，也高于装备了更大参数量 backbone 的 RTMPose-x （65.3 AP），可见这个方法确实很强。

而在速度方面，秉承着 RTMPose 的 Real-Time-Model 之名，相较于 ControlNet 中的标配 OpenPose，有着一百倍以上的速度提升：

在 3090 上，单人推理速度是 OpenPose 的 85 倍，而当画面中的人数来到 9 人时，速度则是 175 倍。

DWPose

如图所示，DWPose 提出的是一种两阶段的蒸馏算法，第一阶段设计了一个 weight decay 策略，同时对特征图和 logits 进行蒸馏，第二阶段则是对学生网络进行自蒸馏。

整体来看，第一阶段的训练跟标准的 RTMPose 训练相比，增加了一个教师网络提供特征图和 Logits 上的监督信号，第二阶段则只有 20% 的训练步长，因此该方法整体训练时长为标准 RTMPose 的 1.2 倍（当然，这是没有把更大的教师网络的训练时间算进去）。

第一阶段

特征蒸馏的部分，用的是经典的万金油 MSE，使用了一层 1x1 卷积把学生特征映射到跟教师网络相同通道数：

Logits 蒸馏的部分也是简单粗暴的 KL 散度，有了教师网络提供更多信息，很多没有标注的点也可以进行蒸馏：

权重衰减策略的部分，受到目标检测蒸馏方法 TADF 的启发，设计了一个随着训练逐渐降低蒸馏损失权重的策略：

第二阶段

自蒸馏的目的在于固定 Backbone 后，微调 Head 以调整特征分布，更好地学习 Backbone 输出的特征到预测目标之间的分布。通俗易懂一点地理解，同时改变 Backbone 和 Head 肯定不如一门心思调 Head 来得高效，我们可以认为第一阶段的学习已经得到了一个足够好的 backbone 了，因此第二阶段专心把 Head 调好。

消融实验

由于 DWPose 中使用了一个新的全身数据集 UBody，因此很需要重点关注的是消融实验部分，需要证明性能提升的来源真的是算法本身，而不是加数据带来的。这部分的实验作者做的非常好。

可以看到，在单一 COCO-Wholebody 上训练的 RTMPose-l 模型精度为 61.1 AP，增加了 UBody 联合训练可以带来一个点的提升，在此基础上再加蒸馏，可以再提一个点。

对于两阶段蒸馏的消融实验可以看到，同样在 RTMPose-l 上，第一阶段蒸馏可以带来 0.8 AP 的提升，单独的第二阶段自蒸馏则是 0.1 AP 的提升，而两个阶段联合则能一共带来 1.0 AP的提升。

自蒸馏部分可以看到，越小的网络上对 Head 做自蒸馏的收益越高，这个不难理解，因为越小的网络 Backbone 提取的特征质量越差，因此对于 Head 的能力要求越高，花费更多的时间去训练 Head 带来的收益也越高。反之，更大的网络 Backbone 提供的特征质量高，特征到目标的映射学习起来难度更低，所以微调 Head 的提升也有限。

由此我们不难推出一个调参的小技巧，对于较小的网络，可以适当增大 Head 参数量和层数，来强化 Head 的映射能力，而对于大的网络，Head 可以简单一点。

最后，第一阶段上的各种设计也做了详细的消融实验，以证明有效性。

后话

RTMPose 设计的初衷就是要拉近 MMPose 与实际业务场景的距离，以产品化和工程化的思维来解决姿态估计任务，比起一篇顶会论文或一种新算法，它更多对标的是 MediaPipe，OpenPose，PaddleDet 等产品。它从诞生之初就是以商业产品的要求被设计出来的，我希望能围绕 RTMPose 建立起一整套的 RTM 生态，把以人为中心的感知任务做得实时且高性能。

因此，我们非常高兴能看到 DWPose 这样优秀的工作出现，能对 ControlNet 为代表的 AIGC 领域带来更多帮助。同时 DWPose 也证明了 RTMPose 这个模型结构还有很大的潜力可以挖掘，我们后续会将 DWPose 的蒸馏策略也集成到 RTMPose 的训练中。

在 RTMPose 直播时我有提过后续的迭代计划，会有手部、脸部、全身、动物等各种任务的覆盖，算法上还会有 One-stage RTMPose，RTMPose3D 等一系列工作。目前人体、手部、脸部的模型已经放出了多数据混合训练的高性能权重，受到了来自业界公司的同学一致好评，也不乏互联网大厂的小伙伴在业务中开始使用 RTMPose。

最后，放一点小彩蛋。

这个季度的迭代计划主要是全身和 One-stage，DWPose 出来后促使我这周也终于挤出时间开始做全身姿态的优化了，目前初步的结果 RTMPose-l 在 384x288 上不用蒸馏已经到了 67.8 AP，推理速度也比之前能进一步加快，而且还远远没有到性能极限。后续等 DWPose 的 PR 合入后还可以再强强联手一波，初步先定一个 COCO-WholeBody 上 70 AP 的小目标吧，欢迎大家保持关注 MMPose 和 RTMPose，还没点 Star 的同学欢迎给 MMPose/RTMPose/DWPose 点个 Star，敬请期待~

公众号后台回复“数据集”获取100+深度学习各方向资源整理