MMPose 中的 YOLOX-Pose:基于 OKS 损失的单阶段多人姿态估计实践指南
【免费下载链接】mmposeOpenMMLab Pose Estimation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmpose
导读
YOLOX-Pose 是 MMPose 基于 YOLO-Pose(CVPRW 2022)思想与 YOLOX 目标检测框架实现的一套自底向上(bottom-up)单阶段多人姿态估计方案,它将人检测与关键点回归统一到一次前向推理中,天然适配实时应用场景。本篇文章以 yoloxpose_coco.md 为骨架,结合仓库中的配置文件与源码,系统讲解 YOLOX-Pose 的算法原理、四种规格模型的配置结构、训练流程与评估方式,读完后你将能够在 MMPose 中完成 YOLOX-Pose 的复现、调参与评估。
一、算法背景:从 YOLO-Pose 到 YOLOX-Pose
YOLO-Pose 是一篇发表于 CVPRW 2022 的工作(论文标题为Yolo-pose: Enhancing YOLO for multi person pose estimation using object keypoint similarity loss),核心思想是在单次前向传播中同时完成所有人实例的检测与关键点回归,属于典型的自底向上方法,避免了 Top-Down 两阶段方案中"先检测、再逐人裁剪推理"所带来的计算冗余。
MMPose 在 YOLO 检测框架上复现并扩展了该思路,实现为YOLOX-Pose,相关说明见 README.md:
- 继承 YOLO-Pose"检测与姿态估计统一"的优势;
- 为更精确地预测关键点位置,为不同关节分别回归偏移量,使用带自适应卷积的独立分支;
- 从而针对每个关键点优化特征提取过程。
算法层面,YOLOX-Pose 的关键创新是引入OKS(Object Keypoint Similarity,目标关键点相似度)损失作为标签分配与回归监督的核心度量,这也是 YOLO-Pose 论文标题所强调的贡献点。
二、整体架构:CSPDarknet + YOLOXPAFPN + YOLOXPoseHead
从 yoloxpose_s_8xb32-300e_coco-640.py 可以看出,YOLOX-Pose 的模型由三部分组成,整体注册为BottomupPoseEstimator:
| 模块 | 类型 | 关键参数 |
|---|---|---|
| backbone | CSPDarknet | out_indices=(2,3,4),spp_kernal_sizes=(5,9,13),激活函数Swish |
| neck | YOLOXPAFPN | in_channels=[128,256,512],out_channels=128,num_csp_blocks=1 |
| head | YOLOXPoseHead | num_keypoints=17,featmap_strides=(8,16,32) |
其中 backbone 的初始化直接复用了 MMDetection 中在 COCO 上预训练好的 YOLOX 权重(prefix='backbone.'),这意味着姿态模型以检测特征为基础进行迁移学习。neck 与 head 中均使用BN(momentum=0.03, eps=0.001)与Swish激活,与 YOLOX 原始训练设定保持一致。
YOLOXPoseHead的底层实现位于 yoloxpose_head.py,其模块YOLOXPoseHeadModule在每一层特征图上初始化了三个平行分支:
- 分类分支(cls branch):
stacked_convs层堆叠卷积后输出num_classes的分类分数; - 回归分支(reg branch):输出 4 维边界框
bbox_preds与 1 维目标度objectness; - 姿态分支(pose branch):采用
stacked_convs * 2层(即两倍深度的)自适应卷积,输出num_keypoints * 2维关键点偏移kpt_offsets与num_keypoints维关键点可见度kpt_vis。
姿态分支使用两倍深度的卷积层,正是"为不同关节分别学习偏移量、独立优化特征提取"这一设计意图的直接体现(见 yoloxpose_head.py)。
此外,head 在初始化时会通过bias_init_with_prob(0.01)设置分类与目标度分支的偏置,以保证训练初期预测稳定(见 yoloxpose_head.py)。
2.1 标签分配:SimOTAAssigner 与 OKS
YOLOX-Pose 的标签分配使用SimOTAAssigner,并设置dynamic_k_indicator='oks'——即以 OKS 作为动态 k 的衡量指标来决定每个真实目标匹配的 anchor 数量:
assigner=dict(type='SimOTAAssigner', dynamic_k_indicator='oks'), overlaps_power=0.5,overlaps_power=0.5用于对 OKS 重叠度做幂次缩放,影响分配时的代价计算。这是"用 OKS 损失增强 YOLO"这一论文思想在实现层的直接落地。
2.2 损失函数组合
配置文件 中定义了六个损失,分工如下:
| 损失 | 类型 | reduction | loss_weight | 作用 |
|---|---|---|---|---|
loss_cls | BCELoss | sum | 1.0 | 分类分数 |
loss_bbox | IoULoss(mode='square',eps=1e-16) | sum | 5.0 | 边界框回归 |
loss_obj | BCELoss(use_target_weight=True) | sum | 1.0 | 目标度 |
loss_oks | OKSLoss | none | 30.0 | 关键点回归(核心) |
loss_vis | BCELoss(use_target_weight=True) | mean | 1.0 | 关键点可见度 |
loss_bbox_aux | L1Loss | sum | 1.0 | 辅助边界框回归 |
其中loss_oks的metainfo指向configs/_base_/datasets/coco.py,用于读取 COCO 17 个关键点的归一化标准差(按类别尺度归一化关键点回归目标),norm_target_weight=True表示按目标权重归一化。OKS 损失权重 30.0 是全部损失中最大的,突出了关键点回归在整个模型目标中的核心地位。
2.3 测试配置
test_cfg=dict( score_thr=0.01, nms_thr=0.65, )推理时保留分数高于 0.01 的预测,并对其做阈值为 0.65 的 NMS 后处理。
三、四种规格模型与 COCO 精度基准
仓库在configs/body_2d_keypoint/yoloxpose/coco/下提供了 tiny / s / m / l 四种规格,对应 YOLOX 系列的宽度与深度缩放因子:
| 配置 | 输入尺寸 | widen_factor | deepen_factor | 预训练权重来源 |
|---|---|---|---|---|
| yoloxpose_tiny_4xb64-300e_coco-416.py | 416×416 | 0.375 | 0.33 | YOLOX-tiny |
| yoloxpose_s_8xb32-300e_coco-640.py | 640×640 | 0.5 | 0.33 | YOLOX-s |
| yoloxpose_m_8xb32-300e_coco-640.py | 640×640 | 0.75 | 0.67 | YOLOX-m |
| yoloxpose_l_8xb32-300e_coco-640.py | 640×640 | 1.0 | 1.0 | YOLOX-l |
在 COCO val2017 上的精度如下(数据来自 yoloxpose_coco.md 与 yoloxpose_coco.yml):
| Arch | Input Size | AP | AP^50 | AP^75 | AR | AR^50 |
|---|---|---|---|---|---|---|
| yoloxpose_tiny | 416×416 | 0.526 | 0.793 | 0.556 | 0.571 | 0.833 |
| yoloxpose_s | 640×640 | 0.641 | 0.872 | 0.702 | 0.682 | 0.902 |
| yoloxpose_m | 640×640 | 0.695 | 0.899 | 0.766 | 0.733 | 0.926 |
| yoloxpose_l | 640×640 | 0.712 | 0.901 | 0.782 | 0.749 | 0.926 |
可以看到,从 tiny 到 l,AP 随模型容量稳步提升:tiny 以 416 输入获得 52.6 AP,适合轻量部署;l 在 640 输入下达到 71.2 AP,同时保持着自底向上方法"一次前向、全员检测"的效率特性。
3.1 变体配置的继承方式
tiny / m / l 三个配置均以 s 配置为基础继承。例如 tiny 配置通过_base_ = './yoloxpose_s_8xb32-300e_coco-640.py'继承全部结构,仅覆盖模型缩放因子、预训练权重、neck 通道数、batch_size 与输入尺寸;而 m / l 则在覆盖widen_factor、deepen_factor、in_channels/out_channels、num_csp_blocks与权重路径。这种_base_继承机制让"换一个规格只需改几行"成为可能。
四、训练配置深度解析
4.1 训练循环与学习率调度
train_cfg = dict( _delete_=True, type='EpochBasedTrainLoop', max_epochs=300, val_interval=10, dynamic_intervals=[(280, 1)])- 总训练 300 epoch,每 10 epoch 验证一次;
- 从第 280 epoch 起,验证间隔动态调整为 1 epoch(配合下文模式切换 hook 使用)。
学习率采用三段式调度:
param_scheduler = [ dict(type='QuadraticWarmupLR', by_epoch=True, begin=0, end=5, convert_to_iter_based=True), dict(type='CosineAnnealingLR', eta_min=0.0002, begin=5, T_max=280, end=280, by_epoch=True, convert_to_iter_based=True), dict(type='ConstantLR', by_epoch=True, factor=1, begin=280, end=300), ]即:前 5 epoch 使用二次曲线 warmup(按 iteration 计),5–280 epoch 余弦退火至eta_min=0.0002,最后 20 epoch 保持恒定学习率。优化器为AdamW(lr=0.004, weight_decay=0.05),其中 BN 参数与 bias 不做权重衰减,并启用clip_grad(max_norm=0.1)防止梯度爆炸。配置同时声明auto_scale_lr = dict(base_batch_size=256),在改变 batch_size 时可自动线性缩放学习率。
4.2 两阶段训练:YOLOX 风格的强增强与模式切换
YOLOX-Pose 完整继承了 YOLOX 的"两阶段训练"策略,这也是训练配置中最值得关注的部分:
阶段一(0–280 epoch):启用Mosaic、BottomupRandomAffine(shift_factor=0.1、rotate_factor=10、scale_factor=(0.75,1.0),采用透视变换模式)、YOLOXMixUp(ratio_range=(0.8,1.6))等强数据增强,见 train_pipeline_stage1。
阶段二(最后 20 epoch):关闭 Mosaic 与 MixUp,仅保留不带随机仿射的BottomupRandomAffine(各 prob 置 0)、YOLOXHSVRandomAug与RandomFlip,让模型在接近真实分布的数据上精调。
两个阶段的切换由自定义 hook 完成:
custom_hooks = [ dict(type='YOLOXPoseModeSwitchHook', num_last_epochs=20, new_train_pipeline=train_pipeline_stage2, priority=48), dict(type='SyncNormHook', priority=48), dict(type='EMAHook', ema_type='ExpMomentumEMA', momentum=0.0002, update_buffers=True, strict_load=False, priority=49), ]YOLOXPoseModeSwitchHook(实现在 mode_switch_hooks.py)在最后 20 epoch 将训练 pipeline 切换到阶段二;SyncNormHook在分布式训练时同步 BN 统计量;EMAHook维护指数滑动平均(momentum 0.0002)的模型副本,验证与测试时使用 EMA 权重,是 YOLO 系列提升稳定性的常用技巧。
数据预处理端,PoseDataPreprocessor使用pad_size_divisor=32将输入 padding 到 32 的倍数,并在 batch 级启用BatchSyncRandomResize(random_size_range=(480,800)),在训练中动态改变输入分辨率以增强尺度鲁棒性。
4.3 数据流与标签编码
data_mode='bottomup'明确声明了自底向上的数据组织方式。训练集使用CocoDataset加载person_keypoints_train2017.json,并经过FilterAnnotations(by_kpt=True, by_box=True)过滤掉不含关键点或边界框的样本。
标签编码由 codec 完成:
codec = dict(type='YOLOXPoseAnnotationProcessor', input_size=input_size)YOLOXPoseAnnotationProcessor实现在 annotation_processors.py,其职责是将 COCO 标注转换为 YOLOX-Pose 需要的输入格式:将类别 ID 减 1 转为 0 起始的标签,并通过label_mapping_table将bbox/bbox_labels/keypoints/keypoints_visible/area映射为bboxes/labels/keypoints/keypoints_visible/areas。它还提供了expand_bbox选项:当开启时,会把边界框扩张到包含所有可见关键点(见 annotation_processors.py),确保关键点不会落在框外——这对后续的检测框与关键点一致性非常重要。
4.4 评估配置
val_evaluator = dict( type='CocoMetric', ann_file=data_root + 'coco/annotations/person_keypoints_val2017.json', score_mode='bbox', nms_mode='none', )评估使用 COCO 官方指标,关键点检测框(bbox)作为打分依据,nms_mode='none'表示在指标内部不做额外的 NMS 处理。验证与测试共用同一套 dataloader 与 evaluator。
五、从零开始:训练与评估实操
在准备好 COCO 数据集(目录结构为data/coco/annotations/person_keypoints_{train,val}2017.json与data/coco/{train,val}2017/)并完成环境安装后,即可使用仓库自带的 tools/train.py 开始训练:
# 单卡训练 s 规格 python tools/train.py configs/body_2d_keypoint/yoloxpose/coco/yoloxpose_s_8xb32-300e_coco-640.py # 8 卡分布式训练 bash tools/dist_train.sh configs/body_2d_keypoint/yoloxpose/coco/yoloxpose_s_8xb32-300e_coco-640.py 8评估与推理可使用 tools/test.py:
python tools/test.py configs/body_2d_keypoint/yoloxpose/coco/yoloxpose_s_8xb32-300e_coco-640.py \ /path/to/yoloxpose_s_8xb32-300e_coco-640-56c79c1f_20230829.pth预训练权重可从配置对应 yml 元数据中记录的地址下载(如 yoloxpose_coco.yml 中的Weights字段)。训练时的 checkpoint 每 10 epoch 保存一次(CheckpointHook(interval=10, max_keep_ckpts=3)),只保留最近 3 份以控制磁盘占用。
若希望用训练好的模型直接对图片推理,可参考 demo/inferencer_demo.py 以及 推理指南 中基于 MMPose Inferencer 的用法。
六、总结与适用场景
YOLOX-Pose 是 MMPose 在"检测与姿态估计一体化"方向上的重要实现:
- 单阶段自底向上:一次前向同时输出全部人体的框、类别、17 个关键点及可见度,推理高效;
- OKS 驱动的标签分配与回归:
SimOTAAssigner+OKSLoss让关键点质量直接参与匹配与监督; - 两阶段训练策略:Mosaic/MixUp 强增强配合最后 20 epoch 的 pipeline 切换与 EMA,训练稳定、收敛良好;
- 多规格可扩展:通过
_base_继承与 widen/deepen 因子即可在 52.6 AP(tiny)到 71.2 AP(l)之间按算力选型。
对于需要实时处理视频流、多人密集场景、且希望同时获得检测框与骨架的开发者而言,YOLOX-Pose 是一个兼顾速度与精度的开箱即用方案。深入阅读 yoloxpose_head.py 与 annotation_processors.py 可以进一步理解其损失计算与标签分配细节,便于在此基础上定制自己的姿态模型。
【免费下载链接】mmposeOpenMMLab Pose Estimation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmpose
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考