• 人工智能
• 计算机视觉
• 深度学习

导读

本文以 hrnet_coco.md 为骨架，系统讲解 ViTPose 仓库中基于 Associative Embedding（关联嵌入）与 HRNet 骨干的 COCO 自底向上（bottom-up）人体姿态估计方案：从算法原理、官方评测指标，到 hrnet_w32_coco_512x512.py / hrnet_w48_coco_512x512.py 的逐段配置解读、损失函数与后处理分组源码剖析，再到训练、测试与单图推理的完整实战命令。读完本文，你将能够理解自底向上姿态估计"先检点、后分组"的完整数据流，并可直接在 COCO 数据集上复现文档中的评测指标。

一、方法背景：为什么选择 Associative Embedding + HRNet

1.1 三种核心算法引用

文档开头以 BibTeX 形式引用了三篇奠基性工作，对应方法、骨干与数据集的组合：

• Associative Embedding（NIPS'2017，Newell et al.）：提出"检测 + 分组"的自底向上范式。网络对每个关键点不仅预测热力图（heatmap），还额外预测一个 tag（嵌入向量）。同一人的各关键点 tag 彼此接近，不同人的 tag 相互远离，据此可将散落的关键点聚合成完整姿态。
• HRNet（CVPR'2019，Sun et al.）：高分辨率表示学习网络。整个网络始终维护高分辨率特征分支，并通过多分支并行与跨分辨率特征融合，避免因下采样丢失空间细节，非常契合小目标与密集人群场景下的关键点定位。
• COCO（ECCV'2014）：评测基准数据集，提供person_keypoints_train2017.json与person_keypoints_val2017.json等标注文件。

三者结合即本文的核心方案：用 HRNet 提取高分辨率特征，由关联嵌入头同时输出 17 类关键点热力图与分组 tag，在 COCO val2017 上验证精度。

1.2 自底向上 vs 自顶向下

与"先检测人框、再逐框回归关键点"的自顶向下（top-down）方法不同，自底向上方法直接在整张图上检测所有关键点再分组。其优势在于：天然规避了自顶向下方法对检测框质量的依赖、无需重复裁剪，但对密集场景的分组正确性要求更高——这正是 AE tag 机制要解决的问题。更完整的说明可参见 associative_embedding 目录 README。

二、官方评测结果（文档核心数据）

文档给出了 HRNet-w32 / HRNet-w48 两种骨干在 COCO val2017 上的两套评测结果。

2.1 单尺度测试（无 multi-scale test）

2.2 多尺度测试（3 个默认尺度 [2, 1, 0.5]）

关键结论：多尺度测试（multi-scale test）显著提升指标——HRNet-w48 的 AP 从 0.665 提升至 0.712（+4.7 个点）。原因在于多尺度推理聚合了不同分辨率下的热力图响应，对尺度变化更鲁棒；代价是推理耗时按尺度数成倍增加。文档中 ckpt/log 链接对应 OpenMMLab 官方发布的权重与训练日志，可用于直接复现或作为预训练起点。

三、配置逐段精读：以 hrnet_w48_coco_512x512.py 为例

下文以 hrnet_w48_coco_512x512.py 为蓝本逐段拆解；hrnet_w32_coco_512x512.py 除骨干宽度（32 vs 48）与samples_per_gpu（24 vs 16）外结构一致。

3.1 基础继承与训练节奏

_base_ = [ '../../../../_base_/default_runtime.py', '../../../../_base_/datasets/coco.py' ] checkpoint_config = dict(interval=50) evaluation = dict(interval=50, metric='mAP', save_best='AP')

• _base_继承两份基础配置：default_runtime.py（日志、checkpoint 间隔、分布式参数等）与 datasets/coco.py（COCO 的dataset_info，含 17 个关键点名称、骨架连接、joint_weights、sigmas等）。
• 训练 300 个 epoch 期间，每 50 个 epoch 保存一次 checkpoint 并在验证集上以mAP为指标评估，同时按save_best='AP'保存最优模型。

3.2 优化器与学习率策略

optimizer = dict(type='Adam', lr=0.0015) optimizer_config = dict(grad_clip=None) lr_config = dict( policy='step', warmup='linear', warmup_iters=500, warmup_ratio=0.001, step=[200, 260]) total_epochs = 300

• 采用Adam优化器，初始学习率 0.0015，不裁剪梯度。
• 前 500 次迭代以 0.001 倍初始学习率线性预热，随后在第 200 与第 260 epoch 处按step策略分档衰减，共训练 300 epoch。

3.3 通道与数据配置

channel_cfg = dict( dataset_joints=17, dataset_channel=[[0, 1, ..., 16]], inference_channel=[0, 1, ..., 16]) data_cfg = dict( image_size=512, base_size=256, base_sigma=2, heatmap_size=[128], num_joints=channel_cfg['dataset_joints'], dataset_channel=channel_cfg['dataset_channel'], inference_channel=channel_cfg['inference_channel'], num_scales=1, scale_aware_sigma=False)

• image_size=512：训练输入图像短边尺寸；base_size=256为基准尺寸（配合test_scale_factor=[1]还原原始尺度）。
• heatmap_size=[128]：输出热力图尺寸为 128x128，即相对于 512 输入做了 4 倍下采样。
• num_scales=1、scale_aware_sigma=False：当前方案为单尺度训练、使用固定高斯 sigma，而非 HigherHRNet 式的多尺度监督。
• 640x640 输入版本（如 hrnet_w48_coco_640x640.py）将base_size提高到 320、heatmap_size提高到 160，以更高分辨率换取更精细定位。

3.4 模型定义：AssociativeEmbedding 检测器

model = dict( type='AssociativeEmbedding', pretrained='https://download.openmmlab.com/mmpose/pretrain_models/hrnet_w48-8ef0771d.pth', backbone=dict( type='HRNet', in_channels=3, extra=dict( stage1=dict(num_modules=1, num_branches=1, block='BOTTLENECK', num_blocks=(4,), num_channels=(64,)), stage2=dict(num_modules=1, num_branches=2, block='BASIC', num_blocks=(4, 4), num_channels=(48, 96)), stage3=dict(num_modules=4, num_branches=3, block='BASIC', num_blocks=(4, 4, 4), num_channels=(48, 96, 192)), stage4=dict(num_modules=3, num_branches=4, block='BASIC', num_blocks=(4, 4, 4, 4), num_channels=(48, 96, 192, 384)))), keypoint_head=dict( type='AESimpleHead', in_channels=48, num_joints=17, num_deconv_layers=0, tag_per_joint=True, with_ae_loss=[True], extra=dict(final_conv_kernel=1), loss_keypoint=dict( type='MultiLossFactory', num_joints=17, num_stages=1, ae_loss_type='exp', with_ae_loss=[True], push_loss_factor=[0.001], pull_loss_factor=[0.001], with_heatmaps_loss=[True], heatmaps_loss_factor=[1.0])), ... )

• 检测器：type='AssociativeEmbedding'对应 mmpose/models/detectors/associative_embedding.py 中的AssociativeEmbedding类。其forward_train依次执行 backbone 提取特征 → head 输出 →keypoint_head.get_loss计算损失；forward_test则执行多尺度/翻转推理、热力图与 tag 聚合，再交给HeatmapParser完成分组。
• 骨干 HRNet：stage1 为单分支 Bottleneck（64 通道），后续逐级扩到 2/3/4 个并行分支（w48 对应 48/96/192/384 通道），全程并行高分辨率特征。
• 关联嵌入头 AESimpleHead：继承自 mmpose/models/heads/ae_simple_head.py 中的DeconvHead。tag_per_joint=True表示每个关键点类别预测独立的 tag 通道，输出通道数为num_joints + num_joints = 34（17 个热力图 + 17 个 tag）；若为 False 则共用 1 个 tag 通道（共 18 通道）。num_deconv_layers=0表示直接以 HRNet 输出分辨率为热力图分辨率。

3.5 测试配置：决定推理与分组行为

test_cfg=dict( num_joints=17, max_num_people=30, scale_factor=[1], with_heatmaps=[True], with_ae=[True], project2image=True, align_corners=False, nms_kernel=5, nms_padding=2, tag_per_joint=True, detection_threshold=0.1, tag_threshold=1, use_detection_val=True, ignore_too_much=False, adjust=True, refine=True, flip_test=True)

各参数含义如下：

• max_num_people=30：每张图最多分组 30 人，热力图topk的 K 值也取 30（见HeatmapParser.top_k）。
• nms_kernel=5、nms_padding=2：对热力图做 5x5 最大池化 NMS，抑制重复峰值。
• detection_threshold=0.1：关键点峰值分数低于该值视为未检出；tag_threshold=1：tag 距离超过该阈值视为不同人（见 mmpose/core/post_processing/group.py 的_match_by_tag中的diff_saved[row][col] < params.tag_threshold判断）。
• adjust=True, refine=True：坐标微调与分组后细化。
• flip_test=True：开启水平翻转测试，热力图取平均、tag 取 concat（tag 维度翻倍），可稳定提升精度。

3.6 数据管线与数据集

train_pipeline = [ dict(type='LoadImageFromFile'), dict(type='BottomUpRandomAffine', rot_factor=30, scale_factor=[0.75, 1.5], scale_type='short', trans_factor=40), dict(type='BottomUpRandomFlip', flip_prob=0.5), dict(type='ToTensor'), dict(type='NormalizeTensor', mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), dict(type='BottomUpGenerateTarget', sigma=2, max_num_people=30), dict(type='Collect', keys=['img', 'joints', 'targets', 'masks'], meta_keys=[]), ]

• 训练增强包括随机仿射（旋转 ±30°、尺度 0.75~1.5、平移 40px）与随机翻转；目标生成采用sigma=2的高斯核，每图最多 30 人。
• 验证/测试管线使用BottomUpGetImgSize+BottomUpResizeAlign，test_scale_factor=[1]单尺度。
• 数据集为BottomUpCocoDataset，实现位于 mmpose/datasets/datasets/bottom_up/bottom_up_coco.py，标注文件指向data/coco/annotations/person_keypoints_{train,val}2017.json。训练时samples_per_gpu=16（w48）或 24（w32），每 GPU 2 个 worker。

四、源码级原理：tag 是怎么"聚人"的

4.1 联合损失：MultiLossFactory

loss_keypoint使用的MultiLossFactory定义在 mmpose/models/losses/multi_loss_factory.py，它按num_stages逐阶段组合两类损失：

• HeatmapLoss（关键点检测）：对输出前num_joints通道与目标热力图计算均方误差，乘heatmaps_loss_factor=[1.0]。
• AELoss（关联嵌入）：作用于 tag 通道，包含两项——
  • pull loss（吸引）：同一人的各关键点 tag 应互相接近，最小化其与均值 tag 的平方距离；
  • push loss（排斥）：不同人的均值 tag 应互相远离，ae_loss_type='exp'时用指数形式push = sum(exp(-diff^2)) - num_tags实现。

两者权重均为0.001。注意 AELoss 中joints记录了各关键点在 heatmap 上的位置坐标（joint[0]），用于从 tag 图采样对应位置的 tag 值——这正是"每个关键点带一个 tag"的监督来源。

4.2 后处理：从热力图到多人姿态

推理分组的完整链路在 mmpose/models/detectors/associative_embedding.py 的forward_test与 mmpose/core/post_processing/group.py 中实现：

1. split_ae_outputs将 head 输出拆分为热力图与 tag；
2. 若开启flip_test，对水平翻转图像再推理一次，热力图平均融合、tag 拼接；
3. aggregate_stage_flip/aggregate_scale完成多尺度聚合（本文单尺度仅 1 个 scale）；
4. HeatmapParser.parse执行最终分组：top_k先用nms_kernel=5的 MaxPool2d 做 NMS，再取每个关键点热力图的前max_num_people个峰值位置及其对应 tag；_match_by_tag按关键点连接顺序（COCO 17 点有专门定义的joint_order）逐类关键点做匈牙利匹配（Munkres 算法），以 tag 距离为代价矩阵，将各关键点并入已有分组或新建分组。

从源码结构可以推断：分组的正确率高度依赖 tag 质量，因此tag_per_joint=True的独立 tag 通道与 pull/push 损失系数是整套方案的关键超参。

五、实战：训练、测试与推理

5.1 单卡训练

python tools/train.py configs/body/2d_kpt_sview_rgb_img/associative_embedding/coco/hrnet_w48_coco_512x512.py

默认在./work_dirs/hrnet_w48_coco_512x512下保存日志与权重。常用参数（详见 tools/train.py）：

• --work-dir：覆盖默认工作目录；
• --resume-from：从指定 checkpoint 恢复训练；
• --no-validate：跳过训练期间验证；
• --gpu-id：指定 GPU（非分布式模式默认 0）；
• --seed/--deterministic：固定随机种子保证可复现；
• --cfg-options：以k=v形式临时覆盖配置，如model.backbone.extra.stage4.num_channels；
• --launcher：none（单卡）/pytorch（分布式）/slurm。

多卡训练可使用仓库提供的 tools/dist_train.sh（如bash tools/dist_train.sh <CONFIG> <GPU_NUM>）。

5.2 测试与评估

python tools/test.py \ configs/body/2d_kpt_sview_rgb_img/associative_embedding/coco/hrnet_w48_coco_512x512.py \ /path/to/hrnet_w48_coco_512x512.pth \ --eval mAP

• --eval mAP：对应evaluation.metric='mAP'，输出 COCO 姿态 AP/AR 系列指标；
• --out：保存结果文件；--fuse-conv-bn：融合 Conv-BN 提升推理速度；
• 评估实现位于 mmpose/datasets/datasets/bottom_up/bottom_up_coco.py，预测结果会先写为result_keypoints.json再经 COCOeval 计算。

5.3 单图推理演示

python demo/bottom_up_img_demo.py \ configs/body/2d_kpt_sview_rgb_img/associative_embedding/coco/hrnet_w48_coco_512x512.py \ /path/to/hrnet_w48_coco_512x512.pth \ --img-path tests/data/coco/000000000785.jpg \ --out-img-root vis_results

• 支持单张图片或图片目录（png/jpg/jpeg/tiff/bmp），见 demo/bottom_up_img_demo.py；
• 关键参数：--kpt-thr（默认 0.3）过滤低分关键点、--pose-nms-thr（默认 0.9）控制姿态级 NMS 的 OKS 阈值、--device指定推理设备；
• 注意：底向上 demo 无需检测框输入，直接把整图喂给模型，返回的结果即为分组后的多人姿态。

六、同目录扩展配置一览

associative_embedding/coco目录下还提供了丰富的变体配置，便于横向对比（参见 目录 README 及对应 .md 评测页）：

• 输入分辨率：512x512 与 640x640 两档（如 hrnet_w48_coco_640x640.py）；
• UDP 变体：*_udp.py配置在 data_cfg 与 test_cfg 中开启use_udp=True（Unbiased Data Processing），可缓解下采样导致的坐标偏差；
• 其他骨干：ResNet 系列（resnet_coco.md）、Hourglass-AE（hourglass_ae_coco.md）、MobilenetV2（mobilenetv2_coco.md）以及多尺度监督的 HigherHRNet（higherhrnet_coco.md）。

七、小结

本文从 hrnet_coco.md 出发，将其中两张评测表所代表的方案完整还原为可运行的配置与源码链路：HRNet 高分辨率骨干 → AESimpleHead 联合输出热力图与 tag → MultiLossFactory 的 heatmap + push/pull 联合监督 → HeatmapParser 的 NMS + 匈牙利匹配分组。无论你是想复现指标、迁移到新数据集，还是深入理解自底向上姿态估计的工程细节，上述配置与源码路径都值得继续研读。

• 人工智能
• 计算机视觉
• 深度学习