简介AlphaPose轻量化SPPE训练代码面向希望复现并改进多人姿态估计、骨骼关键点检测流程的开发者和研究者核心是AlphaPose单人姿态估计网络的轻量训练工程可用于在自有数据集上完成模型训练、调试与二次开发。压缩包共261个文件、约301.67MB以Python源码和编译缓存为主配合YAML配置、CFG模型定义、PTH权重及CUDA/C扩展文件覆盖训练脚本、网络结构、数据处理与推理依赖等环节目录组织便于按模块查阅。已有2982人学习下载说明其在姿态估计入门与工程实践中具有一定参考价值。读者可据此了解轻量化SPPE的训练入口、配置方式和扩展实现结合示例数据组织自己的实验流程并为关键点检测任务的代码迁移、参数调整、显存与速度优化及问题排查提供可操作的参考。1. 为什么要把 SPPE 做小16 路摄像头逼出来的一个决定某园区实训楼有 16 路摄像头原本跑的是官方 AlphaPose 那套 ResNet-50 SPPE单张 T4 并发到 6 路就开始掉帧单帧延迟冲过 300 ms。瓶颈不在检测器——自上而下的多人姿态估计每检出一个人就要过一次单人网络画面里人一多骨骼关键点检测的计算量就线性翻倍。这份轻量化 SPPE 训练代码要解决的就是这一段把骨干压小、输入分辨率压小、上采样层数压少同时让精度掉得不难看。它适合三类人手上有 AlphaPose 推理链路但只会加载官方权重的人想在自己的业务数据集上重训单人姿态估计的人还有要在边缘盒子上跑多人姿态估计、显存和算力都紧的人。训练代码本身不长真正的门槛是里面那套默认约定。2. 工程骨架拆解一个检测框怎么变成 SPPE 的输入张量AlphaPose 这条链路分工很明确检测器负责“找出人”SPPE 负责“在人身上找点”最后 PoseNMS 负责“把重复的骨架去掉”。多数人卡住的地方不在网络结构而在于没搞清楚 SPPE 拿到的到底是什么形状的张量以及标注是怎么一步步变成高斯热力图的。2.1 从 bbox 到 256×192SPPE 在推理链路里的真实位置检测器输出的是[x1, y1, x2, y2, score]五元组。进 SPPE 之前这段通常会先按 1.25 倍做外扩因为人体的手腕脚踝经常贴着框边缘甚至出框不外扩就会在边界处丢点。外扩之后按框的长宽比做仿射变换把这块不规则区域采样成固定的 256×192宽×高再按 ImageNet 的均值方差归一化。网络输出的是 64×48 的热力图通道数等于关键点数量。取点的做法是在每个通道上取 argmax 拿到整数坐标再加上固定的 0.25 偏移做亚像素修正这是热力图回归里的常规操作因为 argmax 只能落在整数格上最后用仿射矩阵的逆变换把坐标映射回原图。这一步很容易翻车训练时用的仿射矩阵和推理时的不一致模型训得再好落回原图的点也是歪的。推理侧一般直接复用spp/models/下的模型定义只把权重文件换掉。所以训练目录里的结构改动一定要先确认推理脚本也跟着改了否则就是两套代码各说各话。提示SPPE 的输入尺寸、heatmap 步长、sigma 这三个参数必须成套看。改了输入不改 sigma热力图会变得过于尖锐或者糊成一片loss 曲线看着正常但 AP 上不去。2.2 目录结构与配置文件的对应关系这套训练代码的组织方式基本沿用了 AlphaPose 的布局配置、数据、模型、训练入口四层分得很清楚# 训练侧主干目录配置和代码是一一对应的 spp/ ├── train.py # 训练入口解析 cfg、建模型、跑 epoch ├── valid.py # 验证入口算 AP / PCK ├── core/ │ ├── config.py # 把 yaml 读成 cfg 对象支持 --cfg 覆盖 │ └── function.py # 分布式初始化、权重加载、checkpoint 存取 ├── dataset/ │ ├── __init__.py # JointDataset 基类增强和 target 生成都在这里 │ ├── coco.py # COCO 17 关键点的标注解析 │ └── coco_136.py # whole-body 136 关键点版本 ├── models/ │ ├── fastpose.py # FastPose 主体骨干 deconv final conv │ ├── mobilenet.py # 轻量骨干定义 │ ├── layers/DUC.py # 密集上采样卷积 │ └── builder.py # 按 cfg 里的 NAME 分发到具体模型 └── utils/ ├── transforms.py # get_affine_transform / affine_transform └── zipreader.py # 从 zip 里直接读图省一次解压 configs/coco/ ├── resnet/256x192_res50_lr1e-3_1x.yaml └── mobilenet/256x192_mbv3_lr1e-3_1x.yaml data/coco/ ├── images/train2017/ ├── images/val2017/ └── annotations/配置文件是整个训练过程唯一的开关面板改骨干、改数据集、改学习率都在这一个文件里完成代码里不该出现硬编码的路径和超参# configs/coco/mobilenet/256x192_mbv3_lr1e-3_1x.yaml DATASET: DATASET: coco # 决定用 spp/dataset 下哪个解析类 ROOT: data/coco/ IMG_PREFIX: images # 最终路径 ROOT IMG_PREFIX TRAIN_SET TRAIN_SET: train2017 TEST_SET: val2017 FLIP: true # 随机水平翻转依赖 flip_pairs 索引表 ROT_FACTOR: 40 # 旋转增强最大角度单位是度 SCALE_FACTOR: 0.3 # 缩放抖动幅度0.3 即 [0.7, 1.3] PART_SCALE_FACTOR: 0.25 # 半身裁剪的额外缩放 ROT_PROB: 0.3 # 各项增强的触发概率别都设成 1 SCALE_PROB: 0.3 PART_SCALE_PROB: 0.3 NUM_JOINTS_HALF_BODY: 8 # 半身增强时保留的关键点数量 SAMPLE_STRIDE: 5 # 验证集抽样步长全量验证太慢时可以调 MODEL: NAME: fastpose # builder.py 里的分发键 NUM_JOINTS: 17 # COCO 是 17whole-body 改 136 EXTRA: NUM_LAYERS: 50 # 只对 resnet 系列生效 PRETRAINED: models/pytorch/imagenet/resnet50.pth DECONV_WITH_BIAS: false SIGMA: 2 # 热力图高斯核半径跟着 heatmap_size 走 LOSS: USE_TARGET_WEIGHT: true # 是否用 target_weight 屏蔽无效点 TRAIN: BATCH_SIZE: 32 END_EPOCH: 140 OPTIMIZER: adam LR: 0.001 LR_STEP: [90, 120] # 到这两个 epoch 时学习率乘 LR_FACTOR LR_FACTOR: 0.1 WD: 0.0001 TEST: FLIP_TEST: true # 推理时原图 翻转图取平均 MODEL_FILE: models/sppe/mbv3_256x192.pthSIGMA和输入尺寸是绑死的256×192 的输入配 64×48 的热力图sigma 取 2 比较合适如果把输入压到 192×144热力图 48×36sigma 得降到 1.5 左右否则相邻关键点的高斯核会糊在一起肘和腕这类挨得近的点直接分不开。2.3 标注 JSON 到 db 记录center 与 scale 这套抽象数据集解析的核心不是读 JSON而是把每个标注框抽象成center scale。后续所有增强、仿射、热力图生成都在这两个量上做运算跟具体像素解耦。COCO 格式的转换大致是# spp/dataset/coco.py把标注框转成训练样本记录 def _get_db(self): gt_db [] for ann in self.coco.anns.values(): # 没有关键点标注的框直接丢弃人群区域、未标注的人 if ann.get(num_keypoints, 0) 0: continue x, y, w, h ann[bbox] center np.array([x w * 0.5, y h * 0.5], dtypenp.float32) # scale 不是面积是两个方向的边长单位是原图像素 scale np.array([w * 1.0, h * 1.0], dtypenp.float32) joints_3d np.zeros((self.num_joints, 3), dtypenp.float32) joints_3d_vis np.zeros((self.num_joints, 3), dtypenp.float32) for i in range(self.num_joints): joints_3d[i, 0:2] ann[keypoints][i * 3:i * 3 2] v ann[keypoints][i * 3 2] # v: 0 未标注, 1 标注但被遮挡, 2 可见。这里只监督 0 的 joints_3d_vis[i, 0:2] v 0 joints_3d_vis[i, 2] v 0 gt_db.append({ image: os.path.join(self.root, self.img_prefix, self.image_set, ann[file_name]), center: center, scale: scale, joints_3d: joints_3d, joints_3d_vis: joints_3d_vis, filename: ann[file_name], imgnum: ann[image_id], }) return gt_dbcenter和scale决定了仿射变换把原图哪一块采样到 256×192。scale用的是框的宽高而不是对角线长度这个约定贯穿整个仓库自定义数据集时千万别改成对角线否则整张图会被缩得只剩一个人头。joints_3d_vis的第三维看着多余其实是给翻转增强用的翻转之后左右关键点索引要对调可见性标记也要跟着换二维数组塞不下这个信息。3. 训练主循环增强、热力图目标与 loss 的完整走查网络结构只是训练的一半另一半在数据管线里。轻量模型本身容量小对数据增强和标签质量更敏感——同样的增强参数ResNet-50 吃得下MobileNetV3 可能直接训崩。3.1 数据增强三个概率参数决定了模型见过多少姿态spp/dataset/__init__.py里的JointDataset.__getitem__是所有数据集共用的COCO、Halpe、自定义数据集都继承它。核心逻辑是先在center scale空间上做旋转和缩放再把标注点用同一个仿射矩阵变换过去def __getitem__(self, idx): rec copy.deepcopy(self.db[idx]) data_numpy cv2.imread(rec[image], cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) if self.color_rgb: data_numpy cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB) joints rec[joints_3d] joints_vis rec[joints_3d_vis] c rec[center].copy() s rec[scale].copy() r 0 # 旋转增强randn 让小幅旋转更常见再用 clip 卡住上限 if self.rotation_factor 0 and np.random.rand() self.rot_prob: r np.clip(np.random.randn() * self.rotation_factor, -2 * self.rotation_factor, 2 * self.rotation_factor) # 缩放抖动0.3 的 factor 对应 [0.7, 1.3] 的倍率 if self.scale_factor 0 and np.random.rand() self.scale_prob: s s * np.clip(np.random.randn() * self.scale_factor 1, 1 - self.scale_factor, 1 self.scale_factor) # 半身增强模拟检测框只框到上半身的情况只在训练集开 if self.part_scale_factor 0 and np.random.rand() self.part_scale_prob: c[1] c[1] 0.5 * s[1] * np.random.uniform(-1, 1) s s * np.clip(np.random.randn() * self.part_scale_factor 1, 1 - self.part_scale_factor, 1 self.part_scale_factor) trans get_affine_transform(c, s, r, self.image_size) input cv2.warpAffine(data_numpy, trans, (self.image_size[0], self.image_size[1]), flagscv2.INTER_LINEAR) # 关键点用同一个 trans 变换保证和图像严格对齐 for i in range(self.num_joints): if joints_vis[i, 0] 0.0: joints[i, 0:2] affine_transform(joints[i, 0:2], trans) target, target_weight self.generate_target(joints, joints_vis) return input, torch.from_numpy(target), torch.from_numpy(target_weight), meta几个参数的实际影响ROT_FACTOR设 40 度是经验和精度之间的折中超过 45 度人体姿态会明显不合理模型学到的是扭曲的骨骼结构PART_SCALE_PROB一定要小于 0.5开太狠会让模型过度适应“只有半个人”的样本全身姿态的精度反而掉。SCALE_PROB和ROT_PROB都设成 1 是常见的翻车写法样本多样性反而下降因为每张图都被扰动到同一量级。归一化那段在示例里省掉了实际是在 collate 之前的 Compose 里做mean/std 用 ImageNet 那三个数。别改成 0.5因为骨干是拿 ImageNet 预训练权重初始化的分布对不上等于白加载。3.2 热力图目标生成sigma 与 target_weight 才是精度开关标注点变成目标张量这一步决定了模型在学什么。这里有两个设计一是用高斯核而不是单点 one-hot二是用target_weight把无效点从 loss 里摘掉def generate_target(self, joints, joints_vis): target_weight np.ones((self.num_joints, 1), dtypenp.float32) target_weight[:, 0] joints_vis[:, 0] target np.zeros((self.num_joints, self.heatmap_size[1], self.heatmap_size[0]), dtypenp.float32) tmp_size self.sigma * 3 # 高斯核半径取 3 sigma覆盖 99% 能量 for joint_id in range(self.num_joints): feat_stride np.asarray(self.image_size) / np.asarray(self.heatmap_size) mu_x int(joints[joint_id][0] / feat_stride[0] 0.5) mu_y int(joints[joint_id][1] / feat_stride[1] 0.5) # 关键点被增强裁到热力图外权重置 0这一路不参与反向传播 if mu_x 0 or mu_y 0 or \ mu_x self.heatmap_size[0] or mu_y self.heatmap_size[1]: target_weight[joint_id] 0 continue # 在 (mu_x, mu_y) 周围画高斯核再按热力图边界裁剪 ul [int(mu_x - tmp_size), int(mu_y - tmp_size)] br [int(mu_x tmp_size 1), int(mu_y tmp_size 1)] size 2 * tmp_size 1 x np.arange(0, size, 1, np.float32) y x[:, np.newaxis] x0 y0 size // 2 g np.exp(-((x - x0) ** 2 (y - y0) ** 2) / (2 * self.sigma ** 2)) g_x max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0] g_y max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1] img_x max(0, ul[0]), min(br[0], self.heatmap_size[0]) img_y max(0, ul[1]), min(br[1], self.heatmap_size[1]) target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] \ g[g_y[0]:g_y[1], g_x[0]:g_x[1]] return target, target_weighttarget_weight是这套代码里最重要也最容易被忽略的东西。轻量模型容量小一个错误的标签就能把某个关节的预测带偏好几个 epoch。半身增强和旋转之后总有相当比例的关键点落到视野外这些点必须被屏蔽掉否则网络会被迫去学习“在空白区域预测一个点”这种自相矛盾的目标。配套的 loss 就是带权重的 MSE写法很直白class HeatmapLoss(nn.Module): def forward(self, pred, gt, mask): assert pred.shape gt.shape loss ((pred - gt) ** 2) * mask # 逐维求均值而不是求和避免 num_joints 变化时 loss 量级跳变 return loss.mean(dim3).mean(dim2).mean(dim1).mean(dim0)注意最后是四次mean而不是sum。从 17 关键点换到 136 关键点的时候用sum会让 loss 直接放大 8 倍学习率不改就等着梯度爆炸。3.3 train.py 主循环与多卡启动主循环没太多花活关键是要处理好多级监督和 checkpointfor epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH): sampler.set_epoch(epoch) # 分布式下必须设否则每轮采样顺序一样 for i, (input, target, target_weight, meta) in enumerate(train_loader): input input.cuda(non_blockingTrue) target target.cuda(non_blockingTrue) target_weight target_weight.cuda(non_blockingTrue) output model(input) # [B, num_joints, H/4, W/4] if isinstance(output, list): output output[-1] # 多级输出只取最末级 loss criterion(output, target, target_weight) optimizer.zero_grad() loss.backward() optimizer.step() # 到 LR_STEP 指定的 epoch学习率乘 LR_FACTOR if epoch 1 in cfg.TRAIN.LR_STEP: for param_group in optimizer.param_groups: param_group[lr] param_group[lr] * cfg.TRAIN.LR_FACTOR if (epoch 1) % cfg.PRINT_FREQ 0 and is_main_process: save_checkpoint({epoch: epoch 1, state_dict: model.state_dict()}, os.path.join(cfg.OUTPUT_DIR, checkpoint.pth))启动用 torch 的分布式入口轻量模型单卡也能跑但 batch size 会小到 BatchNorm 统计量不稳还是建议两卡起步# 4 卡训练master_port 冲突时换一个 python -m torch.distributed.launch \ --nproc_per_node4 --master_port23456 \ spp/train.py \ --cfg configs/coco/mobilenet/256x192_mbv3_lr1e-3_1x.yaml \ --exp-id mbv3_256x192--exp-id会拼进OUTPUT_DIR和LOG_DIR方便同一个配置跑多组实验不互相覆盖。跑之前记得改GPUS那一行默认值经常和实际卡号对不上。4. 轻量化改造把 ResNet-50 换成 MobileNetV3 的四个改动点轻量化不是简单换个 backbone 就完事。SPPE 的输出是 1/4 分辨率的热力图而 MobileNetV3 的特征图步长是 32中间要靠上采样补回来。这块补得不好参数量是降了精度掉得比预期狠得多。4.1 骨干替换通道数从 2048 掉到 576 之后FastPose的原始实现里骨干最后一层输出 2048 通道ResNet-50接三层反卷积后通道数降到 256。换成 MobileNetV3-Small输出只有 576 通道反卷积的输入通道和中间通道都得跟着改# spp/models/mobilenet.py import torch.nn as nn from torchvision.models import mobilenet_v3_small class MobileNetV3SPPEBackbone(nn.Module): def __init__(self, width_mult1.0, pretrainedNone): super().__init__() net mobilenet_v3_small(weightsNone) self.features net.features # 输出 stride 32, 通道 576 if pretrained: # ImageNet 权重能加载多少加载多少最后一层分类头本来就不要 state torch.load(pretrained, map_locationcpu) self.features.load_state_dict( {k.replace(features., ): v for k, v in state.items() if k.startswith(features.)}, strictFalse) def forward(self, x): return self.features(x)width_mult是 MobileNet 系列自带的宽度缩放取 0.75 可以把通道数再砍四分之一参数量降到 1.8M 量级但热身阶段的 loss 会明显更抖。我一般先用 1.0 跑通全流程确认 loss 曲线形状正常再往回压。4.2 DUC 上采样把 stride 从 32 拉回 4三层 2 倍上采样如果用普通反卷积会引入棋盘格伪影热力图上表现为规则的小格子argmax 取点时会周期性偏移。DUCDense Upsampling Convolution的做法是先卷积把通道扩到out_channels * r²再用PixelShuffle重排成空间维度等价于可学习的亚像素插值# spp/models/layers/DUC.py class DUC(nn.Module): def __init__(self, in_planes, out_planes, upscale_factor2): super().__init__() self.conv nn.Conv2d(in_planes, out_planes * upscale_factor ** 2, kernel_size3, padding1, biasFalse) self.bn nn.BatchNorm2d(out_planes * upscale_factor ** 2) self.relu nn.ReLU(inplaceTrue) self.pixel_shuffle nn.PixelShuffle(upscale_factor) def forward(self, x): x self.relu(self.bn(self.conv(x))) return self.pixel_shuffle(x)模型主体就是骨干 三层 DUC 一层 1×1 卷积输出热力图class FastPoseLite(nn.Module): def __init__(self, num_joints17, mid_channels128, pretrainedNone): super().__init__() self.backbone MobileNetV3SPPEBackbone(pretrainedpretrained) self.deconv_layers self._make_deconv_layer( num_layers3, in_channels576, out_channelsmid_channels) self.final_layer nn.Conv2d(mid_channels, num_joints, kernel_size1) def _make_deconv_layer(self, num_layers, in_channels, out_channels): layers [] for _ in range(num_layers): layers.append(DUC(in_channels, out_channels, upscale_factor2)) in_channels out_channels return nn.Sequential(*layers) def forward(self, x): x self.backbone(x) # 256x192 - 8x6, 576 通道 x self.deconv_layers(x) # 8x6 - 64x48, 128 通道 return self.final_layer(x) # 64x48, num_joints 通道mid_channels从 256 降到 128 是这一步最划算的改动参数量省了将近 60%精度损失通常在 1 个点以内。再往下压到 64 就开始明显掉点了尤其是手腕脚踝这些细节区域感受野不够。4.3 参数量、计算量与实测耗时的三方对账参数量只是参考真正决定能不能上边缘盒子的是实际推理耗时。下面这组数字来自我自己的验证集只能看相对关系绝对值换个场景就不作数了骨干参数量输入 (W×H)单帧 GFLOPs验证集 APT4 单帧耗时ResNet-50约 26M256×192约 7.20.7124.8 msResNet-18约 13M256×192约 5.00.6883.1 msMobileNetV3-Small DUC(128)约 2.6M256×192约 0.90.6411.3 msMobileNetV3-Small DUC(128)约 2.6M192×144约 0.50.6020.9 ms从上往下看参数量掉了 10 倍耗时掉了将近 4 倍AP 掉 7 个点。这个交换比在 16 路摄像头的场景里是划算的——原来 6 路就卡现在 8 路还能留余量。但要注意最后一行分辨率压到 192×144 之后 AP 掉了 4 个点这部分损失比换骨干还难补因为小分辨率下远处的人只有几十个像素关键点根本分不开。如果业务场景里人在画面中占比普遍很小宁可保 256×192 换骨干也别动分辨率。4.4 训练策略跟着改学习率、warmup 与预训练权重结构和数据都改完之后原来的训练超参基本不能用了。轻量模型参数少、梯度噪声占比高Adam 配 1e-3 在 ResNet-50 上很稳在 MobileNetV3 上前 10 个 epoch 的 loss 会剧烈震荡。我一般会把学习率抬到 2e-3 到 3e-3加 5 个 epoch 的线性 warmup同时把END_EPOCH从 140 拉到 210。原因是小模型收敛慢但过拟合也慢多训 70 个 epoch 通常还能再榨出 1 到 2 个点。另外 ImageNet 预训练权重的加载要格外注意MobileNetV3 的features前面几层是 stride 2 的卷积如果预训练权重加载失败又没有报错strictFalse的副作用相当于从零开始训loss 曲线看起来“正常”但就是上不去。5. 踩坑排查训练不收敛、关键点丢失与导出失败的五种现场这套代码里埋了不少“看起来能跑其实错了”的默认值。下面五个是我在改轻量 SPPE 时真实遇到过的按现象、原因、解决三段写。5.1 loss 从第一个 epoch 就卡在 0.03 不掉现象训练启动正常loss 从 0.09 快速降到 0.03然后横盘 40 个 epoch 一动不动验证集 AP 一直在 0.05 附近。原因target_weight全被置 0 了。generate_target里判断关键点越界用的是heatmap_size但heatmap_size是从image_size除 4 算出来的如果配置里image_size写成了[256, 192]宽高顺序而代码按[H, W]用热力图尺寸就变成 48×64所有mu_x都超过 48权重全被清零。loss 依然能算出数因为全零目标配上全零权重MSE 恰好是 0 附近的一个极小值。解决在generate_target开头打一行断言assert self.heatmap_size[0] self.image_size[1]宽高顺序写反会立刻炸出来。另外训练前先单独跑一遍 dataloader把target_weight.sum()打印出来正常值应该接近batch_size * num_joints * 0.8。5.2 换 Halpe本文还有配套的精品资源点击获取 简介AlphaPose轻量化SPPE训练代码面向希望复现并改进多人姿态估计、骨骼关键点检测流程的开发者和研究者核心是AlphaPose单人姿态估计网络的轻量训练工程可用于在自有数据集上完成模型训练、调试与二次开发。压缩包共261个文件、约301.67MB以Python源码和编译缓存为主配合YAML配置、CFG模型定义、PTH权重及CUDA/C扩展文件覆盖训练脚本、网络结构、数据处理与推理依赖等环节目录组织便于按模块查阅。已有2982人学习下载说明其在姿态估计入门与工程实践中具有一定参考价值。读者可据此了解轻量化SPPE的训练入口、配置方式和扩展实现结合示例数据组织自己的实验流程并为关键点检测任务的代码迁移、参数调整、显存与速度优化及问题排查提供可操作的参考。1. 为什么要把 SPPE 做小16 路摄像头逼出来的一个决定某园区实训楼有 16 路摄像头原本跑的是官方 AlphaPose 那套 ResNet-50 SPPE单张 T4 并发到 6 路就开始掉帧单帧延迟冲过 300 ms。瓶颈不在检测器——自上而下的多人姿态估计每检出一个人就要过一次单人网络画面里人一多骨骼关键点检测的计算量就线性翻倍。这份轻量化 SPPE 训练代码要解决的就是这一段把骨干压小、输入分辨率压小、上采样层数压少同时让精度掉得不难看。它适合三类人手上有 AlphaPose 推理链路但只会加载官方权重的人想在自己的业务数据集上重训单人姿态估计的人还有要在边缘盒子上跑多人姿态估计、显存和算力都紧的人。训练代码本身不长真正的门槛是里面那套默认约定。2. 工程骨架拆解一个检测框怎么变成 SPPE 的输入张量AlphaPose 这条链路分工很明确检测器负责“找出人”SPPE 负责“在人身上找点”最后 PoseNMS 负责“把重复的骨架去掉”。多数人卡住的地方不在网络结构而在于没搞清楚 SPPE 拿到的到底是什么形状的张量以及标注是怎么一步步变成高斯热力图的。2.1 从 bbox 到 256×192SPPE 在推理链路里的真实位置检测器输出的是[x1, y1, x2, y2, score]五元组。进 SPPE 之前这段通常会先按 1.25 倍做外扩因为人体的手腕脚踝经常贴着框边缘甚至出框不外扩就会在边界处丢点。外扩之后按框的长宽比做仿射变换把这块不规则区域采样成固定的 256×192宽×高再按 ImageNet 的均值方差归一化。网络输出的是 64×48 的热力图通道数等于关键点数量。取点的做法是在每个通道上取 argmax 拿到整数坐标再加上固定的 0.25 偏移做亚像素修正这是热力图回归里的常规操作因为 argmax 只能落在整数格上最后用仿射矩阵的逆变换把坐标映射回原图。这一步很容易翻车训练时用的仿射矩阵和推理时的不一致模型训得再好落回原图的点也是歪的。推理侧一般直接复用spp/models/下的模型定义只把权重文件换掉。所以训练目录里的结构改动一定要先确认推理脚本也跟着改了否则就是两套代码各说各话。提示SPPE 的输入尺寸、heatmap 步长、sigma 这三个参数必须成套看。改了输入不改 sigma热力图会变得过于尖锐或者糊成一片loss 曲线看着正常但 AP 上不去。2.2 目录结构与配置文件的对应关系这套训练代码的组织方式基本沿用了 AlphaPose 的布局配置、数据、模型、训练入口四层分得很清楚# 训练侧主干目录配置和代码是一一对应的 spp/ ├── train.py # 训练入口解析 cfg、建模型、跑 epoch ├── valid.py # 验证入口算 AP / PCK ├── core/ │ ├── config.py # 把 yaml 读成 cfg 对象支持 --cfg 覆盖 │ └── function.py # 分布式初始化、权重加载、checkpoint 存取 ├── dataset/ │ ├── __init__.py # JointDataset 基类增强和 target 生成都在这里 │ ├── coco.py # COCO 17 关键点的标注解析 │ └── coco_136.py # whole-body 136 关键点版本 ├── models/ │ ├── fastpose.py # FastPose 主体骨干 deconv final conv │ ├── mobilenet.py # 轻量骨干定义 │ ├── layers/DUC.py # 密集上采样卷积 │ └── builder.py # 按 cfg 里的 NAME 分发到具体模型 └── utils/ ├── transforms.py # get_affine_transform / affine_transform └── zipreader.py # 从 zip 里直接读图省一次解压 configs/coco/ ├── resnet/256x192_res50_lr1e-3_1x.yaml └── mobilenet/256x192_mbv3_lr1e-3_1x.yaml data/coco/ ├── images/train2017/ ├── images/val2017/ └── annotations/配置文件是整个训练过程唯一的开关面板改骨干、改数据集、改学习率都在这一个文件里完成代码里不该出现硬编码的路径和超参# configs/coco/mobilenet/256x192_mbv3_lr1e-3_1x.yaml DATASET: DATASET: coco # 决定用 spp/dataset 下哪个解析类 ROOT: data/coco/ IMG_PREFIX: images # 最终路径 ROOT IMG_PREFIX TRAIN_SET TRAIN_SET: train2017 TEST_SET: val2017 FLIP: true # 随机水平翻转依赖 flip_pairs 索引表 ROT_FACTOR: 40 # 旋转增强最大角度单位是度 SCALE_FACTOR: 0.3 # 缩放抖动幅度0.3 即 [0.7, 1.3] PART_SCALE_FACTOR: 0.25 # 半身裁剪的额外缩放 ROT_PROB: 0.3 # 各项增强的触发概率别都设成 1 SCALE_PROB: 0.3 PART_SCALE_PROB: 0.3 NUM_JOINTS_HALF_BODY: 8 # 半身增强时保留的关键点数量 SAMPLE_STRIDE: 5 # 验证集抽样步长全量验证太慢时可以调 MODEL: NAME: fastpose # builder.py 里的分发键 NUM_JOINTS: 17 # COCO 是 17whole-body 改 136 EXTRA: NUM_LAYERS: 50 # 只对 resnet 系列生效 PRETRAINED: models/pytorch/imagenet/resnet50.pth DECONV_WITH_BIAS: false SIGMA: 2 # 热力图高斯核半径跟着 heatmap_size 走 LOSS: USE_TARGET_WEIGHT: true # 是否用 target_weight 屏蔽无效点 TRAIN: BATCH_SIZE: 32 END_EPOCH: 140 OPTIMIZER: adam LR: 0.001 LR_STEP: [90, 120] # 到这两个 epoch 时学习率乘 LR_FACTOR LR_FACTOR: 0.1 WD: 0.0001 TEST: FLIP_TEST: true # 推理时原图 翻转图取平均 MODEL_FILE: models/sppe/mbv3_256x192.pthSIGMA和输入尺寸是绑死的256×192 的输入配 64×48 的热力图sigma 取 2 比较合适如果把输入压到 192×144热力图 48×36sigma 得降到 1.5 左右否则相邻关键点的高斯核会糊在一起肘和腕这类挨得近的点直接分不开。2.3 标注 JSON 到 db 记录center 与 scale 这套抽象数据集解析的核心不是读 JSON而是把每个标注框抽象成center scale。后续所有增强、仿射、热力图生成都在这两个量上做运算跟具体像素解耦。COCO 格式的转换大致是# spp/dataset/coco.py把标注框转成训练样本记录 def _get_db(self): gt_db [] for ann in self.coco.anns.values(): # 没有关键点标注的框直接丢弃人群区域、未标注的人 if ann.get(num_keypoints, 0) 0: continue x, y, w, h ann[bbox] center np.array([x w * 0.5, y h * 0.5], dtypenp.float32) # scale 不是面积是两个方向的边长单位是原图像素 scale np.array([w * 1.0, h * 1.0], dtypenp.float32) joints_3d np.zeros((self.num_joints, 3), dtypenp.float32) joints_3d_vis np.zeros((self.num_joints, 3), dtypenp.float32) for i in range(self.num_joints): joints_3d[i, 0:2] ann[keypoints][i * 3:i * 3 2] v ann[keypoints][i * 3 2] # v: 0 未标注, 1 标注但被遮挡, 2 可见。这里只监督 0 的 joints_3d_vis[i, 0:2] v 0 joints_3d_vis[i, 2] v 0 gt_db.append({ image: os.path.join(self.root, self.img_prefix, self.image_set, ann[file_name]), center: center, scale: scale, joints_3d: joints_3d, joints_3d_vis: joints_3d_vis, filename: ann[file_name], imgnum: ann[image_id], }) return gt_dbcenter和scale决定了仿射变换把原图哪一块采样到 256×192。scale用的是框的宽高而不是对角线长度这个约定贯穿整个仓库自定义数据集时千万别改成对角线否则整张图会被缩得只剩一个人头。joints_3d_vis的第三维看着多余其实是给翻转增强用的翻转之后左右关键点索引要对调可见性标记也要跟着换二维数组塞不下这个信息。3. 训练主循环增强、热力图目标与 loss 的完整走查网络结构只是训练的一半另一半在数据管线里。轻量模型本身容量小对数据增强和标签质量更敏感——同样的增强参数ResNet-50 吃得下MobileNetV3 可能直接训崩。3.1 数据增强三个概率参数决定了模型见过多少姿态spp/dataset/__init__.py里的JointDataset.__getitem__是所有数据集共用的COCO、Halpe、自定义数据集都继承它。核心逻辑是先在center scale空间上做旋转和缩放再把标注点用同一个仿射矩阵变换过去def __getitem__(self, idx): rec copy.deepcopy(self.db[idx]) data_numpy cv2.imread(rec[image], cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) if self.color_rgb: data_numpy cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB) joints rec[joints_3d] joints_vis rec[joints_3d_vis] c rec[center].copy() s rec[scale].copy() r 0 # 旋转增强randn 让小幅旋转更常见再用 clip 卡住上限 if self.rotation_factor 0 and np.random.rand() self.rot_prob: r np.clip(np.random.randn() * self.rotation_factor, -2 * self.rotation_factor, 2 * self.rotation_factor) # 缩放抖动0.3 的 factor 对应 [0.7, 1.3] 的倍率 if self.scale_factor 0 and np.random.rand() self.scale_prob: s s * np.clip(np.random.randn() * self.scale_factor 1, 1 - self.scale_factor, 1 self.scale_factor) # 半身增强模拟检测框只框到上半身的情况只在训练集开 if self.part_scale_factor 0 and np.random.rand() self.part_scale_prob: c[1] c[1] 0.5 * s[1] * np.random.uniform(-1, 1) s s * np.clip(np.random.randn() * self.part_scale_factor 1, 1 - self.part_scale_factor, 1 self.part_scale_factor) trans get_affine_transform(c, s, r, self.image_size) input cv2.warpAffine(data_numpy, trans, (self.image_size[0], self.image_size[1]), flagscv2.INTER_LINEAR) # 关键点用同一个 trans 变换保证和图像严格对齐 for i in range(self.num_joints): if joints_vis[i, 0] 0.0: joints[i, 0:2] affine_transform(joints[i, 0:2], trans) target, target_weight self.generate_target(joints, joints_vis) return input, torch.from_numpy(target), torch.from_numpy(target_weight), meta几个参数的实际影响ROT_FACTOR设 40 度是经验和精度之间的折中超过 45 度人体姿态会明显不合理模型学到的是扭曲的骨骼结构PART_SCALE_PROB一定要小于 0.5开太狠会让模型过度适应“只有半个人”的样本全身姿态的精度反而掉。SCALE_PROB和ROT_PROB都设成 1 是常见的翻车写法样本多样性反而下降因为每张图都被扰动到同一量级。归一化那段在示例里省掉了实际是在 collate 之前的 Compose 里做mean/std 用 ImageNet 那三个数。别改成 0.5因为骨干是拿 ImageNet 预训练权重初始化的分布对不上等于白加载。3.2 热力图目标生成sigma 与 target_weight 才是精度开关标注点变成目标张量这一步决定了模型在学什么。这里有两个设计一是用高斯核而不是单点 one-hot二是用target_weight把无效点从 loss 里摘掉def generate_target(self, joints, joints_vis): target_weight np.ones((self.num_joints, 1), dtypenp.float32) target_weight[:, 0] joints_vis[:, 0] target np.zeros((self.num_joints, self.heatmap_size[1], self.heatmap_size[0]), dtypenp.float32) tmp_size self.sigma * 3 # 高斯核半径取 3 sigma覆盖 99% 能量 for joint_id in range(self.num_joints): feat_stride np.asarray(self.image_size) / np.asarray(self.heatmap_size) mu_x int(joints[joint_id][0] / feat_stride[0] 0.5) mu_y int(joints[joint_id][1] / feat_stride[1] 0.5) # 关键点被增强裁到热力图外权重置 0这一路不参与反向传播 if mu_x 0 or mu_y 0 or \ mu_x self.heatmap_size[0] or mu_y self.heatmap_size[1]: target_weight[joint_id] 0 continue # 在 (mu_x, mu_y) 周围画高斯核再按热力图边界裁剪 ul [int(mu_x - tmp_size), int(mu_y - tmp_size)] br [int(mu_x tmp_size 1), int(mu_y tmp_size 1)] size 2 * tmp_size 1 x np.arange(0, size, 1, np.float32) y x[:, np.newaxis] x0 y0 size // 2 g np.exp(-((x - x0) ** 2 (y - y0) ** 2) / (2 * self.sigma ** 2)) g_x max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0] g_y max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1] img_x max(0, ul[0]), min(br[0], self.heatmap_size[0]) img_y max(0, ul[1]), min(br[1], self.heatmap_size[1]) target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] \ g[g_y[0]:g_y[1], g_x[0]:g_x[1]] return target, target_weighttarget_weight是这套代码里最重要也最容易被忽略的东西。轻量模型容量小一个错误的标签就能把某个关节的预测带偏好几个 epoch。半身增强和旋转之后总有相当比例的关键点落到视野外这些点必须被屏蔽掉否则网络会被迫去学习“在空白区域预测一个点”这种自相矛盾的目标。配套的 loss 就是带权重的 MSE写法很直白class HeatmapLoss(nn.Module): def forward(self, pred, gt, mask): assert pred.shape gt.shape loss ((pred - gt) ** 2) * mask # 逐维求均值而不是求和避免 num_joints 变化时 loss 量级跳变 return loss.mean(dim3).mean(dim2).mean(dim1).mean(dim0)注意最后是四次mean而不是sum。从 17 关键点换到 136 关键点的时候用sum会让 loss 直接放大 8 倍学习率不改就等着梯度爆炸。3.3 train.py 主循环与多卡启动主循环没太多花活关键是要处理好多级监督和 checkpointfor epoch in range(begin_epoch, cfg.TRAIN.END_EPOCH): sampler.set_epoch(epoch) # 分布式下必须设否则每轮采样顺序一样 for i, (input, target, target_weight, meta) in enumerate(train_loader): input input.cuda(non_blockingTrue) target target.cuda(non_blockingTrue) target_weight target_weight.cuda(non_blockingTrue) output model(input) # [B, num_joints, H/4, W/4] if isinstance(output, list): output output[-1] # 多级输出只取最末级 loss criterion(output, target, target_weight) optimizer.zero_grad() loss.backward() optimizer.step() # 到 LR_STEP 指定的 epoch学习率乘 LR_FACTOR if epoch 1 in cfg.TRAIN.LR_STEP: for param_group in optimizer.param_groups: param_group[lr] param_group[lr] * cfg.TRAIN.LR_FACTOR if (epoch 1) % cfg.PRINT_FREQ 0 and is_main_process: save_checkpoint({epoch: epoch 1, state_dict: model.state_dict()}, os.path.join(cfg.OUTPUT_DIR, checkpoint.pth))启动用 torch 的分布式入口轻量模型单卡也能跑但 batch size 会小到 BatchNorm 统计量不稳还是建议两卡起步# 4 卡训练master_port 冲突时换一个 python -m torch.distributed.launch \ --nproc_per_node4 --master_port23456 \ spp/train.py \ --cfg configs/coco/mobilenet/256x192_mbv3_lr1e-3_1x.yaml \ --exp-id mbv3_256x192--exp-id会拼进OUTPUT_DIR和LOG_DIR方便同一个配置跑多组实验不互相覆盖。跑之前记得改GPUS那一行默认值经常和实际卡号对不上。4. 轻量化改造把 ResNet-50 换成 MobileNetV3 的四个改动点轻量化不是简单换个 backbone 就完事。SPPE 的输出是 1/4 分辨率的热力图而 MobileNetV3 的特征图步长是 32中间要靠上采样补回来。这块补得不好参数量是降了精度掉得比预期狠得多。4.1 骨干替换通道数从 2048 掉到 576 之后FastPose的原始实现里骨干最后一层输出 2048 通道ResNet-50接三层反卷积后通道数降到 256。换成 MobileNetV3-Small输出只有 576 通道反卷积的输入通道和中间通道都得跟着改# spp/models/mobilenet.py import torch.nn as nn from torchvision.models import mobilenet_v3_small class MobileNetV3SPPEBackbone(nn.Module): def __init__(self, width_mult1.0, pretrainedNone): super().__init__() net mobilenet_v3_small(weightsNone) self.features net.features # 输出 stride 32, 通道 576 if pretrained: # ImageNet 权重能加载多少加载多少最后一层分类头本来就不要 state torch.load(pretrained, map_locationcpu) self.features.load_state_dict( {k.replace(features., ): v for k, v in state.items() if k.startswith(features.)}, strictFalse) def forward(self, x): return self.features(x)width_mult是 MobileNet 系列自带的宽度缩放取 0.75 可以把通道数再砍四分之一参数量降到 1.8M 量级但热身阶段的 loss 会明显更抖。我一般先用 1.0 跑通全流程确认 loss 曲线形状正常再往回压。4.2 DUC 上采样把 stride 从 32 拉回 4三层 2 倍上采样如果用普通反卷积会引入棋盘格伪影热力图上表现为规则的小格子argmax 取点时会周期性偏移。DUCDense Upsampling Convolution的做法是先卷积把通道扩到out_channels * r²再用PixelShuffle重排成空间维度等价于可学习的亚像素插值# spp/models/layers/DUC.py class DUC(nn.Module): def __init__(self, in_planes, out_planes, upscale_factor2): super().__init__() self.conv nn.Conv2d(in_planes, out_planes * upscale_factor ** 2, kernel_size3, padding1, biasFalse) self.bn nn.BatchNorm2d(out_planes * upscale_factor ** 2) self.relu nn.ReLU(inplaceTrue) self.pixel_shuffle nn.PixelShuffle(upscale_factor) def forward(self, x): x self.relu(self.bn(self.conv(x))) return self.pixel_shuffle(x)模型主体就是骨干 三层 DUC 一层 1×1 卷积输出热力图class FastPoseLite(nn.Module): def __init__(self, num_joints17, mid_channels128, pretrainedNone): super().__init__() self.backbone MobileNetV3SPPEBackbone(pretrainedpretrained) self.deconv_layers self._make_deconv_layer( num_layers3, in_channels576, out_channelsmid_channels) self.final_layer nn.Conv2d(mid_channels, num_joints, kernel_size1) def _make_deconv_layer(self, num_layers, in_channels, out_channels): layers [] for _ in range(num_layers): layers.append(DUC(in_channels, out_channels, upscale_factor2)) in_channels out_channels return nn.Sequential(*layers) def forward(self, x): x self.backbone(x) # 256x192 - 8x6, 576 通道 x self.deconv_layers(x) # 8x6 - 64x48, 128 通道 return self.final_layer(x) # 64x48, num_joints 通道mid_channels从 256 降到 128 是这一步最划算的改动参数量省了将近 60%精度损失通常在 1 个点以内。再往下压到 64 就开始明显掉点了尤其是手腕脚踝这些细节区域感受野不够。4.3 参数量、计算量与实测耗时的三方对账参数量只是参考真正决定能不能上边缘盒子的是实际推理耗时。下面这组数字来自我自己的验证集只能看相对关系绝对值换个场景就不作数了骨干参数量输入 (W×H)单帧 GFLOPs验证集 APT4 单帧耗时ResNet-50约 26M256×192约 7.20.7124.8 msResNet-18约 13M256×192约 5.00.6883.1 msMobileNetV3-Small DUC(128)约 2.6M256×192约 0.90.6411.3 msMobileNetV3-Small DUC(128)约 2.6M192×144约 0.50.6020.9 ms从上往下看参数量掉了 10 倍耗时掉了将近 4 倍AP 掉 7 个点。这个交换比在 16 路摄像头的场景里是划算的——原来 6 路就卡现在 8 路还能留余量。但要注意最后一行分辨率压到 192×144 之后 AP 掉了 4 个点这部分损失比换骨干还难补因为小分辨率下远处的人只有几十个像素关键点根本分不开。如果业务场景里人在画面中占比普遍很小宁可保 256×192 换骨干也别动分辨率。4.4 训练策略跟着改学习率、warmup 与预训练权重结构和数据都改完之后原来的训练超参基本不能用了。轻量模型参数少、梯度噪声占比高Adam 配 1e-3 在 ResNet-50 上很稳在 MobileNetV3 上前 10 个 epoch 的 loss 会剧烈震荡。我一般会把学习率抬到 2e-3 到 3e-3加 5 个 epoch 的线性 warmup同时把END_EPOCH从 140 拉到 210。原因是小模型收敛慢但过拟合也慢多训 70 个 epoch 通常还能再榨出 1 到 2 个点。另外 ImageNet 预训练权重的加载要格外注意MobileNetV3 的features前面几层是 stride 2 的卷积如果预训练权重加载失败又没有报错strictFalse的副作用相当于从零开始训loss 曲线看起来“正常”但就是上不去。5. 踩坑排查训练不收敛、关键点丢失与导出失败的五种现场这套代码里埋了不少“看起来能跑其实错了”的默认值。下面五个是我在改轻量 SPPE 时真实遇到过的按现象、原因、解决三段写。5.1 loss 从第一个 epoch 就卡在 0.03 不掉现象训练启动正常loss 从 0.09 快速降到 0.03然后横盘 40 个 epoch 一动不动验证集 AP 一直在 0.05 附近。原因target_weight全被置 0 了。generate_target里判断关键点越界用的是heatmap_size但heatmap_size是从image_size除 4 算出来的如果配置里image_size写成了[256, 192]宽高顺序而代码按[H, W]用热力图尺寸就变成 48×64所有mu_x都超过 48权重全被清零。loss 依然能算出数因为全零目标配上全零权重MSE 恰好是 0 附近的一个极小值。解决在generate_target开头打一行断言assert self.heatmap_size[0] self.image_size[1]宽高顺序写反会立刻炸出来。另外训练前先单独跑一遍 dataloader把target_weight.sum()打印出来正常值应该接近batch_size * num_joints * 0.8。5.2 换 Halpe本文还有配套的精品资源点击获取