PyTorch关键点检测实战：从自建数据集到模型训练与推理-北京尧图网络

简介这份资源面向具备一定深度学习基础、希望动手实践关键点检测的开发者与研究人员围绕PyTorch中的Keypoint R-CNN展开解决在自建数据集上训练关键点检测模型时数据组织、模型配置与训练评估等实际问题。压缩包共116个文件约8.55MB包含39张jpg示例图像、34个txt与31个json标注文件、8个py脚本、2个ipynb笔记本及1个md说明覆盖数据准备、标签转换、模型训练到推理部署的完整流程。已有166人学习下载。读者可借助其中的Notebook与脚本理解如何将自建标注转换为模型可读格式、调整关键点输出头与超参数、组合位置损失与分类损失进行训练并通过验证集监控性能、导出ONNX用于多平台推理适合作为姿态估计或面部特征点检测项目的实操参考。1. 从一堆标注图到能跑的关键点模型这份 PyTorch 资源到底给了什么手里有一批自己标的图每张图上几个关键点想训一个能预测这些点的模型翻了一圈发现分类、检测的教程满地都是关键点检测的完整流程却少得可怜。这份keypoint_detection.zip就是冲着这个缺口来的——它围绕 PyTorch 里现成的keypointrcnn_resnet50_fpn搭了一条从自建数据集到训练、推理的链路。关键点检测和普通目标检测最大的区别在于它输出的不是框而是框内每个点的坐标所以数据标注格式、损失计算、后处理全都不一样。这份资源适合已经会用 PyTorch 跑通分类或检测、现在要往关键点方向走的人也适合手上有自定义关键点需求比如零件定位、姿态估计、面部特征点但不想从零写模型的从业者。它不解决标注工具的问题也不承诺开箱即用的精度给的是能改、能调、能复现的骨架。2. 关键点检测的数据集长什么样从标注格式到 Dataset 类2.1 为什么关键点标注不能直接套检测的格式目标检测的标注是每张图若干个框加类别关键点检测要在每个框里再挂一组点每个点有 x、y 和可见性状态。PyTorch 的keypointrcnn_resnet50_fpn期望的 target 结构里keypoints是一个形状为[num_objects, num_keypoints, 3]的张量最后一维是(x, y, visibility)。visibility 通常用 0 表示未标注、1 表示标注但不可见、2 表示可见这个约定来自 COCO模型内部会据此决定哪些点参与损失计算。很多人第一次训关键点模型时 loss 一直不降八成是把 visibility 全填了 0 或者把坐标归一化搞反了。常见做法是标注阶段就用 COCO 格式存 JSON每张图的annotations里带keypoints数组顺序是[x1, y1, v1, x2, y2, v2, ...]这样后面转 Dataset 时不用再猜顺序。2.2 自建数据集的目录结构与 JSON 字段假设你已经用某个标注工具导出了 COCO 格式的 JSON目录一般长这样dataset/ images/ img_0001.jpg img_0002.jpg annotations/ instances_train.json instances_val.jsoninstances_train.json里几个关键字段必须对齐images列表里每项有id、file_name、width、heightannotations列表里每项有image_id、category_id、bbox[x, y, w, h]、keypoints、num_keypointscategories里定义你的关键点类别keypoints字段列出每个点的名字和顺序。顺序一旦定了就不能改因为模型输出的通道顺序是固定的。我一般会在 Dataset 初始化时把categories[0][keypoints]读出来存成self.keypoint_names后面可视化时直接按这个顺序画省得对不上。2.3 写一个能直接用的 KeypointDataset下面这个 Dataset 类是我从这份资源里拆出来、又补了边界处理的版本可以直接抄import os import json import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms.functional as F class KeypointDataset(Dataset): def __init__(self, root, ann_file, transformsNone): self.root root self.transforms transforms with open(ann_file, r) as f: data json.load(f) # 按 image_id 把标注归组避免每张图遍历整个 annotations self.imgs {img[id]: img for img in data[images]} self.anns {} for ann in data[annotations]: self.anns.setdefault(ann[image_id], []).append(ann) self.ids list(self.imgs.keys()) # 关键点数量从 categories 里取不写死 self.num_keypoints len(data[categories][0][keypoints]) def __len__(self): return len(self.ids) def __getitem__(self, idx): img_id self.ids[idx] info self.imgs[img_id] img Image.open(os.path.join(self.root, info[file_name])).convert(RGB) anns self.anns.get(img_id, []) boxes, keypoints, labels [], [], [] for ann in anns: x, y, w, h ann[bbox] # 过滤掉宽高为 0 的脏标注否则 ROIAlign 会报错 if w 1 or h 1: continue boxes.append([x, y, x w, y h]) kps ann[keypoints] # 转成 [K, 3] 的形状 kps torch.tensor(kps, dtypetorch.float32).view(-1, 3) keypoints.append(kps) labels.append(ann[category_id]) if len(boxes) 0: # 没有有效标注时返回空 target训练循环里要跳过 target { boxes: torch.zeros((0, 4), dtypetorch.float32), keypoints: torch.zeros((0, self.num_keypoints, 3), dtypetorch.float32), labels: torch.zeros((0,), dtypetorch.int64), } else: target { boxes: torch.tensor(boxes, dtypetorch.float32), keypoints: torch.stack(keypoints), labels: torch.tensor(labels, dtypetorch.int64), } if self.transforms: img, target self.transforms(img, target) else: img F.to_tensor(img) return img, target逻辑上分三步先把 JSON 按image_id归组避免每次__getitem__都全量扫描再把每个 annotation 的 bbox 从[x, y, w, h]转成[x1, y1, x2, y2]因为模型内部用的是后者最后把 keypoints 从扁平数组 reshape 成[K, 3]。参数上要注意num_keypoints是从 categories 动态读的不是写死的这样换数据集时不用改代码。transforms这里留了口子因为关键点检测的增强必须同步变换图像和坐标不能只用torchvision的图像增强。2.4 同步增强为什么不能用普通的 RandomHorizontalFlip图像翻转了关键点坐标也得跟着翻否则模型学到的就是错位的映射。常见做法是自定义一个 transform接收(img, target)内部同时处理图像和target[keypoints]以及target[boxes]。水平翻转时新的 x 坐标是width - 1 - x同时如果关键点有左右对称的语义比如左眼和右眼还要交换对应的通道顺序。这一步是很多人翻车的地方——图像翻了、点没翻或者点翻了、左右没换训出来的模型在验证集上看着 loss 还行实际预测全是歪的。我一般会在增强后抽几张图可视化一遍确认点和图对得上再开始正式训练。3. 模型加载与训练循环从预训练权重到 loss 收敛3.1 为什么选 keypointrcnn_resnet50_fpn 而不是自己搭自己从 ResNet 加几个反卷积头搭关键点分支不是不行但 ROI 对齐、关键点热图或回归的损失设计、正负样本分配这些细节调起来非常耗时。keypointrcnn_resnet50_fpn是 torchvision 里已经封装好的两阶段模型第一阶段 RPN 出候选框第二阶段在 ROI 上同时做分类、框回归和关键点回归。它的关键点分支输出的是每个点的(x, y)偏移和可见性 logits损失函数内部已经处理了 visibility 的加权。用预训练权重的好处是 backbone 已经学过通用特征自建数据集哪怕只有几百张图也能在几十个 epoch 内看到合理结果。常见做法是加载 COCO 预训练权重然后把 box predictor 和 keypoint predictor 的最后一层换成自己关键点数量对应的输出。3.2 替换预测头与构建模型import torch import torchvision from torchvision.models.detection.keypoint_rcnn import KeypointRCNNPredictor from torchvision.models.detection.faster_rcnn import FastRCNNPredictor def build_model(num_keypoints, num_classes): # 加载 COCO 预训练权重backbone 和 RPN 直接复用 model torchvision.models.detection.keypointrcnn_resnet50_fpn( weightsCOCO_V1, weights_backboneIMAGENET1K_V1 ) # 替换框分类头COCO 是 91 类自建数据集通常只有几类 in_features model.roi_heads.box_predictor.cls_score.in_features model.roi_heads.box_predictor FastRCNNPredictor(in_features, num_classes) # 替换关键点头输出通道 num_keypoints * 2坐标 num_keypoints可见性 in_features_kp model.roi_heads.keypoint_predictor.kps_score_lowres.in_channels model.roi_heads.keypoint_predictor KeypointRCNNPredictor(in_features_kp, num_keypoints) return model这里有两个参数必须对上num_classes要包含背景类比如你只有一类目标就传 2num_keypoints必须和 Dataset 里读出来的一致。weightsCOCO_V1会下载预训练权重如果网络环境不允许可以提前把权重文件放到缓存目录或者用weightsNone从头训但那样需要的数据量和 epoch 数会成倍增加。替换预测头之后新加的层是随机初始化的训练时建议对 backbone 用较小的学习率对新加的 head 用较大的学习率常见做法是分组设置 param_groups。3.3 训练循环里必须处理的三个细节第一个细节是空 target。如果某张图过滤后没有有效标注target[boxes]是空张量直接送进模型会报错。训练循环里要判断if target[boxes].shape[0] 0: continue。第二个细节是模型在 train 模式下返回的是 loss 字典不是预测结果所以model(images, targets)的返回值要直接取losses model(images, targets)然后sum(losses.values()).backward()。第三个细节是学习率预热前几百个 iteration 用线性预热否则随机初始化的 head 会产生很大的梯度把预训练好的 backbone 带偏。def train_one_epoch(model, optimizer, data_loader, device, epoch): model.train() total_loss 0 for images, targets in data_loader: images [img.to(device) for img in images] targets [{k: v.to(device) for k, v in t.items()} for t in targets] # 跳过没有有效标注的样本 targets [t for t in targets if t[boxes].shape[0] 0] if len(targets) 0: continue loss_dict model(images, targets) losses sum(loss for loss in loss_dict.values()) optimizer.zero_grad() losses.backward() # 梯度裁剪防止早期梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm10.0) optimizer.step() total_loss losses.item() return total_loss / max(len(data_loader), 1)clip_grad_norm_的max_norm设成 10 是个经验值关键点回归的梯度有时候会很大不裁的话 loss 容易变成 nan。total_loss除以len(data_loader)只是粗略平均更严谨的做法是记录有效 batch 数。训练过程中要定期在验证集上跑model.eval()并关掉梯度看关键点预测的可视化结果光看 loss 曲线判断不了点有没有对齐。3.4 学习率分组与优化器配置params [p for p in model.parameters() if p.requires_grad] # 按参数名分组backbone 用小学习率head 用大学习率 backbone_params [p for n, p in model.named_parameters() if backbone in n and p.requires_grad] head_params [p for n, p in model.named_parameters() if backbone not in n and p.requires_grad] optimizer torch.optim.SGD([ {params: backbone_params, lr: 1e-4}, {params: head_params, lr: 1e-3}, ], momentum0.9, weight_decay1e-4) scheduler torch.optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1)backbone 的 1e-4 和 head 的 1e-3 差一个数量级是为了让新初始化的 head 快速跟上同时不破坏预训练特征。StepLR每 10 个 epoch 降一次学习率如果数据集很小几百张可以改成CosineAnnealingLR衰减更平滑。batch size 受显存限制关键点检测因为 ROI 上要多算一路分支显存占用比纯检测高常见做法是 batch size 设 2 到 4配合梯度累积模拟大 batch。4. 推理与可视化把预测点画回原图4.1 eval 模式下的输出结构模型切到eval()之后model(images)返回的是一个列表每个元素对应一张图里面有boxes、labels、scores、keypoints、keypoints_scores。keypoints的形状是[num_detections, K, 3]最后一维是(x, y, visibility_score)注意这里的第三个值不是 0/1/2 的可见性标签而是模型预测的可见性分数通常大于 0 就认为该点可见。scores是每个检测框的置信度可视化时一般只保留scores 0.7的框否则图上会糊一堆低置信度的点。4.2 可视化脚本与坐标还原import matplotlib.pyplot as plt import matplotlib.patches as patches import numpy as np def visualize_prediction(image, prediction, threshold0.7, keypoint_namesNone): # image 是 tensor转回 numpy 方便画 img_np image.permute(1, 2, 0).cpu().numpy() fig, ax plt.subplots(1, figsize(10, 10)) ax.imshow(img_np) boxes prediction[boxes].cpu().numpy() scores prediction[scores].cpu().numpy() keypoints prediction[keypoints].cpu().numpy() for box, score, kps in zip(boxes, scores, keypoints): if score threshold: continue x1, y1, x2, y2 box rect patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth2, edgecolorr, facecolornone) ax.add_patch(rect) for i, (kx, ky, kv) in enumerate(kps): if kv 0.5: continue ax.plot(kx, ky, o, markersize6, colorlime) if keypoint_names: ax.text(kx, ky, keypoint_names[i], fontsize8, coloryellow) plt.axis(off) plt.show()坐标不需要额外还原因为模型内部处理时已经考虑了图像原始尺寸输出的 keypoints 就是原图坐标系下的值。但如果你在 Dataset 里做了 resize 增强就必须在 target 里同步缩放关键点坐标否则训练和推理的尺度对不上。threshold设 0.7 是个起点实际项目里要根据验证集上的表现调宁可高一点漏检也别低到满屏假点。keypoint_names从 Dataset 里传过来画图时直接标名字排查哪个点偏了非常直观。4.3 用 OKS 而不是 mAP 来评估关键点检测的评估指标是 OKSObject Keypoint Similarity不是检测里的 IoU。OKS 会根据每个关键点的标准差对距离做归一化公式里有个k_i是每个点的尺度因子COCO 里对每个关键点都有预设值。自建数据集如果关键点语义和 COCO 不同k_i要自己估常见做法是用验证集上预测点和真值的距离统计来反推。torchvision 没有内置 OKS 计算需要自己写核心是exp(-d^2 / (2 * s^2 * k_i^2))其中s是目标框面积的平方根。评估时按 OKS 阈值 0.5 到 0.95 算 AP和 COCO 官方口径一致这样和别人报的结果才有可比性。5. 避坑与排查训关键点模型最容易翻车的五件事5.1 loss 一直不降最后变成 nan现象是训练前几个 iteration loss 正常突然跳到 nan 然后再也下不来。原因通常是关键点坐标没有归一化或者归一化方式不对导致回归目标的数值范围过大梯度爆炸。解决方法是检查target[keypoints]里 x、y 是否在图像尺寸范围内如果图像是 4000 像素宽坐标就是几千的量级虽然模型内部有处理但配合大学习率仍然容易炸。把学习率降到 1e-4 以下加上梯度裁剪通常能稳住。5.2 预测的点全部挤在框中心现象是可视化时所有关键点都落在检测框的正中间像没学一样。原因是关键点分支的损失权重太低或者keypoint_predictor替换后没有正确初始化。检查model.roi_heads.keypoint_predictor是否真的被替换成了KeypointRCNNPredictor以及num_keypoints是否和 Dataset 一致。另一个可能是 visibility 全填了 0导致所有点都不参与损失模型只能输出均值。把 visibility 按实际标注情况填 1 或 2重新训。5.3 验证集上框对了但点全歪现象是检测框位置很准但框内的关键点位置系统性偏移。原因多半是增强时图像和关键点没有同步变换比如只做了 resize 没缩放坐标或者做了 padding 没平移坐标。排查方法是关掉所有增强只用ToTensor如果点就准了说明问题出在增强环节。解决方法是把增强写成一个接收(img, target)的函数内部对图像和坐标做同样的几何变换翻转时记得交换对称点的通道顺序。5.4 训练时显存溢出现象是 batch size 设为 2 仍然 OOM。关键点检测比纯检测多一路 ROI 分支显存占用大概是 1.5 倍。解决方法是把model.roi_heads.detections_per_img从默认的 100 降到 20 到 50减少第二阶段处理的候选框数量同时把model.rpn.post_nms_top_n_train从 2000 降到 1000。如果还不够用梯度累积batch size 设 1累积 4 次再optimizer.step()等效 batch size 为 4。5.5 推理时点画出来了但位置对不上原图现象是可视化时点画在图像外面或者整体偏移。原因是推理时传入的图像经过了 resize但可视化时用的是原图。解决方法是推理前记录原图尺寸把预测的 keypoints 按比例缩放回去或者干脆推理时不 resize直接送原图让模型内部处理。如果用了torchvision.transforms.Resize记得同时保存缩放比例在可视化前把坐标除以比例。6. 把关键点数量改掉之后换数据集时最省事的验证套路换一个关键点数量不同的数据集时最容易被忽略的是KeypointRCNNPredictor的输出通道和 Dataset 里num_keypoints的联动。我一般会先写一个最小验证脚本不训练只做前向传播确认输入输出形状对得上。具体做法是构造一个 batch size 为 1 的假数据图像用torch.rand(3, 256, 256)target 里 boxes 给一个[10, 10, 100, 100]keypoints 给torch.rand(1, K, 3)然后跑一遍model.train()看 loss 能不能算出来再跑一遍model.eval()看输出 keypoints 的最后一维是不是 K。这一步花不了两分钟但能挡掉后面几个小时的无效训练。def sanity_check(model, num_keypoints, devicecpu): model.to(device) model.train() images [torch.rand(3, 256, 256).to(device)] targets [{ boxes: torch.tensor([[10., 10., 100., 100.]]).to(device), keypoints: torch.rand(1, num_keypoints, 3).to(device) * 100, labels: torch.tensor([1]).to(device), }] loss_dict model(images, targets) print(train loss keys:, loss_dict.keys()) model.eval() with torch.no_grad(): preds model(images) print(pred keypoints shape:, preds[0][keypoints].shape) assert preds[0][keypoints].shape[1] num_keypoints, 关键点数量不匹配这个脚本里keypoints乘了 100 是为了模拟真实坐标范围如果直接torch.rand得到 0 到 1 的值可能掩盖归一化相关的问题。assert那行是硬检查形状不对直接报错比训到一半才发现强。验证通过之后再跑一个 epoch 看 loss 是否在合理范围关键点回归的 loss 初始值一般在 0.1 到 1 之间如果上来就是几十说明坐标尺度或者学习率有问题。另一个省事的习惯是固定随机种子把torch.manual_seed、np.random.seed、random.seed都设上这样两次运行的 loss 曲线能对齐排查问题时变量少一个。数据加载的num_workers在调试阶段设 0避免多进程带来的随机性等流程跑通了再调大加速。从那以后我每次换数据集都先跑一遍这个 sanity check再开始正式训练省下来的时间够多调好几轮参。希望帮到你。本文还有配套的精品资源点击获取

PyTorch关键点检测实战：从自建数据集到模型训练与推理

相关推荐

PyTorch keypoint_rcnn 训练自建数据集：从标注到推理全流程

STM32F072C8T6最小核心板设计：从选型到PCB布局全攻略

定制Claude Code Mods：从黑盒到可控AI编程工具

更多资讯

大模型Skills是什么？从设计到落地的完整指南

基于STM32与HX711的智能体重秤：从原理图到低功耗实现

MATLAB异常检测工具箱实战：从算法选型到参数调优的完整指南

AI时代前端生态迁移：从Vue到React的双语能力构建

机器人嵌入式岗位重新定价：STM32贬值，系统能力吃香

Oracle 9i补丁9.2.0.8在Windows 32位上的安装与避坑全解析

今日更新

声子晶体仿真：三角晶格与六角晶格的能带计算实战

度假村安保定岗定编：岗位梳理、编制测算与排班设计

COMSOL等离子体模块实战：氩气双层DBD仿真建模全解析

本周精选

交通拥堵预测毕设：基于Python+LSTM的车辆流量预测系统全流程解析

MATLAB/Simulink搭建10机39节点电力系统暂态稳定仿真实战

无人机航拍三维重建全流程：从SfM到网格生成的避坑指南

本月热门

我发现了一个新思路：用 Remotion + Claude Code 像写代码一样自动化生成短视频

Windows下 Codex 中 Chrome 和 Computer Use 插件不可用问题排查及解决参考方式：TaoToken 统一 Key 配置与验证

2026 大模型集体涨价：用 Python 做企业 Token 成本测算与选型避坑（附配置）

读完文章，想聊聊您的网站？