企业官网建设流程全解析

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

本文围绕 Associative Embedding（AE）README 展开，系统讲解 AE 这一最流行的 2D Bottom-Up 姿态估计范式：先检测图像中所有关键点，再通过"关联标签（tag）"将它们分组为个体实例，并结合 ViTPose 仓库中 COCO-WholeBody 的 HRNet/HigherHRNet 配置、源码实现与结果，给出可复用的训练、评测与推理方案。读完本文，你将理解 tag 分组的核心原理、push/pull 损失的计算方式、HeatmapParser 后处理链路，并掌握如何基于仓库配置在 133 关键点的全身数据集上训练与部署 AE 模型。

一、AE 算法核心原理：检测 + 分组

Associative Embedding（AE）出自 Newell 等人 2017 年的论文Associative embedding: End-to-end learning for joint detection and grouping（NIPS'2017），是 2D 人体姿态估计中最主流的bottom-up（自底向上）方法之一。与 top-down 方法先检测人体框、再对每个人单独估计关键点不同，bottom-up 方法直接处理整张图像：

1. 检测阶段：一次性检测出图像中所有人的所有关键点（heatmap 输出）；
2. 分组阶段：把这些零散的关键点按人体实例归组，构成完整姿态。

分组的关键在于 AE 提出的"关联标签（tag）"机制：网络为每个检测到的关键点额外预测一个 tag 值。同一个人的所有关键点 tag 彼此接近，不同人的 tag 彼此远离，因此可以依据 tag 的相似性把关键点聚类为不同个体。这一"先全部检测、再按 tag 关联"的设计让 AE 天然适合拥挤、遮挡场景，无需依赖目标检测器。

仓库中该算法对应的注册实现为 AssociativeEmbedding 检测器，其 docstring 明确给出了与上述 README 一致的定义："Associative embedding pose detectors"。

1.1 tag 的训练信号：push 损失与 pull 损失

tag 的"同类相近、异类相远"特性不是凭空出现的，而是由专门的AE 损失监督学习得到。仓库中的实现位于 multi_loss_factory.py：

• pull 损失（拉近损失）：把同一人各关键点的 tag 拉向它们的均值，公式为pull = mean((tmp - mean(tag))^2)，对应代码见 AELoss.singleTagLoss 中的pull = pull + torch.mean((tmp - tags[-1].expand_as(tmp))**2)；
• push 损失（推远损失）：把不同人的 tag 均值互相推开。当ae_loss_type='exp'时，push = sum(exp(-(diff)^2)) - num_tags，等价于指数形式的排斥项；当ae_loss_type='max'时则采用 hinge 形式clamp(1 - |diff|, min=0)。两种类型的选择见 AELoss。

最终MultiLossFactory会按每个 stage 分别组合 heatmap 损失（均方误差）、push 损失与 pull 损失，并用heatmaps_loss_factor、push_loss_factor、pull_loss_factor加权，见 MultiLossFactory.forward。

1.2 tag 的维度：tag_per_joint

关于 tag 的通道数，配置项tag_per_joint决定：

• tag_per_joint=True：每个关键点独立预测自己的 tag，tag 通道数等于关键点数量（dim_tag = num_joints）；
• tag_per_joint=False：所有关键点共享单一 tag，dim_tag = 1。

该逻辑在 AESimpleHead 与 AEHigherResolutionHead 中均有体现：out_channels = num_joints + dim_tag，即网络的输出通道被划分为"关键点热图部分 + tag 部分"两部分。在推理时由split_ae_outputs（见 bottom_up_eval.py）按with_heatmaps与with_ae拆分。

二、仓库内的 AE 实现链路与后处理

2.1 检测器整体流程

AssociativeEmbedding.forward_test 展示了完整的推理链路：

1. 对每个测试尺度（test_scale_factor）前向网络，得到多尺度 heatmaps 与 tags；
2. 若开启flip_test=True，将图像水平翻转再次前向，并用flip_index恢复关键点左右对应关系（flip_feature_maps），随后与原始输出做平均聚合（aggregate_stage_flip）；
3. 多尺度之间通过aggregate_scale聚合（heatmap 取平均，tag 采用unsqueeze_concat拼接以保留多尺度信息）；
4. 调用 HeatmapParser.parse 完成关键点分组；
5. 通过get_group_preds把热图坐标映射回原图坐标。

2.2 HeatmapParser：从热图到人体姿态

HeatmapParser 是 AE 后处理的核心，其流水线为：

• Heatmap NMS：先用MaxPool2d(nms_kernel, stride=1, padding=nms_padding)对热图做非极大值抑制，保留局部峰值，见 nms；
• Top-K 选取：对每个关键点通道取置信度最高的max_num_people（默认 30）个候选位置及其对应 tag 值，见 top_k；
• tag 匹配：_match_by_tag用Munkres（匈牙利）算法求解候选关键点与已有 person 分组之间的最优匹配，匹配代价由 tag 距离决定，并受detection_threshold（关键点置信度门槛，默认 0.1）与tag_threshold（tag 距离门槛，默认 1）约束，见 _match_by_tag；
• adjust 与 refine：adjust=True时对坐标做 ±0.25 的亚像素微调（adjust）；refine=True时利用"检测到关键点的平均 tag"回头检索漏检关键点（refine）。

值得注意的是，对于 17 关键点的 COCO 人体模型，_Params还内置了固定的joint_order处理顺序（group.py），其他关键点数量则按自然顺序处理。

三、COCO-WholeBody：133 关键点的全身任务

AE 在仓库中落地于COCO-WholeBody全身姿态数据集（Jin 等人，ECCV'2020）。该数据集在 COCO 人体 17 点基础上扩展为133 个关键点，数据集定义见 configs/base/datasets/coco_wholebody.py，其分组结构为：

该划分在数据集类 BottomUpCocoWholeBodyDataset 的 docstring 与body_num/foot_num/face_num/left_hand_num/right_hand_num属性中均有体现，其_do_python_keypoint_eval会分别对 body、foot、face、left hand、right hand 以及 wholebody 六组分别调用 COCOeval 计算 mAP（bottom_up_coco_wholebody.py）。

评测指标方面，仓库配置统一使用metric='mAP'、key_indicator='AP'，即全身姿态的整体 AP/AR，并细分为 Body、Foot、Face、Hand 与 Whole 五档（见下文结果表）。

四、配置详解：HRNet 与 HigherHRNet 两种 AE 模型

仓库在该目录下提供了两组模型家族共 8 个训练配置（输入 512×512 与 640×640 两种分辨率）：

4.1 HRNet + AESimpleHead 配置逐项解读

以 hrnet_w32_coco_wholebody_512x512.py 为例，逐段拆解 AE 训练配置：

训练与优化器设置

_base_ = [ '../../../../_base_/default_runtime.py', '../../../../_base_/datasets/coco_wholebody.py' ] checkpoint_config = dict(interval=50) evaluation = dict(interval=50, metric='mAP', key_indicator='AP') 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，线性 warmup 500 轮（比例 0.001），在第 200、260 epoch 阶梯式下降，总计训练 300 epoch；
• 每 50 个 epoch 保存一次 checkpoint 并做一次验证，验证指标为 mAP，主指标为 AP。

通道与数据配置

channel_cfg = dict( num_output_channels=133, dataset_joints=133, dataset_channel=[list(range(133))], inference_channel=list(range(133))) 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, )

• num_joints=133对应 COCO-WholeBody 全部关键点；heatmap_size=[128]表示输出热图尺寸为输入 512 的 1/4；
• scale_aware_sigma=False表示不使用与人体框尺寸挂钩的自适应高斯 sigma，而统一使用base_sigma=2生成目标热图。

模型结构：AssociativeEmbedding + HRNet + AESimpleHead

model = dict( type='AssociativeEmbedding', pretrained='https://download.openmmlab.com/mmpose/' 'pretrain_models/hrnet_w32-36af842e.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=(32, 64)), stage3=dict(num_modules=4, num_branches=3, block='BASIC', num_blocks=(4, 4, 4), num_channels=(32, 64, 128)), stage4=dict(num_modules=3, num_branches=4, block='BASIC', num_blocks=(4, 4, 4, 4), num_channels=(32, 64, 128, 256)))), keypoint_head=dict( type='AESimpleHead', in_channels=32, num_joints=133, 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=133, 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], supervise_empty=False)), ... )

关键点：

• backbone为标准 HRNet-w32（stage1 单分支 BOTTLENECK 64 通道，stage2–stage4 逐步扩展为 2/3/4 分支、32/64/128/256 通道），输出 32 通道高分辨率特征；
• AESimpleHead不接反卷积（num_deconv_layers=0），直接在 HRNet 输出上做 1×1 卷积（final_conv_kernel=1），输出 133（热图）+ 133（tag）= 266 通道；
• 损失采用MultiLossFactory，ae_loss_type='exp'，heatmap 损失权重 1.0，push/pull 损失权重各 0.001，且supervise_empty=False表示不对空通道施加监督（见 HeatmapLoss 中empty_mask的处理）；
• pretrained加载 ImageNet 预训练的 HRNet-w32 权重。

test_cfg：后处理与测试策略

test_cfg=dict( num_joints=133, 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)

这些参数与上文 HeatmapParser 一一对应：单尺度测试（scale_factor=[1]）、热图与 tag 均输出、NMS 核 5/填充 2、检测置信度门槛 0.1、tag 距离门槛 1、开启 flip test、坐标微调与漏检修复。

数据流水线

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 倍、平移 40 像素、50% 概率水平翻转；
• BottomUpGenerateTarget以 sigma=2 生成热图目标，同时输出joints（用于 AE 损失的 tag 监督）与masks；
• 验证/测试流水线使用BottomUpGetImgSize与BottomUpResizeAlign完成尺度对齐。这些变换的实现见 bottom_up_transform.py。

数据加载

data_root = 'data/coco' data = dict( workers_per_gpu=2, train_dataloader=dict(samples_per_gpu=24), val_dataloader=dict(samples_per_gpu=1), test_dataloader=dict(samples_per_gpu=1), train=dict(type='BottomUpCocoWholeBodyDataset', ann_file=f'{data_root}/annotations/coco_wholebody_train_v1.0.json', img_prefix=f'{data_root}/train2017/', ...), val=dict(type='BottomUpCocoWholeBodyDataset', ann_file=f'{data_root}/annotations/coco_wholebody_val_v1.0.json', img_prefix=f'{data_root}/val2017/', ...), test=dict(...同上 val...), )

训练 batch size 为 24（单卡），推理 batch size 为 1（AssociativeEmbedding.forward_test中assert img.size(0) == 1）。数据集类型为BottomUpCocoWholeBodyDataset（bottom_up_coco_wholebody.py）。

4.2 HigherHRNet：多尺度高分辨率 head

与 HRNet 版本的区别在于 higherhrnet_w32_coco_wholebody_512x512.py 使用HigherHRNet的尺度感知设计：

data_cfg = dict( image_size=512, base_size=256, base_sigma=2, heatmap_size=[128, 256], # 双分辨率热图 num_scales=2, scale_aware_sigma=False, ) keypoint_head=dict( type='AEHigherResolutionHead', in_channels=32, num_joints=133, tag_per_joint=True, extra=dict(final_conv_kernel=1, ), num_deconv_layers=1, num_deconv_filters=[32], num_deconv_kernels=[4], num_basic_blocks=4, cat_output=[True], with_ae_loss=[True, False], # 原始分辨率分支算 AE 损失，反卷积分支不算 loss_keypoint=dict( type='MultiLossFactory', num_joints=133, num_stages=2, ae_loss_type='exp', with_ae_loss=[True, False], push_loss_factor=[0.001, 0.001], pull_loss_factor=[0.001, 0.001], with_heatmaps_loss=[True, True], heatmaps_loss_factor=[1.0, 1.0], supervise_empty=False)), test_cfg=dict( ... with_heatmaps=[True, True], with_ae=[True, False], ...)

• 通过一层步长 2 的反卷积（num_deconv_layers=1，kernel 4，后接 4 个 BasicBlock）把 HRNet 输出上采样到原始分辨率，得到[128, 256]双尺度热图，缓解小尺度人体关键点定位困难；
• cat_output=[True]表示反卷积输入会与前一阶段输出做通道拼接（见 AEHigherResolutionHead.forward）；
• 原始分辨率分支计算 AE 损失（with_ae_loss=[True, False]），反卷积分支只计算热图损失，测试时同理只对原始分辨率分支取 tag。

4.3 实验设置说明：+预训练策略

根据仓库中 hrnet_coco-wholebody.md 与 higherhrnet_coco-wholebody.md 的 Note：表中+表示模型先在原始 COCO 人体关键点数据集上预训练，再在 COCO-WholeBody 上微调，仓库作者发现该策略能带来更好的性能。

五、COCO-WholeBody 实验结果

以下结果均来自仓库文档，评测于 COCO-WholeBody v1.0 val 集，未使用多尺度测试。

5.1 HRNet + AESimpleHead（结果表）

5.2 HigherHRNet + AEHigherResolutionHead（结果表）

可见 HigherHRNet 借助高分辨率分支在 Face/Hand 等精细部位上普遍优于同级 HRNet。以上各配置的 ckpt 与 log 下载入口以仓库中对应模型的.md/.yml文件为准（如 hrnet_coco-wholebody.yml、higherhrnet_coco-wholebody.yml）。

六、训练、评测与推理实操

6.1 训练

使用仓库 tools/train.py 训练 AE 模型，参数与单卡/多卡用法如下：

# 单卡训练 python tools/train.py \ configs/wholebody/2d_kpt_sview_rgb_img/associative_embedding/coco-wholebody/hrnet_w32_coco_wholebody_512x512.py \ --work-dir work_dirs/hrnet_w32_coco_wholebody_512x512 # 多卡分布式训练（8 卡） bash tools/dist_train.sh \ configs/wholebody/2d_kpt_sview_rgb_img/associative_embedding/coco-wholebody/hrnet_w32_coco_wholebody_512x512.py \ 8 \ --work-dir work_dirs/hrnet_w32_coco_wholebody_512x512

前置条件：需按 data_preparation.md 将 COCO-WholeBody 数据置于data/coco/（coco_wholebody_train_v1.0.json、coco_wholebody_val_v1.0.json与对应图像目录）。

6.2 评测

使用 tools/test.py 在 val 集上复现仓库指标：

python tools/test.py \ configs/wholebody/2d_kpt_sview_rgb_img/associative_embedding/coco-wholebody/hrnet_w32_coco_wholebody_512x512.py \ ${CHECKPOINT_FILE} \ --out work_dirs/hrnet_w32_coco_wholebody_512x512/result.pkl

评测过程会按 Body / Foot / Face / Left Hand / Right Hand / WholeBody 六个分组分别输出 AP/AR 指标（见上文_do_python_keypoint_eval）。

6.3 图像推理与可视化

仓库提供 bottom-up 专用推理脚本 demo/bottom_up_img_demo.py（无需检测器，直接对整图推理）：

python demo/bottom_up_img_demo.py \ ${MMPOSE_CONFIG_FILE} ${MMPOSE_CHECKPOINT_FILE} \ --img-path ${IMG_PATH} \ --out-img-root vis_results \ --device=cuda:0 \ --kpt-thr 0.3 \ --pose-nms-thr 0.9

脚本内部调用inference_bottom_up_pose_model完成单图 bottom-up 推理，并通过vis_pose_result渲染 133 关键点姿态，--kpt-thr控制显示关键点置信度门槛，--pose-nms-thr为姿态级 OKS NMS 阈值。

七、源码级验证：单元测试

仓库测试 tests/test_models/test_bottom_up_forward.py 直接验证了 AE 检测器的完整链路：

• 以AssociativeEmbedding+ ResNet-18 backbone +AESimpleHead+MultiLossFactory构造模型（17 关键点）；
• 构造(1, 3, 256, 256)随机输入、目标热图(N, 17, H/32, W/32)、mask 与joints；
• 分别验证forward(return_loss=True)返回包含heatmap_loss/push_loss/pull_loss的损失字典，以及forward(return_loss=False)与forward_dummy的推理路径。

该测试从源码层面印证了本文所述"heatmap 损失 + push/pull AE 损失联合训练、tag 分组推理"的完整机制，也说明 AE 组件是仓库中可独立复用、被 CI 覆盖的核心模块。

八、小结

本文以 Associative Embedding README 为骨架，完整覆盖了 AE 的两阶段（先检测、后分组）思想与 tag 关联机制，并结合仓库源码深入剖析了 push/pull 损失、tag_per_joint、HeatmapParser 分组流水线、HRNet/HigherHRNet 两类 head 的差异、COCO-WholeBody 133 关键点配置、四组官方结果以及训练/评测/推理的完整命令。对于需要在拥挤场景或全身姿态任务中落地 bottom-up 方案的开发者，这套配置与代码即是现成的起点——只需准备数据、下载对应 checkpoint 即可复现仓库结果，或将AssociativeEmbedding检测器与AESimpleHead/AEHigherResolutionHead复用到自定义关键点数据集上。

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