RTMO:迈向高性能的单阶段实时多人姿态估计

摘要

https://arxiv.org/pdf/2312.07526.pdf
实时多人姿态估计在平衡速度和精度方面提出了重大挑战。虽然两阶段自上而下的方法随着图像中人数增加而变慢，但现有的单阶段方法往往无法同时提供高精度和实时性能。本文介绍了RTMO，这是一种单阶段姿态估计框架，通过在YOLO架构中使用双1-D热图来表示关键点，实现了坐标分类的无缝集成，同时保持了高速度，并实现了与自上而下方法相当的精度。我们提出了一个动态坐标分类器和定制的热图学习损失函数，专门用于解决坐标分类与密集预测模型之间的不兼容性问题。RTMO在速度上比最先进的单阶段姿态估计器快约9倍，同时在COCO数据集上的准确率提高了1.1%。我们的最大模型RTMO-l在COCO val2017上达到了74.8%的AP，同时在单个V100 GPU上的运行速度为141 FPS，这证明了其效率和准确性。代码和模型可在https://github.com/openmmlab/mmpose/tree/dev-1.x/projects/rtmo上找到。

1、简介

多人姿态估计(MPPE)在计算机视觉领域至关重要，其应用范围从增强现实到体育分析。对于需要即时反馈的应用，如运动员定位的指导，实时处理尤为关键。虽然已经出现了许多实时姿态估计技术 [3,15,16,30] ，但如何在速度和精度之间取得平衡仍然是一个挑战。

当前的实时姿态估计方法分为两类：自上而下法 [3,15] 和单阶段法 [16,30]。自上而下方法使用预先训练的检测器在主体周围创建边界框，然后对每个人进行姿态估计。一个关键的局限性是它们的推理时间与图像中的人数成比例（参见图1）。另一方面，单阶段方法直接预测图像中所有人体的关键点位置。然而，当前的实时单阶段方法 [16,30,33] 在准确性方面落后于自上而下的方法（参见图1）。这些方法依赖于YOLO架构，直接回归关键点的坐标，这影响了性能，因为这种技术类似于对每个关键点使用Dirac delta分布，忽略了固有的歧义和不确定性 [20]。

另外，坐标分类方法使用双1-D热图，通过将关键点位置的概率分布在跨越整个图像的两个集合的bin上，增加了空间分辨率。这提供了更准确的预测，同时几乎没有额外的计算成本 [15,22]。然而，直接将坐标分类应用于像单阶段姿态估计这样的密集预测场景会导致bin利用率低下，因为图像和每个人只占据一个较小的区域。此外，传统的Kullback-Leibler散度（KLD）损失将所有样本视为相等，这对于网格中实例难度差异很大的单阶段姿态估计来说是不理想的。

因此，RTMO实现了与实时自上而下方法相当的精度，并超过了其他轻量级单阶段方法，如图1所示。此外，当处理图像中的多个实例时，RTMO在速度上超过了具有相似精度的自上而下方法。值得注意的是，RTMO-1模型在COCO val2017数据集上达到了74.8%的平均精度（AP），并在NVIDIA V100 GPU上表现出141帧每秒（FPS）的速度。在CrowdPose基准测试中 [18]，RTMO-1达到了73.2%的AP，创下单阶段方法的新纪录。这项工作的关键贡献包括：

• 一种创新的坐标分类技术，专门针对密集预测场景，利用坐标bins进行精确的关键点定位，同时解决了实例大小和复杂性的差异所带来的挑战。
• 一种新的实时单阶段MPPE方法，无缝地将坐标分类与YOLO架构相结合，在现有自上而下和单阶段MPPE方法中实现了性能和速度之间的最佳平衡。

2、相关工作

2.1. 单阶段姿态估计器

受单阶段目标检测算法的进步 [7, 9, 24, 40, 51] 的启发，一系列单阶段姿态估计方法已经出现 [10, 30, 34, 39, 51]。这些方法在单个前向传播中执行MPPE，并直接从预定的根位置回归实例特定的关键点。其他方法，如PETR [37] 和ED-Pose [46]，将姿态估计视为集合预测问题，为关键点回归建立了完全端到端的框架。除了基于回归的解决方案，还有像FCPose [31]、InsPose [35]和CID [42]这样的技术，利用动态卷积或注意力机制生成实例特定的热图来进行关键点定位。

与两阶段姿态估计方法相比，单阶段方法省去了预处理（例如自上而下方法的行人检测）和后处理（例如自下而上方法的键点分组）的需要。这带来了两个好处：1）一致的推理时间，与图像中的实例数量无关；2）简化的管道，便于部署和实际使用。尽管有这些优点，现有的单阶段方法在平衡高准确性与实时推理方面仍面临挑战。高准确度模型 [42,46] 通常依赖于资源密集的后备，如HRNet [38] 或Swin [25]，这使得实时估计变得具有挑战性。相反，实时模型 [30,33] 在性能上做出了妥协。我们的模型解决了这一权衡问题，同时提供了高准确度和快速的实时推理。

2.2. 坐标分类

SimCC [22] 和 RTMPose [15] 已经采用坐标分类进行姿态估计。这些方法基于水平和垂直轴上的亚像素bins对关键点进行分类，实现了空间区分，而无需依赖高分辨率特征图。这有效地平衡了准确性和速度。然而，对于密集预测方法，跨越整个图像的bins是不切实际的，因为需要大量的bins来减少量化误差，这导致许多bins对于单个实例来说是多余的，从而导致效率低下。DFL [20] 在每个锚点的预定范围内设置bins，可能无法覆盖大实例的关键点，并且对于小实例可能产生显著的量化误差。我们的方法根据每个实例的大小分配bins，将它们放置在局部区域内，优化了bin的使用，确保了关键点的覆盖范围，并最小化了量化误差。

2.3. Transformer增强的姿态估计

基于Transformer的架构在姿态估计中变得无处不在，利用最先进的Transformer后端来提高准确性，如ViTPose [45]，或结合CNN和Transformer编码器来捕捉空间关系 [47]。TokenPose [21] 和 Poseur [32] 证明了基于令牌的关键点嵌入在热图和回归方法中的有效性，利用视觉线索和解剖学约束。PETR [37] 和 ED-Pose [46] 等框架在端到端多人体姿态估计中引入了Transformers，而RTMPose [15] 将自注意力与基于SimCC的框架相结合 [22]，用于关键点依赖性分析。虽然位置编码在注意力中是标准的，用于指示查询和键的位置，但我们创新地使用它来形成每个空间bin的表示向量，使计算bin-keypoint相似性成为可能，从而促进准确的定位预测。

3. 方法论

在我们的模型中，我们采用类似YOLO的架构，如图2所示。主干是CSPDarknet [9]，我们使用Hybrid Encoder [28]处理最后三个特征图，产生具有16和32的相应下采样率的P_{4}和P_{5}空间特征。这些特征图中的每个像素映射到原始图像平面上均匀分布的网格单元格。网络头利用每个空间级别的双卷积块，为每个网格单元格生成分数和相应的姿态特征。这些姿态特征用于预测边界框、关键点坐标和可见性。通过动态坐标分类器生成1-D热图预测的细节在3.1节中给出，而基于MLE的热图损失在3.2节中提出。完整的训练和推理过程在3.3节中概述。
3.1. 动态坐标分类器

与每个网格单元格关联的姿态特征封装了关键点相对于网格的位移。之前的工作[10、30、34]直接回归这些位移，因此在性能上存在不足。我们的研究探索将坐标分类与单阶段姿态估计框架相结合，以提高关键点定位精度。现有坐标分类方法的明显限制是它们静态的bin分配策略。为了解决这个问题，我们引入了动态坐标分类器（DCC），它动态地分配范围并形成两个1-D热图中的bin表示，有效地解决了密集预测上下文中坐标分类的不兼容性。

动态bin分配。顶部的姿态估计器中使用的坐标分类技术将bins分配到整个输入图像上[15、22]。这种方法导致单阶段方法中的bin浪费，因为每个对象只占据图像的一小部分。DFL [20]在每个锚点附近的预定范围内设置bins，这可能会错过较大实例中的关键点，并在较小实例中导致显著的量化误差。DCC通过动态分配bins来解决这些限制，以确保与每个实例的边界框对齐，确保局部覆盖范围。首先使用逐点卷积层回归边界框，然后将其扩大1.25倍，以覆盖所有关键点，即使在预测不准确的情况下也是如此。这些扩展的边界框在水平和垂直轴上均匀地划分为B_{x}和B_{y}个bins。每个水平bin的x坐标计算如下：
$$x_i=x_l+\left(x_r-x_l\right)\frac{i-1}{B_x-1},$$
其中，$x_r$和$x_l$是边界框的左右两边，索引$i$从1变化到$B_x$。y轴的bins计算方式类似。

动态bin编码。在DCC的上下文中，由于每个网格的预测边界框不同，每个bin的位置在网格之间会有所变化。这与以前的方法[15, 22]不同，其中bin坐标是固定的。DCC不是像这些方法那样在网格之间使用共享的bin表示，而是实时生成定制的表示。具体来说，我们将每个bin的坐标编码为位置编码，以创建针对bin的特定表示。我们使用正弦位置编码[41]定义为：

$${\left[\mathbf{PE}\left(x_i\right)\right]}_c=\{\begin{array}{ll}\sin \left(\frac{x_i}{t^{c/C}}\right),& \text{ for even }c\\ \cos \left(\frac{x_i}{t^{(c-1)/C}}\right),& \text{ for odd }c\end{array},$$
其中，t表示温度，c是索引，C表示总维度数。我们使用全连接层来完善位置编码的适应性，该全连接层应用可学习的线性变换$\varphi$，从而优化其在DCC中的有效性。

The primary objective of DCC is to accurately predict keypoint occurrence probabilities at each bin, informed by bin coordinates and keypoint features. Keypoint features are extracted from the pose feature and refined via a Gated Attention Unit (GAU) module [13] following RTMPose [15], to enhance inter-keypoint consistency. The probability heatmap is generated by multiplying the keypoint features \boldsymbol{f}{k} with the positional encodings of each bin \boldsymbol{P E}\left(x{i}\right) , followed by a softmax:

\hat{p}{k}\left(x{i}\right)=\frac{e^{\boldsymbol{f}{k} \cdot \boldsymbol{\phi}\left(\boldsymbol{P E}\left(x{i}\right)\right)}}{\sum_{j=1}^{B_{x}} e^{\boldsymbol{f}{k} \cdot \boldsymbol{\phi}\left(\boldsymbol{P E}\left(x{j}\right)\right)}}

where \boldsymbol{f}_{k} is the k -th keypoint’s feature vector.
3.2. MLE for Coordinate Classification

In classification tasks, one-hot targets and cross-entropy loss are commonly utilized. Label smoothing, like Gaussian label smoothing used in SimCC [22] and RTMPose [15], along with KLD, can improve performance. The Gaussian mean \mu_{x}, \mu_{y} and variance \sigma^{2} are set to the annotated coordinates and a predefined parameter. The target distribution is defined as:

p_{k}\left(x_{i} \mid \mu_{x}\right)=\frac{1}{\sqrt{2 \pi} \sigma} e^{{-\frac{\left(x_{i}-\mu_{x}\right)}{2}}{2 \sigma^{2}}} \sim \mathcal{N}\left(x_{i} ; \mu_{x}, \sigma^{2}\right)
Importantly, we note that p_{k}\left(x_{i} \mid \mu_{x}\right) is mathematically identical to the likelihood p_{k}\left(\mu_{x} \mid x_{i}\right) of the annotation \mu_{x} under a Gaussian error model with true value x_{i} . This symmetrical property arises because the Gaussian distribution is symmetric with respect to its mean. Treating the predicted \hat{p}{k}\left(x{i}\right) as the prior of x_{i} , the annotation likelihood for the k -th keypoint is:

\begin{aligned}
P\left(\mu_{x}\right) & =\sum_{i=1}^{B_{x}} P\left(\mu_{x} \mid x_{i}\right) P\left(x_{i}\right) \
& =\sum_{i=1}^{B_{x}} \frac{1}{\sqrt{2 \pi} \sigma} e^{{-\frac{\left(x_{i}-\mu_{x}\right)}{2}}{2 \sigma^{2}}} \hat{p}{k}\left(x{i}\right) .
\end{aligned}

Maximizing this likelihood models the true distribution of the annotations.

In practice, we employ a Laplace distribution for P\left(\mu_{x} \mid x_{i}\right) and a negative log-likelihood loss:

\mathcal{L}{\text {mle }}^{(x)}=-\log \left[\sum{i=1}^{B_{x}} \frac{1}{\hat{\sigma}} e^{-\frac{\left|x_{i}-\mu_{x}\right|}{2 \hat{\sigma} s}} \hat{p}{k}\left(x{i}\right)\right],

3.3. Training and Inference

Training. Our model, adhering to a YOLO-like structure, employs dense grid prediction for human detection and pose estimation. It is crucial for the model to differentiate between positive and negative grids. We extend SimOTA [9] for training, assigning positive grids based on grid scores, bounding box regression and pose estimation accuracy. The head’s score branch classifies these grids, supervised by varifocal loss [50] \mathcal{L}_{c l s} , with target scores being the Object Keypoint Similarity (OKS) between the predicted pose and the assigned ground truth for each grid.

Positive grid tokens yield bounding box, keypoint coordinates, and visibility predictions. Keypoint coordinates are derived via the DCC, while other predictions come from pointwise convolution layers. The losses applied are IoU loss for bounding boxes \mathcal{L}{\text {bbox }} , MLE loss for keypoints \mathcal{L}{\text {mle }} , and BCE loss for visibility \mathcal{L}_{\text {vis }} .

Given the DCC’s computational demands, we implement a pointwise convolution layer for preliminary coordinate regression, similar to YOLO-Pose [30], to mitigate out-ofmemory issues. This regressed keypoints \mathrm{kpt}{\text {reg }} serves as a proxy in SimOTA for positive grid selection, with the decoded keypoints \mathrm{kpt}{\text {dec }} later used to calculate OKS. The regression branch’s loss is OKS loss [30]:

\mathcal{L}{\text {proxy }}=1-\text { OKS }\left(\mathrm{kpt}{\text {reg }}, \mathrm{kpt}_{\text {dec }}\right) .

The total loss for the proposed model is

\mathcal{L}=\lambda_{1} \mathcal{L}{\text {bbox }}+\lambda{2} \mathcal{L}{\text {mle }}+\lambda{3} \mathcal{L}{\text {proxy }}+\lambda{4} \mathcal{L}{\text {cls }}+\mathcal{L}{\text {vis }},

with hyperparameters \lambda_{1}, \lambda_{2}, \lambda_{3} , and \lambda_{4} set at \lambda_{1}=\lambda_{2}= 5, \lambda_{3}=10 , and \lambda_{4}=2 .

推理。在推理阶段，我们的模型使用0.1的分数阈值和网格过滤的非最大抑制。然后，它从选定的网格解码姿态特征为热图，利用热图的积分来导出关键点坐标。这种选择性解码方法最小化了需要处理的特征数量，从而降低了计算成本。

4、实验

4.1、设置

数据集。实验主要在COCO2017关键点检测基准[23]上进行，该基准包含约250K个人实例，具有17个关键点。在val2017和test-dev集上与最先进的方法进行了性能比较。为了探索模型的上限性能，训练还扩展到了其他数据集：CrowdPose [18]、AIC [43]、MPII [1]、JHMDB [14]、Halpe [8]和PoseTrack18 [2]。这些注释被转换为COCO格式。RTMO在CrowdPose基准[18]上进一步评估，该基准以拥挤和遮挡的场景而闻名，由20K张图像和约80K个人组成，具有14个关键点。对于这两个数据集，基于OKS的平均精度（AP）用作评估指标。

实现细节。在训练过程中，我们采用了YOLOX [9]的图像增强管道，包括mosaic增强、随机颜色调整、几何变换和MixUp [49]。训练图像的尺寸调整为[480,800]。COCO和CrowdPose数据集的训练周期数分别为600和700。训练过程分为两个阶段：第一阶段涉及使用姿势标注训练代理分支和DCC，第二阶段将代理分支的目标转向从DCC解码出的姿势。使用AdamW优化器[26]，权重衰减为0.05，在Nvidia GeForce RTX 3090 GPU上以批量大小256进行训练。两个训练阶段的初始学习率分别为$4\times 10^{-3}$和$5\times 10^{-4}$，通过余弦退火衰减到$2\times 10^{-4}$。对于推理，图像被调整为640。使用ONNXRuntime在Intel Xeon Gold CPU上测量CPU延迟。使用ONNXRuntime和TensorRT在NVIDIA V100 GPU上测试GPU延迟，并使用半精度浮点数（FP16）格式。

4.2、基准结果

COCO 为了评估RTMO与其他实时姿态估计器，我们在COCO val 2017数据集上测量了AP和推理延迟。对于单阶段方法，我们考虑了KAPAO [33]、YOLOv8-Pose [16]和YOLOXPose - YOLO-Pose [30]在YOLOX [9]上的改编。对于自上而下的方法，我们选择了RLE [19]、SimCC [22]和RTMPose [15]进行比较。RTMDetnano [29]，一个高效的对象检测模型，作为自上而下模型的“人类检测器”。由于自上而下模型在图像中出现更多人时会变慢，我们根据人数将COCO val2017集进行分区，并相应地评估自上而下模型的运行速度。如Fig. 3所示，RTMO系列在性能和速度上均超越了类似的轻量级单阶段方法。与自上而下模型相比，RTMO-m和RTMO-1与RTMPose-m和RTMPose1的精度相当，在图像中出现更多人时速度更快。使用ONNXRuntime时，当有大约四个人时，RTMO与RTMPose的速度相当；使用TensorRT FP16时，当有两个人或更多人时，RTMO更快。这表明了RTMO在多人物场景中的优势。重要的是，尽管处理中的令牌数量随着图像中的人数而变化，但推理延迟的差异很小。例如，在有超过10人的子集中，GPU上的RTMO-1的延迟仅比一个子集高出约0.1毫秒，占总体延迟的约0.5%。

我们对RTMO在COCO test-dev上与领先的一阶段姿态估计器进行了评估，结果如表1所示。RTMO在速度和精度方面表现出显著的优势。具体来说，使用ResNet-50 [11]骨干网的RTMO-s性能优于PETR [36]，速度快10倍，同时保持相似的准确性。与轻量级模型如KAPAO和YOLO-Pose相比，RTMO在不同模型大小上的准确性均优于其他模型。当在COCO train2017上进行训练时，RTMO-l在所有测试模型中排名第二。表现最好的模型是使用Swin-L [25]骨干网的ED-Pose [46]，但该模型相当庞大且不利于部署。使用相同ResNet-50骨干网的RTMO在AP上超越ED-Pose 1.1%，并且速度更快。此外，将ED-Pose转换为ONNX格式导致延迟慢于其PyTorch模型，每帧约1.5秒。相比之下，RTMO-l的ONNX模型仅需19.1ms即可处理一张图像。通过在更多人体姿态数据集上进一步训练，RTMO-1在一阶段姿态估计器中以准确性排名第一。

CrowdPose 为了评估RTMO在具有挑战性的场景下的性能，我们在CrowdPose [18]基准上进行测试，该基准以密集人群、显著的人员重叠和遮挡为特点。结果总结如表2所示。在自下而上和单阶段方法中，RTMO-s的准确性可与DEKR [10]相媲美，但其参数仅为DEKR的15%。当在CrowdPose数据集上训练时，尽管模型大小较小，但RTMO-1超过了使用Swin-L [25]骨干网的ED-Pose [46]。值得注意的是，RTMO-l主要在中等和困难样本上超过ED-Pose，这表明其在挑战性场景下的有效性。此外，通过使用更多的训练数据，RTMO-1达到了81.7%的AP，突显了该模型的容量。

4.3、定量结果

RTMO采用坐标分类方法，在具有挑战性的多人物场景中表现出强大的性能，其中人物相对较小，且频繁出现遮挡。图4表明，即使在这些困难的条件下，RTMO也能生成空间上准确的热图，为每个关键点提供稳健且具有上下文感知的预测。

4.4、消融实验

分类与回归。为了评估坐标分类相对于回归的有效性，我们将模型中的1-D热图生成替换为用于回归的全连接层，并通过OKS损失[30]进行监督。表3比较了两种方法的性能。使用DCC模块和MLE损失，坐标分类在COCO上的性能优于回归2.1% AP。

坐标分类的损失函数。与其他使用KLD作为损失函数的具有坐标分类的人体姿态估计方法不同，我们的研究指出KLD对于RTMO来说是不够的。表3中的证据表明，与我们的MLE损失相比，使用KLD会导致较低的准确性。我们将这种改进归因于MLE损失函数中可学习的方差的包含，这有助于平衡硬样本和易样本之间的学习。具体来说，在一阶段姿态估计器框架中，每个网格的难度级别会因分配实例的姿态和大小以及网格与实例之间的相对位置等因素而有所不同。KLD无法考虑这种变化，因此在此上下文中效果较差。

动态策略在坐标分类中。我们首先采用类似于DFL [20]的静态坐标分类策略，其中bins分布在每个网格的固定范围内。该方法在COCO数据集上超过了回归方法，但在CrowdPose上表现不佳。通过引入动态bin分配（DBA）策略，该基线在两个数据集上的性能都下降了。这是合理的，因为不同样本中每个bin的语义没有相应的表示调整。这个问题通过结合动态bin编码（DBE）得到了纠正。使用DBE，我们的DCC方法在两个数据集上都超过了静态策略的效率。此外，没有动态bin编码（DBE），附近bins的概率可能会有很大差异，如图5所示，这与相邻空间位置应该有相似概率的期望相矛盾。相比之下，结合DBE会导致输出热图更加平滑，表明通过启用更好地捕获附近位置之间相似性的表示向量，解码器训练得到了改进。

特征图选择。特征金字塔 [17] 利用多尺度特征检测不同大小的对象；更深层次的特征通常用于检测更大的对象。我们最初的模型受到YOLOX的启发，使用P3、P4、P5特征，步幅分别为8、16和32像素。然而，P3在模型头部占用了78.5%的FLOPs，同时占用了10.7%的正确检测。为了提高效率，我们专注于P4和P5。如表4所示，省略P3可以显著提高速度，同时只有很小的精度损失，这表明P4和P5足以用于多人姿态估计。这表明P3在检测较小实例方面的作用可以通过其余特征得到充分的补偿。

5、结论

总之，我们的RTMO模型在一阶段多人姿态估计中显著提高了速度与精度之间的权衡。通过在基于YOLO的框架中集成坐标分类，我们实现了实时处理和高精度。我们的方法具有动态坐标分类器和基于最大似然估计的损失函数，有效提高了密集预测模型中的位置精度。这一突破不仅增强了姿态估计，而且为视觉检测任务密集预测范围的未来发展奠定了坚实的基础。