PyTorch keypoint_rcnn 训练自建数据集：从标注到推理全流程-北京尧图网络

简介这份资源面向深度学习与计算机视觉方向的学习者和开发者聚焦如何用PyTorch中的Keypoint R-CNN在自建数据集上训练关键点检测模型可应用于姿态估计、人脸特征点定位等场景。压缩包共116个文件约8.55MB包含8个py脚本与2个ipynb笔记本用于模型搭建和训练流程演示34个txt与31个json承担标注与配置数据39个jpg提供图像样本另有md说明与gitmodules子模块配置整体结构便于按数据准备、模型配置、训练评估的脉络查阅。资源已有166人学习下载内容覆盖从数据组织、预处理增强到模型头部调整、损失组合、训练监控与ONNX导出的完整链路读者可据此复现关键点检测训练流程理解自建数据集标注格式与调参排错思路适合具备一定PyTorch基础、希望动手实践关键点检测项目的中高级开发者参考。1. 用 keypoint_rcnn 训练自建数据集从标注到推理的完整路径手里有一批自己拍的图片想标出每张图里目标的几个关键点比如零件的定位孔、手势的指尖、动物的关节然后训练一个能自动预测这些点坐标的模型——这件事用 PyTorch 里的keypoint_rcnn就能干而且不用从零搭网络。它本质上是把目标检测框架 Faster R-CNN 的头部换成了关键点回归分支先框出目标再在框内预测每个关键点的热力图最后取峰值坐标。相比直接回归坐标热力图方式对遮挡和小目标更稳这也是它比一堆自写回归网络更值得用的原因。适合谁手头有几十到几百张标注图、想快速跑通关键点检测闭环的开发者或者已经用过 YOLO 做检测、想扩展到关键点任务的熟手。整条链路是标注工具出 COCO 格式 → 自定义 Dataset 读图读标注 → 改keypoint_rcnn的头部通道数 → 训练调参 → 推理可视化。下面按这个顺序拆开讲参数和坑都会落到具体位置。2. 数据标注与 COCO 格式自建数据集的第一道坎2.1 为什么关键点任务优先选 COCO 标注格式关键点检测的标注比纯检测多一层结构每张图除了bbox还要有keypoints数组长度是3 * KK 是关键点个数每三个数是一组(x, y, v)v0表示未标注v1表示标注但不可见v2表示可见。keypoint_rcnn官方实现直接吃 COCO 格式省掉自己写解析器的功夫。常见做法是用 labelme 或 CVAT 标点再转 COCO。labelme 导出的是它自己的 JSON需要写脚本转。转换时最容易翻车的地方是v值labelme 不区分可见性转出来全是 1如果图里有遮挡点模型会学到错误信号。我的习惯是标注时对遮挡点直接不标转换脚本里把缺失点填(0,0,0)让损失函数忽略它们。COCO 的目录结构建议这样组织annotations放单个 JSON图片按训练验证分文件夹dataset/ annotations/ instances_train.json instances_val.json images/ train/ val/JSON 顶层是images、annotations、categories三个数组。categories里每个类要有keypoints字段和skeleton字段后者定义点之间的连线只影响可视化不影响训练但漏了某些版本会报错。2.2 从 labelme 到 COCO 的转换脚本下面这个脚本把 labelme 的多个 JSON 合成一个 COCO JSON关键点是处理v值和bbox计算import json, os, glob # 关键点名称顺序必须固定训练和推理都依赖这个顺序 KEYPOINTS [p1, p2, p3, p4] SKELETON [[0, 1], [1, 2], [2, 3]] def labelme_to_coco(json_dir, out_path, category_nametarget): images, annotations [], [] ann_id 1 for img_id, jf in enumerate(sorted(glob.glob(os.path.join(json_dir, *.json)))): data json.load(open(jf, encodingutf-8)) h, w data[imageHeight], data[imageWidth] images.append({id: img_id, file_name: os.path.basename(jf).replace(.json, .jpg), width: w, height: h}) # labelme 的 shapes 里可能混有矩形框和点这里只取点 pts {s[label]: s[points][0] for s in data[shapes] if s[shape_type] point} kps [] for name in KEYPOINTS: if name in pts: x, y pts[name] kps.extend([x, y, 2]) # 2 表示可见 else: kps.extend([0, 0, 0]) # 缺失点损失会忽略 xs [kps[i] for i in range(0, len(kps), 3) if kps[i2] 0] ys [kps[i1] for i in range(0, len(kps), 3) if kps[i2] 0] if not xs: continue # 整张图没有有效点跳过 x0, y0, x1, y1 min(xs), min(ys), max(xs), max(ys) # bbox 外扩 10 像素给关键点回归留上下文 pad 10 bbox [max(0, x0-pad), max(0, y0-pad), min(w, x1pad), min(h, y1pad)] bbox [bbox[0], bbox[1], bbox[2]-bbox[0], bbox[3]-bbox[1]] # COCO 是 xywh annotations.append({ id: ann_id, image_id: img_id, category_id: 1, bbox: bbox, area: bbox[2]*bbox[3], iscrowd: 0, num_keypoints: sum(1 for i in range(0, len(kps), 3) if kps[i2] 0), keypoints: kps }) ann_id 1 coco { images: images, annotations: annotations, categories: [{id: 1, name: category_name, keypoints: KEYPOINTS, skeleton: SKELETON}] } json.dump(coco, open(out_path, w, encodingutf-8), ensure_asciiFalse) print(fimages{len(images)}, annotations{len(annotations)}) labelme_to_coco(raw_json, dataset/annotations/instances_train.json)逻辑说明bbox用有效关键点的外接矩形再外扩是因为keypoint_rcnn的 ROI 裁剪依赖 bbox如果 bbox 贴得太紧边缘关键点的热力图会被裁掉一半训练时表现为边缘点永远学不准。pad10是经验值图分辨率大就调到 20。num_keypoints字段必须和实际v0的数量一致某些版本的 COCO 评估会校验它。参数说明KEYPOINTS顺序一旦定了就不能改改顺序等于让模型重新学category_id从 1 开始0 在 COCO 里通常留给背景。2.3 自定义 Dataset 与数据增强的边界keypoint_rcnn期望target里包含boxes、labels、keypoints三个张量。keypoints形状是[N, K, 3]第三维是(x, y, v)。写 Dataset 时最容易忽略的是增强随机裁剪、旋转会改变关键点坐标必须同步变换否则标签和图像对不上模型直接学废。我一般用torchvision的transforms只做颜色抖动和水平翻转翻转时手动镜像关键点 x 坐标并交换左右对称点的顺序。如果关键点没有左右对称关系翻转后 x 坐标用w - x即可。随机缩放可以做但要把 bbox 和关键点一起缩放且缩放后v0的点保持(0,0,0)。import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms.functional as F class KeypointDataset(Dataset): def __init__(self, coco_json, img_root, flipFalse): self.coco json.load(open(coco_json, encodingutf-8)) self.img_root img_root self.flip flip self.img_map {im[id]: im for im in self.coco[images]} # 按 image_id 聚合标注 self.ann_map {} for a in self.coco[annotations]: self.ann_map.setdefault(a[image_id], []).append(a) def __len__(self): return len(self.coco[images]) def __getitem__(self, idx): img_info self.coco[images][idx] img Image.open(os.path.join(self.img_root, img_info[file_name])).convert(RGB) w, h img.size anns self.ann_map.get(img_info[id], []) boxes, labels, kps [], [], [] for a in anns: x, y, bw, bh a[bbox] boxes.append([x, y, xbw, ybh]) labels.append(a[category_id]) k torch.tensor(a[keypoints], dtypetorch.float32).view(-1, 3) kps.append(k) target { boxes: torch.tensor(boxes, dtypetorch.float32), labels: torch.tensor(labels, dtypetorch.int64), keypoints: torch.stack(kps) if kps else torch.zeros((0, 4, 3)) } img F.to_tensor(img) if self.flip and torch.rand(1) 0.5: img torch.flip(img, dims[2]) target[boxes][:, [0, 2]] w - target[boxes][:, [2, 0]] target[keypoints][:, :, 0] w - target[keypoints][:, :, 0] # 注意有左右对称点时要在这里交换顺序 return img, target逻辑说明boxes用xyxy格式这是keypoint_rcnn内部期望的keypoints的v值原样保留损失函数会根据v0决定是否计入。翻转时boxes的 x 坐标交换要写成[:, [0,2]] w - [:, [2,0]]直接减会搞反左右边界。参数说明flip概率 0.5 是常规值如果关键点有强左右语义比如左手右手翻转后必须交换对应点的索引否则标签语义错乱这个坑很隐蔽训练 loss 会正常下降但验证时左右点互换。3. 改造 keypoint_rcnn 头部通道数与关键点数的对应3.1 模型结构里哪些层和关键点数强相关keypoint_rcnn的头部由两部分组成keypoint_head是一个反卷积加卷积的堆叠输出通道数等于关键点数 Kbox_head负责分类和 bbox 回归。用torchvision.models.detection.keypointrcnn_resnet50_fpn时通过num_keypoints参数控制 K但预训练权重是按 COCO 的 17 个点训的直接改 K 会导致头部权重形状不匹配。常见做法是加载预训练权重时用strictFalse让头部随机初始化骨干网络继续用预训练特征。这样小数据集也能收敛因为骨干已经会提特征了。如果数据集和 COCO 差异极大比如医学影像骨干也可以解冻微调但学习率要调小。import torchvision from torchvision.models.detection.keypointrcnn_resnet50_fpn import KeypointRCNN def build_model(num_keypoints, num_classes2, pretrainedTrue): model torchvision.models.detection.keypointrcnn_resnet50_fpn( pretrainedpretrained, # 骨干用 COCO 预训练 num_classesnum_classes, # 背景 1 类目标 num_keypointsnum_keypoints, # 你的 K min_size800, max_size1333 # 输入短边缩放范围 ) return model model build_model(num_keypoints4) # 检查头部输出通道 print(model.roi_heads.keypoint_predictor.kps_score_lowres.out_channels) # 应为 4逻辑说明num_classes包含背景所以单类目标是 2。min_size/max_size控制图像缩放小目标多就把min_size调大但显存会涨。kps_score_lowres是最终输出热力图的层它的out_channels必须等于 K打印出来核对一下改错了训练时会在 loss 计算处报形状错误。参数说明pretrainedTrue时头部会被替换strictFalse加载是内部行为如果自己写加载逻辑记得过滤掉keypoint_head和keypoint_predictor的键。3.2 损失函数构成与训练时的关注指标keypoint_rcnn的损失由四块组成loss_objectness、loss_rpn_box_reg、loss_classifier、loss_box_reg、loss_keypoint。前四块是检测部分的最后一块是关键点热力图交叉熵。训练时如果loss_keypoint不下降先看检测部分是否正常——如果 bbox 都框不准关键点无从谈起。我一般先冻结关键点头部只训检测部分几个 epoch等loss_classifier稳定后再解冻一起训。这样能避免早期关键点损失把骨干带偏。监控时重点看loss_keypoint和验证集上的 OKSObject Keypoint SimilarityOKS 是 COCO 的关键点评估指标比单纯看 loss 靠谱。import torch.optim as optim params [p for p in model.parameters() if p.requires_grad] optimizer optim.SGD(params, lr0.005, momentum0.9, weight_decay0.0005) lr_scheduler optim.lr_scheduler.StepLR(optimizer, step_size3, gamma0.1) model.train() for epoch in range(10): for images, targets in dataloader: images [img for img in images] targets [{k: v for k, v in t.items()} for t in targets] loss_dict model(images, targets) losses sum(loss for loss in loss_dict.values()) optimizer.zero_grad() losses.backward() optimizer.step() lr_scheduler.step() print(fepoch {epoch}, keypoint_loss{loss_dict[loss_keypoint].item():.4f})逻辑说明model在train()模式下返回 loss 字典eval()模式下返回预测结果别搞混。学习率 0.005 是检测任务的常用起点小数据集可以降到 0.001。StepLR每 3 个 epoch 降一次防止后期震荡。参数说明weight_decay0.0005 是标准值如果过拟合严重加到 0.001。batch size 受显存限制keypoint_rcnn输入分辨率高单卡 batch 2 到 4 比较现实太小可以用梯度累积。4. 训练避坑自建数据集上最容易翻车的五个点4.1 关键点全标成 v2 导致遮挡点学不准现象训练 loss 正常下降但推理时被遮挡的关键点位置乱飘OKS 明显低于可见点。原因标注时把所有点都标成v2模型被迫对遮挡点也输出确定坐标学到的是噪声。解决标注阶段对遮挡点要么不标v0要么标v1让损失函数对它们降权。COCO 的 OKS 评估本身会对v0的点跳过所以不影响指标。4.2 bbox 贴太紧导致边缘点热力图被裁现象位于目标边缘的关键点比如轮廓角点始终学不准其他点正常。原因ROI Align 按 bbox 裁剪特征bbox 如果刚好卡在关键点上该点的热力图目标落在裁剪边界外监督信号丢失。解决转换脚本里给 bbox 外扩 10 到 20 像素或者训练时对 bbox 做随机抖动。这个坑很隐蔽因为 loss 不会报错只是某些点精度上不去。4.3 图像增强没同步变换关键点现象训练 loss 震荡大验证集精度远低于训练集可视化发现预测点整体偏移。原因做了随机裁剪或旋转但关键点坐标没跟着变标签和图像错位。解决所有几何增强必须同步作用于关键点翻转时注意左右对称点的索引交换。颜色增强不影响坐标可以放心用。4.4 学习率过大导致关键点损失爆炸现象训练几个 step 后loss_keypoint变成 NaN。原因关键点热力图用的是交叉熵输出未归一化时数值范围大学习率过高容易梯度爆炸。解决把学习率降到 0.001 甚至 0.0005或者在关键点头部前加梯度裁剪。torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm10)是后悔药级别的保险。4.5 验证集评估时忘了切 eval 模式现象验证时 loss 正常但预测结果乱七八糟或者报错说返回类型不对。原因model.eval()没调用模型还在 train 模式返回的是 loss 字典而不是预测框。解决验证和推理前必须model.eval()并用torch.no_grad()包住否则显存会持续累积。5. 推理与可视化把预测点画回原图5.1 单张图推理的完整流程推理时模型返回每个目标的boxes、labels、scores、keypoints。keypoints形状是[N, K, 3]第三维是(x, y, v)这里的v是预测的置信度不是标注的可见性。取点时用v过滤低置信度点阈值一般设 0.5。import cv2 import numpy as np model.eval() img cv2.imread(test.jpg) rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) tensor torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0 with torch.no_grad(): pred model([tensor.to(device)])[0] # 取分数最高的目标 if len(pred[scores]) 0: best pred[scores].argmax() kps pred[keypoints][best].cpu().numpy() # [K, 3] for x, y, v in kps: if v 0.5: cv2.circle(img, (int(x), int(y)), 3, (0, 255, 0), -1) # 按 skeleton 连线 for i, j in SKELETON: if kps[i][2] 0.5 and kps[j][2] 0.5: cv2.line(img, (int(kps[i][0]), int(kps[i][1])), (int(kps[j][0]), int(kps[j][1])), (255, 0, 0), 2) cv2.imwrite(out.jpg, img)逻辑说明输入张量要归一化到[0,1]这是预训练模型的期望。pred[keypoints]里的坐标是原图尺度因为模型内部会做缩放还原不用手动映射。连线时两端点都要过置信度阈值否则会出现半截线。参数说明置信度阈值 0.5 是起点误检多就调高到 0.7漏检多就降到 0.3。cv2.circle的半径按图分辨率调大图用 5。5.2 用 OKS 量化评估而不是只看可视化可视化只能看个大概真正判断模型好坏要用 OKS。OKS 的计算需要每类关键点的标准差COCO 给了一套默认值自建数据集如果没有可以用关键点标注的方差近似。评估脚本可以基于pycocotools改把keypoints结果按 COCO 格式写回 JSON 再调COCOeval。一个容易忽略的点OKS 对 bbox 尺度敏感如果预测框比真实框大很多OKS 会被拉低。所以检测部分的精度直接影响关键点指标调参时不能只盯loss_keypoint。6. 小数据集上的进阶技巧从能跑到能用自建数据集通常只有几十到几百张图直接训keypoint_rcnn很容易过拟合。我踩过的坑是训练集 OKS 0.85验证集只有 0.4可视化一看验证图的点全挤在目标中心。后来靠三招把验证 OKS 拉到 0.7 以上。第一招是冻结骨干的前两层。keypoint_rcnn的 backbone 是 ResNet50FPN浅层学的是边缘纹理这些特征通用性强冻结它们能减少在小数据上的过拟合。代码上遍历model.backbone.body.layer1和layer2的参数设requires_gradFalse只训后面层和头部。代价是收敛慢一点但验证集更稳。第二招是关键点热力图的高斯核尺寸要和目标尺度匹配。默认实现里高斯核标准差是按 bbox 面积算的小目标上核太大热力图糊成一片峰值定位不准。如果目标普遍偏小可以在配置里把keypoint_head的kps_score_lowres上采样倍率调小或者自己改损失里的 sigma。这个改动需要动源码但效果明显。第三招是测试时增强TTA。推理时把图水平翻转一次两次预测的关键点取平均x 坐标翻转回来。这个技巧对对称目标特别有效OKS 通常能涨 2 到 3 个点。代价是推理时间翻倍看场景取舍。# 测试时水平翻转增强 def predict_tta(model, tensor, w): with torch.no_grad(): p1 model([tensor])[0] p2 model([torch.flip(tensor, dims[2])])[0] # p2 的 x 坐标翻转回原图 p2[keypoints][..., 0] w - p2[keypoints][..., 0] # 简单平均实际可按 score 加权 merged (p1[keypoints] p2[keypoints]) / 2 return merged逻辑说明翻转后预测的 x 坐标要映射回原图坐标系w - x是水平镜像公式。两次预测的目标顺序可能不一致严格做法是按 bbox IoU 匹配后再平均简单场景下取分数最高的目标平均也够用。参数说明TTA 只建议在推理阶段用训练阶段用会拖慢且收益小。如果关键点有强左右语义翻转后要交换对应点索引再平均否则左右点会被平均到中间。最后说个习惯每次改完数据或模型先拿 5 张训练图过一遍推理确认能复现标注点再跑全量训练。这个检查花两分钟能省掉几小时的无效训练。希望帮到你。本文还有配套的精品资源点击获取

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

相关推荐

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

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

Spring Boot表单如何发送PUT/DELETE请求？HiddenHttpMethodFilter原理与配置

更多资讯

大模型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 成本测算与选型避坑（附配置）

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