企业官网建设流程全解析

MMPose 实战：用 PVT / PVTv2 视觉 Transformer 骨干做 COCO Top-Down 关键点热图估计

【免费下载链接】mmposeOpenMMLab Pose Estimation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmpose

本文基于 MMPose 模型库中 COCO Top-Down 热图方案的 PVT 文档（pvt_coco.md），完整解读 pose_pvt-s 与 pose_pvtv2-b2 两个模型的基准结果、训练配置与数据流程，并结合源码剖析 PyramidVisionTransformer 骨干的四阶段金字塔结构、空间缩减注意力（SRA）与 MixFFN 实现，以及官方预训练权重的键名转换逻辑。读完后，你可以直接复现这两个模型的训练与评测，并理解 Transformer 骨干如何无缝嵌入 MMPose 的 Top-Down 热图估计框架。

一、背景：为什么在关键点估计中使用 PVT

PVT（Pyramid Vision Transformer，ICCV 2021）是一类无卷积主干的密集预测骨干：它用四级 Transformer 编码器逐级下采样特征，输出多尺度金字塔特征图，与卷积骨干（如 ResNet、HRNet）在语义分割、目标检测、姿态估计等密集任务上保持相同的接口。PVTv2（CVMJ 2022）在此基础上做了三项改进：更大的 Patch Embedding（7×7 起步）、混合前馈网络（MixFFN，用卷积替代线性层并插入 3×3 深度卷积编码位置信息），以及取消绝对位置编码、改用逐阶段 LayerNorm。

在 MMPose 中，PVT / PVTv2 被注册为骨干模块PyramidVisionTransformer与PyramidVisionTransformerV2（见 mmpose/models/backbones/pvt.py），并配有两份开箱即用的 COCO Top-Down 热图配置：

• td-hm_pvt-s_8xb64-210e_coco-256x192.py
• td-hm_pvtv2-b2_8xb64-210e_coco-256x192.py

两模型论文引用信息（来自原文档，可直接用于 LaTeX）：

@inproceedings{wang2021pyramid, title={Pyramid vision transformer: A versatile backbone for dense prediction without convolutions}, author={Wang, Wenhai and Xie, Enze and Li, Xiang and Fan, Deng-Ping and Song, Kaitao and Liang, Ding and Lu, Tong and Luo, Ping and Shao, Ling}, booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision}, pages={568--578}, year={2021} }

@article{wang2022pvt, title={PVT v2: Improved baselines with Pyramid Vision Transformer}, author={Wang, Wenhai and Xie, Enze and Li, Xiang and Fan, Deng-Ping and Song, Kaitao and Liang, Ding and Lu, Tong and Luo, Ping and Shao, Ling}, journal={Computational Visual Media}, pages={1--10}, year={2022}, publisher={Springer} }

数据集为 COCO（ECCV 2014）：

@inproceedings{lin2014microsoft, title={Microsoft coco: Common objects in context}, author={Lin, Tsung-Yi and Maire, Michael and Belongie, Serge and Hays, James and Perona, Pietro and Ramanan, Deva and Doll{\'a}r, Piotr and Zitnick, C Lawrence}, booktitle={European conference on computer vision}, pages={740--755}, year={2014}, organization={Springer} }

二、COCO val2017 基准结果

原文档给出的结果基于COCO val2017，且使用在 val2017 上人体检测 AP 为56.4的检测器生成人体框（对应配置中的COCO_val2017_detections_AP_H_56_person.json，这一点在配置的bbox_file字段中可以得到印证）。完整结果表如下：

注：原文档中每个模型还附有官方权重（ckpt）与训练日志（log）的 openmmlab 下载地址，属于外部链接，此处按规范略去；可从 MMPose 文档站或model-index.yml中按模型名检索。

可以看到 PVTv2-b2 相比 PVT-s 在 AP 上提升约 0.023，与 PVTv2 论文所述“更强基线”的结论一致。

三、训练配置详解

以 td-hm_pvt-s 配置 为主线（PVTv2 版本与其几乎相同，仅骨干不同，见第四节的差异表）。

3.1 训练策略：210 epoch、Adam 5e-4、余弦之外的 MultiStepLR

train_cfg = dict(max_epochs=210, val_interval=10) # optimizer optim_wrapper = dict(optimizer=dict(type='Adam', lr=5e-4)) # learning policy param_scheduler = [ dict(type='LinearLR', begin=0, end=500, start_factor=0.001, by_epoch=False), # warm-up dict(type='MultiStepLR', begin=0, end=210, milestones=[170, 200], gamma=0.1, by_epoch=True) ] # automatically scaling LR based on the actual training batch size auto_scale_lr = dict(base_batch_size=512) # hooks default_hooks = dict(checkpoint=dict(save_best='coco/AP', rule='greater'))

要点说明：

• 优化器：Adam，初始学习率 5e-4；
• 学习率策略：前 500 个 iteration 线性热身（从 0.001×lr 开始），之后按 epoch 在第 170、200 epoch 各乘以 0.1；
• auto_scale_lr：以 512 为基准 batch（8 卡 × 每卡 64），当你更换 GPU 数量时框架会自动按比例缩放学习率，避免小 batch 下 lr 过大；
• checkpoint：以 COCO AP 越大越好为准则保存最佳权重。

3.2 Codec 与网络结构：MSRA 热图 + 4 级金字塔

# codec settings codec = dict( type='MSRAHeatmap', input_size=(192, 256), heatmap_size=(48, 64), sigma=2) # model settings norm_cfg = dict(type='SyncBN', requires_grad=True) model = dict( type='TopdownPoseEstimator', data_preprocessor=dict( type='PoseDataPreprocessor', mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], bgr_to_rgb=True), backbone=dict( type='PyramidVisionTransformer', num_layers=[3, 4, 6, 3], init_cfg=dict(type='Pretrained', checkpoint='.../pvt_small.pth')), # 官方 PVT 图像分类预训练权重 neck=dict(type='FeatureMapProcessor', select_index=3), head=dict( type='HeatmapHead', in_channels=512, out_channels=17, loss=dict(type='KeypointMSELoss', use_target_weight=True), decoder=codec), test_cfg=dict(flip_test=True, flip_mode='heatmap', shift_heatmap=True))

从结构链路上看，数据流为：

1. 输入256×192（w×h）人体裁剪图，用 ImageNet 均值方差做 RGB 归一化（PoseDataPreprocessor）；
2. Backbone输出 4 级特征金字塔，最后一级通道数为 512、空间尺寸为输入的 1/16（即 16×12）；
3. Neck用FeatureMapProcessor(select_index=3)从 4 级特征中只挑选第 4 级送入 Head——这是把多尺度骨干“降格”为单尺度输出的标准做法，实现见 mmpose/models/necks/fmap_proc_neck.py，还支持concat（多尺度上采样拼接）与scale_factor等选项，若需要更高分辨率的热图可以在此调整；
4. HeadHeatmapHead将 512 通道特征上采样到 48×64 并回归 17 通道（COCO 17 关键点）热图，损失为带目标点权重（use_target_weight=True）的KeypointMSELoss；
5. 测试增强：flip_test=True且flip_mode='heatmap'，即对水平翻转的热图做翻转融合，shift_heatmap=True在反变换时修正 1 像素位移。

input_size=(192, 256)与heatmap_size=(48, 64)恰好满足 4 倍下采样关系，这是 MSRA 热图方案的标准配置。

3.3 数据管线与数据加载器

dataset_type = 'CocoDataset' data_mode = 'topdown' data_root = 'data/coco/' # pipelines train_pipeline = [ dict(type='LoadImage'), dict(type='GetBBoxCenterScale'), dict(type='RandomFlip', direction='horizontal'), dict(type='RandomHalfBody'), dict(type='RandomBBoxTransform'), dict(type='TopdownAffine', input_size=codec['input_size']), dict(type='GenerateTarget', encoder=codec), dict(type='PackPoseInputs') ] val_pipeline = [ dict(type='LoadImage'), dict(type='GetBBoxCenterScale'), dict(type='TopdownAffine', input_size=codec['input_size']), dict(type='PackPoseInputs') ] train_dataloader = dict( batch_size=64, num_workers=2, persistent_workers=True, sampler=dict(type='DefaultSampler', shuffle=True), dataset=dict( type=dataset_type, data_root=data_root, data_mode=data_mode, ann_file='annotations/person_keypoints_train2017.json', data_prefix=dict(img='train2017/'), pipeline=train_pipeline)) val_dataloader = dict( batch_size=32, num_workers=2, persistent_workers=True, drop_last=False, sampler=dict(type='DefaultSampler', shuffle=False, round_up=False), dataset=dict( type=dataset_type, data_root=data_root, data_mode=data_mode, ann_file='annotations/person_keypoints_val2017.json', bbox_file='data/coco/person_detection_results/' 'COCO_val2017_detections_AP_H_56_person.json', data_prefix=dict(img='val2017/'), test_mode=True, pipeline=val_pipeline)) test_dataloader = val_dataloader val_evaluator = dict( type='CocoMetric', ann_file=data_root + 'annotations/person_keypoints_val2017.json') test_evaluator = val_evaluator

关键细节：

• 训练管线使用RandomFlip（水平翻转）、RandomHalfBody（半身裁剪）、RandomBBoxTransform（中心/尺度扰动）三类 Top-Down 常用增强；
• GenerateTarget按 MSRA 方案把关键点渲染为高斯热图（sigma=2）；
• 验证/测试集直接加载检测器预生成的人体框文件COCO_val2017_detections_AP_H_56_person.json，这解释了结果表中“detector having human AP of 56.4”的前提——Top-Down 方案只负责在给定人体框内估计姿态，框质量由检测器决定；
• 评估器为CocoMetric，输出 AP/AP50/AP75/AR 等指标，与save_best='coco/AP'呼应。

3.4 训练与测试命令

数据准备（COCO train2017/val2017 与标注）可参考文档 prepare_datasets.md；准备好data/coco/后：

# 单卡训练（8 卡时按 auto_scale_lr 自动调整 lr） python tools/train.py configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_pvt-s_8xb64-210e_coco-256x192.py # 多卡训练 bash tools/dist_train.sh configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_pvt-s_8xb64-210e_coco-256x192.py 8 # 测试（需要训练好的权重路径，结果按 coco 协议评估） python tools/test.py configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_pvt-s_8xb64-210e_coco-256x192.py \ <path-to-ckpt> --out results.json --eval coco

四、源码深读：MMPose 中的 PVT 实现

4.1 骨干主类与四阶段金字塔

PyramidVisionTransformer 的核心是把输入切成 patch 序列，然后经过 4 个阶段的PatchEmbed + PVTEncoderLayer 堆叠，每阶段通道数按embed_dims × num_heads递增。其关键构造参数（默认值即 PVT-s 规格）：

forward（pvt.py#L541-L554）对每个阶段执行：PatchEmbed得到 NLC 序列与hw_shape→ 依次过该阶段所有 encoder 层 → 可选阶段末 LayerNorm（norm_after_stage）→ 转回 NCHW，按out_indices收集为多尺度特征列表。这与测试用例 tests/test_models/test_backbones/test_pvt.py 中的断言一致：对 32×32 输入，四阶段输出形状分别为(1, 64, 8, 8)、(1, 128, 4, 4)、(1, 320, 2, 2)、(1, 512, 1, 1)，即空间尺寸逐级 2 倍下采样、通道逐级 ×2 增长。

4.2 空间缩减注意力（SRA）与 MixFFN

• SpatialReductionAttention（pvt.py#L101-L197）：当sr_ratio > 1时，先用一个sr_ratio × sr_ratio的 stride 卷积对 Key/Value 序列做空间缩减再加 LayerNorm，而 Query 保留完整分辨率——这让注意力复杂度从 O(N²) 降到 O(N·N/s²)，是 PVT 能做高分辨率密集预测的关键。batch_first的布局转换在 forward 中通过 transpose 完成；
• MixFFN（pvt.py#L22-L98）：用 1×1 Conv 替代 Linear，use_conv=True时（PVTv2）在激活函数后插入 3×3 Depth-wise Conv 以编码局部位置信息；
• AbsolutePositionEmbedding（pvt.py#L290-L354）：可学习的位置编码，resize_pos_embed用双线性插值适配不同输入分辨率，因此骨干支持任意输入尺寸（测试用例中 33×33、112×137 等非整除尺寸均通过）；
• PVTEncoderLayer：Pre-Norm 结构，norm1 → SRA → norm2 → MixFFN，DropPath 按全局drop_path_rate线性分配到各层。

4.3 PVT 与 PVTv2 的差异

PyramidVisionTransformerV2 只是 PVT 的子类，用 5 个参数覆盖了 V2 论文的全部改进：

class PyramidVisionTransformerV2(PyramidVisionTransformer): def __init__(self, **kwargs): super().__init__( patch_sizes=[7, 3, 3, 3], # 更大 patch（V1: [4, 2, 2, 2]） paddings=[3, 1, 1, 1], # 配套 padding use_abs_pos_embed=False, # 去掉绝对位置编码 norm_after_stage=True, # 每阶段末尾加 LayerNorm use_conv_ffn=True, # 使用 MixFFN（含 3x3 DWConv） **kwargs)

对照两份配置：PVTv2-b2 的 backbone 显式写了embed_dims=64（PVT-s 走默认值），其余超参一致，这正是 b2 规格的来源。

4.4 官方预训练权重的键名转换

PVT 的 ImageNet 预训练权重来自原 PVT 仓库，键名与 mmcls 风格不同，因此配置里init_cfg.type='Pretrained'时，init_weights 会调用pvt_convert（mmpose/models/utils/ckpt_convert.py#L12-L82）做键名映射：

• 丢弃分类head、norm.*、cls_token等与分类头相关的参数；
• patch_embed{i}.proj→layers.{i-1}.0.projection；
• block{i}.{j}→layers.{i-1}.1.{j(+偏移)}，其中偏移量由权重的 key 里是否含pos_embed自动推断（有绝对位置编码时，各 stage 层序列的 0 号位置是 pos_embed）；
• 原仓库将 Q 与 KV 分开存储（attn.q/attn.kv），转换器用torch.cat([q, kv], dim=0)拼成 mmcvMultiheadAttention所需的in_proj权重；
• 4 阶段的pos_embed中额外包含 cls token，会被切掉（v[:, 1:, :]）。

这就是“从源码结构看”该实现能直接加载原始 PVT 权重而无需手动改 key 的原因。

4.4 单元测试验证

tests/test_models/test_backbones/test_pvt.py 覆盖了：

• 单个PVTEncoderLayer的前向形状（56×56 token 序列输入输出形状不变）；
• pretrain_img_size传入 3 元组时断言失败（必须为 int 或 2 元组）；
• 常规 224×224 与非整除尺寸（33×33、112×137）的多尺度输出形状；
• PVTv2 子类参数校验（pretrained必须为 str 或 None）。

如果你修改了骨干实现，跑一遍该测试文件即可快速回归。

五、如何复用与扩展这套配置

1. 更换输入分辨率：同时修改codec的input_size与heatmap_size（保持 4 倍关系）、head的通道不变；由于位置编码可插值，骨干本身无需改动；
2. 使用多尺度特征：把neck换成select_index=(0, 1, 2, 3), concat=True的FeatureMapProcessor，再调大head的in_channels为四阶段通道之和，即可获得类似 FPN 的融合特征；
3. 切换骨干：只需替换backbone.type与对应超参（如PyramidVisionTransformerV2），TopdownPoseEstimator的接口对任意输出 NCHW 特征列表的骨干都适用；
4. 迁移其他数据集：保留 backbone/neck/head 不变，替换dataset_type、ann_file、bbox_file与out_channels（关键点数）即可，流程与 COCO 完全一致。

六、小结

MMPose 的 PVT 方案展示了 Transformer 骨干在 Top-Down 姿态估计中的完整落地路径：官方 ImageNet 权重经pvt_convert键名转换后加载 → 四阶段金字塔输出取第 4 级 →HeatmapHead回归 48×64 热图 → 210 epoch Adam 训练、MultiStepLR 衰减、flip test 提升精度。pose_pvt-s 取得 AP 0.714，pose_pvtv2-b2 取得 AP 0.737（val2017，56.4 AP 检测器），与同表其他卷积骨干构成一组可对照的基线。源码位于 mmpose/models/backbones/pvt.py，配套测试与配置均已包含在仓库中，可直接复现或二次开发。

【免费下载链接】mmposeOpenMMLab Pose Estimation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmpose