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 热图估计框架。一、背景为什么在关键点估计中使用 PVTPVTPyramid Vision TransformerICCV 2021是一类无卷积主干的密集预测骨干它用四级 Transformer 编码器逐级下采样特征输出多尺度金字塔特征图与卷积骨干如 ResNet、HRNet在语义分割、目标检测、姿态估计等密集任务上保持相同的接口。PVTv2CVMJ 2022在此基础上做了三项改进更大的 Patch Embedding7×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.pytd-hm_pvtv2-b2_8xb64-210e_coco-256x192.py两模型论文引用信息来自原文档可直接用于 LaTeXinproceedings{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} }数据集为 COCOECCV 2014inproceedings{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字段中可以得到印证。完整结果表如下ArchInput SizeAPAP50AP75ARAR50配置文件pose_pvt-s256×1920.7140.8960.7940.7730.936td-hm_pvt-spose_pvtv2-b2256×1920.7370.9050.8120.7910.942td-hm_pvtv2-b2注原文档中每个模型还附有官方权重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、余弦之外的 MultiStepLRtrain_cfg dict(max_epochs210, val_interval10) # optimizer optim_wrapper dict(optimizerdict(typeAdam, lr5e-4)) # learning policy param_scheduler [ dict(typeLinearLR, begin0, end500, start_factor0.001, by_epochFalse), # warm-up dict(typeMultiStepLR, begin0, end210, milestones[170, 200], gamma0.1, by_epochTrue) ] # automatically scaling LR based on the actual training batch size auto_scale_lr dict(base_batch_size512) # hooks default_hooks dict(checkpointdict(save_bestcoco/AP, rulegreater))要点说明优化器Adam初始学习率 5e-4学习率策略前 500 个 iteration 线性热身从 0.001×lr 开始之后按 epoch 在第 170、200 epoch 各乘以 0.1auto_scale_lr以 512 为基准 batch8 卡 × 每卡 64当你更换 GPU 数量时框架会自动按比例缩放学习率避免小 batch 下 lr 过大checkpoint以 COCO AP 越大越好为准则保存最佳权重。3.2 Codec 与网络结构MSRA 热图 4 级金字塔# codec settings codec dict( typeMSRAHeatmap, input_size(192, 256), heatmap_size(48, 64), sigma2) # model settings norm_cfg dict(typeSyncBN, requires_gradTrue) model dict( typeTopdownPoseEstimator, data_preprocessordict( typePoseDataPreprocessor, mean[123.675, 116.28, 103.53], std[58.395, 57.12, 57.375], bgr_to_rgbTrue), backbonedict( typePyramidVisionTransformer, num_layers[3, 4, 6, 3], init_cfgdict(typePretrained, checkpoint.../pvt_small.pth)), # 官方 PVT 图像分类预训练权重 neckdict(typeFeatureMapProcessor, select_index3), headdict( typeHeatmapHead, in_channels512, out_channels17, lossdict(typeKeypointMSELoss, use_target_weightTrue), decodercodec), test_cfgdict(flip_testTrue, flip_modeheatmap, shift_heatmapTrue))从结构链路上看数据流为输入256×192w×h人体裁剪图用 ImageNet 均值方差做 RGB 归一化PoseDataPreprocessorBackbone输出 4 级特征金字塔最后一级通道数为 512、空间尺寸为输入的 1/16即 16×12Neck用FeatureMapProcessor(select_index3)从 4 级特征中只挑选第 4 级送入 Head——这是把多尺度骨干“降格”为单尺度输出的标准做法实现见 mmpose/models/necks/fmap_proc_neck.py还支持concat多尺度上采样拼接与scale_factor等选项若需要更高分辨率的热图可以在此调整HeadHeatmapHead将 512 通道特征上采样到 48×64 并回归 17 通道COCO 17 关键点热图损失为带目标点权重use_target_weightTrue的KeypointMSELoss测试增强flip_testTrue且flip_modeheatmap即对水平翻转的热图做翻转融合shift_heatmapTrue在反变换时修正 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(typeLoadImage), dict(typeGetBBoxCenterScale), dict(typeRandomFlip, directionhorizontal), dict(typeRandomHalfBody), dict(typeRandomBBoxTransform), dict(typeTopdownAffine, input_sizecodec[input_size]), dict(typeGenerateTarget, encodercodec), dict(typePackPoseInputs) ] val_pipeline [ dict(typeLoadImage), dict(typeGetBBoxCenterScale), dict(typeTopdownAffine, input_sizecodec[input_size]), dict(typePackPoseInputs) ] train_dataloader dict( batch_size64, num_workers2, persistent_workersTrue, samplerdict(typeDefaultSampler, shuffleTrue), datasetdict( typedataset_type, data_rootdata_root, data_modedata_mode, ann_fileannotations/person_keypoints_train2017.json, data_prefixdict(imgtrain2017/), pipelinetrain_pipeline)) val_dataloader dict( batch_size32, num_workers2, persistent_workersTrue, drop_lastFalse, samplerdict(typeDefaultSampler, shuffleFalse, round_upFalse), datasetdict( typedataset_type, data_rootdata_root, data_modedata_mode, ann_fileannotations/person_keypoints_val2017.json, bbox_filedata/coco/person_detection_results/ COCO_val2017_detections_AP_H_56_person.json, data_prefixdict(imgval2017/), test_modeTrue, pipelineval_pipeline)) test_dataloader val_dataloader val_evaluator dict( typeCocoMetric, ann_filedata_root annotations/person_keypoints_val2017.json) test_evaluator val_evaluator关键细节训练管线使用RandomFlip水平翻转、RandomHalfBody半身裁剪、RandomBBoxTransform中心/尺度扰动三类 Top-Down 常用增强GenerateTarget按 MSRA 方案把关键点渲染为高斯热图sigma2验证/测试集直接加载检测器预生成的人体框文件COCO_val2017_detections_AP_H_56_person.json这解释了结果表中“detector having human AP of 56.4”的前提——Top-Down 方案只负责在给定人体框内估计姿态框质量由检测器决定评估器为CocoMetric输出 AP/AP50/AP75/AR 等指标与save_bestcoco/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 规格参数默认值含义embed_dims64基础嵌入维度num_stages4金字塔阶段数num_layers[3, 4, 6, 3]各阶段 Transformer 层数num_heads[1, 2, 5, 8]各阶段注意力头数也决定各阶段通道 64/128/320/512patch_sizes/strides[4, 2, 2, 2]各阶段 patch 卷积核与步长逐级 2× 下采样sr_ratios[8, 4, 2, 1]各阶段空间缩减注意力SRA的缩减比mlp_ratios[8, 8, 4, 4]FFN 隐层相对嵌入维度的倍数drop_path_rate0.1随机深度stochastic depth总比率按层线性分配use_abs_pos_embedTrue是否使用绝对位置编码out_indices(0, 1, 2, 3)输出哪些阶段MMPose 默认取全部 4 级forwardpvt.py#L541-L554对每个阶段执行PatchEmbed得到 NLC 序列与hw_shape→ 依次过该阶段所有 encoder 层 → 可选阶段末 LayerNormnorm_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与 MixFFNSpatialReductionAttentionpvt.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 完成MixFFNpvt.py#L22-L98用 1×1 Conv 替代 Linearuse_convTrue时PVTv2在激活函数后插入 3×3 Depth-wise Conv 以编码局部位置信息AbsolutePositionEmbeddingpvt.py#L290-L354可学习的位置编码resize_pos_embed用双线性插值适配不同输入分辨率因此骨干支持任意输入尺寸测试用例中 33×33、112×137 等非整除尺寸均通过PVTEncoderLayerPre-Norm 结构norm1 → SRA → norm2 → MixFFNDropPath 按全局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], # 更大 patchV1: [4, 2, 2, 2] paddings[3, 1, 1, 1], # 配套 padding use_abs_pos_embedFalse, # 去掉绝对位置编码 norm_after_stageTrue, # 每阶段末尾加 LayerNorm use_conv_ffnTrue, # 使用 MixFFN含 3x3 DWConv **kwargs)对照两份配置PVTv2-b2 的 backbone 显式写了embed_dims64PVT-s 走默认值其余超参一致这正是 b2 规格的来源。4.4 官方预训练权重的键名转换PVT 的 ImageNet 预训练权重来自原 PVT 仓库键名与 mmcls 风格不同因此配置里init_cfg.typePretrained时init_weights 会调用pvt_convertmmpose/models/utils/ckpt_convert.py#L12-L82做键名映射丢弃分类head、norm.*、cls_token等与分类头相关的参数patch_embed{i}.proj→layers.{i-1}.0.projectionblock{i}.{j}→layers.{i-1}.1.{j(偏移)}其中偏移量由权重的 key 里是否含pos_embed自动推断有绝对位置编码时各 stage 层序列的 0 号位置是 pos_embed原仓库将 Q 与 KV 分开存储attn.q/attn.kv转换器用torch.cat([q, kv], dim0)拼成 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。如果你修改了骨干实现跑一遍该测试文件即可快速回归。五、如何复用与扩展这套配置更换输入分辨率同时修改codec的input_size与heatmap_size保持 4 倍关系、head的通道不变由于位置编码可插值骨干本身无需改动使用多尺度特征把neck换成select_index(0, 1, 2, 3), concatTrue的FeatureMapProcessor再调大head的in_channels为四阶段通道之和即可获得类似 FPN 的融合特征切换骨干只需替换backbone.type与对应超参如PyramidVisionTransformerV2TopdownPoseEstimator的接口对任意输出 NCHW 特征列表的骨干都适用迁移其他数据集保留 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.714pose_pvtv2-b2 取得 AP 0.737val201756.4 AP 检测器与同表其他卷积骨干构成一组可对照的基线。源码位于 mmpose/models/backbones/pvt.py配套测试与配置均已包含在仓库中可直接复现或二次开发。【免费下载链接】mmposeOpenMMLab Pose Estimation Toolbox and Benchmark.项目地址: https://gitcode.com/GitHub_Trending/mm/mmpose创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考