企业官网建设流程全解析

简介：这份资源面向具备一定深度学习基础、希望动手实践关键点检测的开发者与研究人员，围绕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.json

instances_train.json里几个关键字段必须对齐：images列表里每项有id、file_name、width、height；annotations列表里每项有image_id、category_id、bbox（[x, y, w, h]）、keypoints、num_keypoints；categories里定义你的关键点类别，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, transforms=None): 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, dtype=torch.float32).view(-1, 3) keypoints.append(kps) labels.append(ann['category_id']) if len(boxes) == 0: # 没有有效标注时返回空 target，训练循环里要跳过 target = { 'boxes': torch.zeros((0, 4), dtype=torch.float32), 'keypoints': torch.zeros((0, self.num_keypoints, 3), dtype=torch.float32), 'labels': torch.zeros((0,), dtype=torch.int64), } else: target = { 'boxes': torch.tensor(boxes, dtype=torch.float32), 'keypoints': torch.stack(keypoints), 'labels': torch.tensor(labels, dtype=torch.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( weights='COCO_V1', weights_backbone='IMAGENET1K_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要包含背景类，比如你只有一类目标，就传 2；num_keypoints必须和 Dataset 里读出来的一致。weights='COCO_V1'会下载预训练权重，如果网络环境不允许，可以提前把权重文件放到缓存目录，或者用weights=None从头训，但那样需要的数据量和 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_norm=10.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}, ], momentum=0.9, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.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, threshold=0.7, keypoint_names=None): # 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, linewidth=2, edgecolor='r', facecolor='none') ax.add_patch(rect) for i, (kx, ky, kv) in enumerate(kps): if kv < 0.5: continue ax.plot(kx, ky, 'o', markersize=6, color='lime') if keypoint_names: ax.text(kx, ky, keypoint_names[i], fontsize=8, color='yellow') plt.axis('off') plt.show()

坐标不需要额外还原，因为模型内部处理时已经考虑了图像原始尺寸，输出的 keypoints 就是原图坐标系下的值。但如果你在 Dataset 里做了 resize 增强，就必须在 target 里同步缩放关键点坐标，否则训练和推理的尺度对不上。threshold设 0.7 是个起点，实际项目里要根据验证集上的表现调，宁可高一点漏检，也别低到满屏假点。keypoint_names从 Dataset 里传过来，画图时直接标名字，排查哪个点偏了非常直观。

4.3 用 OKS 而不是 mAP 来评估

关键点检测的评估指标是 OKS（Object 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, device='cpu'): 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，再开始正式训练，省下来的时间够多调好几轮参。希望帮到你。

本文还有配套的精品资源，点击获取