从单类到多类：YOLOv5关键点检测模型的扩展实践

1. 从单类到多类的关键点检测需求

在目标检测的实际应用中，我们经常会遇到需要同时识别多个类别并定位关键点的场景。比如在智能交通系统中，不仅要检测车辆，还需要识别车牌上的四个角点；在人脸识别场景中，可能需要同时检测不同年龄段的人脸及其五官特征点。这时候，传统的单分类关键点检测模型就显得力不从心了。

我最早接触这个问题是在开发一个停车场管理系统时。当时使用的YOLOv5车牌检测模型只能识别单一类型的车牌，但当场景中出现摩托车、电动车等不同车型时，系统就会漏检。更麻烦的是，不同车型的车牌位置和形状差异很大，简单的单类模型根本无法适应这种多样性。

多分类关键点检测的核心难点在于三个方面：首先是数据结构的兼容性，单类模型通常假设所有目标具有相同数量的关键点；其次是网络输出层的设计，需要同时容纳类别信息和关键点坐标；最后是损失函数的平衡，要避免分类损失和关键点回归损失相互干扰。这三个问题不解决，模型很容易出现关键点错位或者分类错误的情况。

2. 数据层面的改造实践

2.1 数据标注格式的调整

原始的单分类关键点检测通常使用固定格式的标注，比如车牌检测可能是："class, x, y, w, h, x1, y1, x2, y2, x3, y3, x4, y4"。但当扩展到多分类时，不同类别的关键点数量可能不同。比如人脸可能有68个关键点，而车牌只需要4个。

我在改造时采用了动态关键点数量的方案。具体做法是在标注文件中添加一个字段表示关键点数量，格式变为："class, num_keypoints, x, y, w, h, x1, y1, ..., xn, yn"。这样在数据加载时，就可以根据num_keypoints动态解析后续的关键点坐标。

python复制# 示例标注文件内容
0 4 0.5 0.5 0.3 0.2 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6  # 车牌，4个关键点
1 5 0.3 0.3 0.4 0.5 0.2 0.2 0.4 0.2 0.5 0.3 0.4 0.4 0.3 0.5  # 交通标志，5个关键点

2.2 数据加载器的修改

YOLOv5的数据加载主要在utils/plate_datasets.py中的LoadImagesAndLabels类。改造时需要特别注意两点：一是要正确解析动态关键点，二是要保持数据增强的兼容性。

对于马赛克增强这类操作，需要确保不同类别的关键点都能被正确处理。我的经验是，先统一将所有关键点转换为相对于图像尺寸的绝对坐标，进行增强后再转换回相对坐标。这样可以避免不同类别关键点数量不同带来的处理困难。

python复制def load_image_and_labels(self, index):
    # 原始单类数据加载逻辑
    label = self.labels[index]
    cls = label[0]
    num_kpts = int(label[1])  # 新增：获取关键点数量
    box = label[2:6]
    kpts = label[6:6+2*num_kpts]  # 动态获取关键点
    
    # 数据增强处理
    if self.augment:
        # 转换为绝对坐标进行增强
        img_h, img_w = img.shape[:2]
        box[0::2] *= img_w
        box[1::2] *= img_h
        kpts[0::2] *= img_w
        kpts[1::2] *= img_h
        
        # 执行马赛克增强等操作
        img, box, kpts = mosaic_augmentation(img, box, kpts)
        
        # 转换回相对坐标
        new_h, new_w = img.shape[:2]
        box[0::2] /= new_w
        box[1::2] /= new_h
        kpts[0::2] /= new_w
        kpts[1::2] /= new_h
    
    return img, cls, box, kpts

3. 网络结构的改造方案

3.1 输出层的重新设计

YOLOv5的检测头(Detect层)需要做较大改动。原始单分类模型的输出维度是nc+5+k2，其中nc是类别数，5表示bbox的4个坐标加1个置信度，k2表示k个关键点的坐标。在多分类场景下，这个设计有两个问题：一是假设所有类别有相同数量的关键点，二是关键点坐标与类别信息没有明确关联。

我的解决方案是在Detect层中，为每个类别分配独立的关键点输出通道。具体实现是在model/yolo_plate.py中修改Detect类的初始化：

python复制class Detect(nn.Module):
    def __init__(self, nc=80, anchors=(), ch=(), kpt_counts=[]):  # kpt_counts是每个类别的关键点数量列表
        super(Detect, self).__init__()
        self.nc = nc  # 类别数
        self.kpt_counts = kpt_counts  # 每个类别的关键点数量
        # 计算最大关键点数量，用于确定输出维度
        max_kpts = max(kpt_counts) if kpt_counts else 0
        self.no = nc + 5 + max_kpts * 2  # 每个anchor的输出维度
        # 其余初始化代码保持不变...

3.2 关键点预测的改进

在前向传播时，需要根据实际类别选择对应的关键点数量。这里有个技巧：虽然我们按照最大关键点数量分配了输出维度，但在计算损失时，只使用当前类别实际需要的部分。

python复制def forward(self, x):
    # 原始前向传播逻辑
    z = []
    for i in range(self.nl):
        x[i] = self.m[i](x[i])  # conv
        bs, _, ny, nx = x[i].shape
        x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
        
        if not self.training:  # 推理时处理
            # 对关键点做sigmoid并缩放
            x[i][..., 5+self.nc:] = x[i][..., 5+self.nc:].sigmoid() * 4 - 2
            # 根据类别选择关键点
            cls_pred = x[i][..., 5:5+self.nc].argmax(-1, keepdim=True)
            kpt_mask = torch.zeros_like(x[i][..., 5+self.nc:])
            for cls in range(self.nc):
                kpt_count = self.kpt_counts[cls]
                kpt_mask[cls_pred == cls, :2*kpt_count] = 1
            x[i][..., 5+self.nc:] *= kpt_mask
            
        z.append(x[i].view(bs, -1, self.no))
    return x if self.training else (torch.cat(z, 1), x)

4. 损失函数的适配优化

4.1 多任务损失平衡

多分类关键点检测的损失函数需要同时考虑四部分：分类损失、bbox回归损失、目标置信度损失和关键点回归损失。在utils/plate_loss.py中，我们需要修改__call__方法：

python复制def __call__(self, p, targets):
    device = targets.device
    lcls, lbox, lobj, lkpt = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
    
    # 构建目标
    tcls, tbox, indices, anchors, tkpts, kpt_masks = self.build_targets(p, targets)
    
    # 计算各项损失
    for i, pi in enumerate(p):
        b, a, gj, gi = indices[i]  # 图片索引、anchor索引、网格坐标
        tobj = torch.zeros_like(pi[..., 0], device=device)
        
        n = b.shape[0]  # 目标数量
        if n:
            ps = pi[b, a, gj, gi]  # 预测子集
            
            # Bbox回归
            pxy = ps[:, :2].sigmoid() * 2. - 0.5
            pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
            pbox = torch.cat((pxy, pwh), 1)
            iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)
            lbox += (1.0 - iou).mean()
            
            # 目标置信度
            tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype)
            
            # 分类损失
            if self.nc > 1:
                t = torch.full_like(ps[:, 5:5+self.nc], self.cn, device=device)
                t[range(n), tcls[i]] = self.cp
                lcls += self.BCEcls(ps[:, 5:5+self.nc], t)
            
            # 关键点损失
            cls_idx = tcls[i]  # 当前目标的类别
            kpt_count = self.kpt_counts[cls_idx]  # 该类别的关键点数量
            pkpts = ps[:, 5+self.nc:5+self.nc+2*kpt_count].sigmoid() * 4 - 2
            lkpt += self.kpt_loss(pkpts, tkpts[i][:, :2*kpt_count], kpt_masks[i][:, :kpt_count])
        
        # 置信度损失
        lobj += self.BCEobj(pi[..., 4], tobj) * self.balance[i]
    
    # 损失加权
    lbox *= self.hyp['box']
    lobj *= self.hyp['obj']
    lcls *= self.hyp['cls']
    lkpt *= self.hyp['kpt']
    
    bs = tobj.shape[0]
    loss = lbox + lobj + lcls + lkpt
    return loss * bs, torch.cat((lbox, lobj, lcls, lkpt, loss)).detach()

4.2 关键点损失的特殊处理

不同类别的关键点可能有不同的重要性。比如车牌的关键点对定位精度要求很高，而人脸的一些边缘特征点则可以容忍较大误差。我们可以通过调整关键点损失权重来实现这一点：

python复制def kpt_loss(self, pred, target, mask):
    # pred: [n, 2*k] 预测的关键点
    # target: [n, 2*k] 真实关键点
    # mask: [n, k] 关键点可见性掩码
    
    k = pred.shape[1] // 2
    loss = 0
    for i in range(k):
        p = pred[:, 2*i:2*i+2]
        t = target[:, 2*i:2*i+2]
        m = mask[:, i]
        
        # 根据关键点索引调整权重
        if i in [0, 2]:  # 重要关键点
            weight = 1.2
        else:
            weight = 0.8
            
        loss += weight * (m * ((p - t).abs()).sum(-1)).mean()
    
    return loss / max(1, k)

5. 后处理与推理优化

5.1 多分类NMS的改造

原始YOLOv5的NMS处理只考虑了bbox和类别，我们需要扩展它以支持关键点。在utils/general.py中修改non_max_suppression函数：

python复制def non_max_suppression_kpt(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, kpt_counts=[]):
    """支持多分类关键点的NMS实现"""
    nc = prediction.shape[2] - 5 - max(kpt_counts)*2  # 类别数
    xc = prediction[..., 4] > conf_thres  # 候选框
    
    # 设置
    max_wh = 4096  # 最大宽高
    max_det = 300  # 每张图最大检测数
    output = [torch.zeros((0, 6 + max(kpt_counts)*2), device=prediction.device)] * prediction.shape[0]
    
    for xi, x in enumerate(prediction):
        x = x[xc[xi]]  # 筛选置信度
        
        if not x.shape[0]:
            continue
            
        # 计算置信度
        x[:, 5:5+nc] *= x[:, 4:5]  # conf = obj_conf * cls_conf
        
        # 转换bbox格式
        box = xywh2xyxy(x[:, :4])
        
        # 处理多标签情况
        if nc > 1:
            i, j = (x[:, 5:5+nc] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j+5, None], x[i, 5+nc:5+nc+2*max(kpt_counts)], j[:, None].float()), 1)
        else:
            # 单类别处理
            conf, j = x[:, 5:5+nc].max(1, keepdim=True)
            x = torch.cat((box, conf, x[:, 5+nc:5+nc+2*max(kpt_counts)], j.float()), 1)[conf.view(-1) > conf_thres]
        
        # 按类别筛选
        if classes is not None:
            x = x[(x[:, -1:].expand(-1, len(classes)) == torch.tensor(classes, device=x.device)).any(1)]
        
        # NMS处理
        c = x[:, -1:] * max_wh  # 类别偏移
        boxes, scores = x[:, :4] + c, x[:, 4]  # 添加类别偏移的boxes
        i = torchvision.ops.nms(boxes, scores, iou_thres)
        if i.shape[0] > max_det:
            i = i[:max_det]
        
        # 后处理：根据实际类别筛选关键点
        output_xi = x[i]
        for det in output_xi:
            cls = int(det[-1])
            kpt_count = kpt_counts[cls]
            det[6+2*kpt_count:-1] = 0  # 将多余关键点置零
        
        output[xi] = output_xi
    
    return output

5.2 推理性能优化技巧

在实际部署中，我发现几个提升多分类关键点检测效率的技巧：

1. 动态输出裁剪：根据置信度阈值提前裁剪低置信度的预测，减少NMS计算量
2. 关键点量化：将关键点坐标从浮点数转换为定点数，可以减少内存占用和加速计算
3. 类别分组处理：对关键点数量相同的类别进行分组处理，可以减少条件判断的开销

python复制def optimize_inference(model, img_size=640):
    """模型推理优化"""
    # 设置模型为评估模式
    model.eval()
    
    # 创建示例输入
    img = torch.zeros((1, 3, img_size, img_size))
    
    # 前向传播跟踪
    traced_model = torch.jit.trace(model, img)
    
    # 应用优化
    torch.jit.optimized_execution(traced_model.graph)
    
    # 量化关键点输出
    for m in traced_model.modules():
        if isinstance(m, Detect):
            # 将关键点输出量化为int8
            m.no = m.nc + 5 + max(m.kpt_counts)*2
            m.forward = torch.jit.script(m.forward)
    
    return traced_model

6. 实际应用中的挑战与解决方案

6.1 类别不平衡问题

在多分类场景中，某些类别的样本可能远多于其他类别。比如在交通场景中，汽车的数量可能远多于特种车辆。这会导致模型偏向于数量多的类别。

我采用的解决方案是：

1. 样本重加权：在数据加载时为不同类别设置不同的采样权重
2. 损失函数调整：为少数类别分配更高的分类损失权重
3. 困难样本挖掘：在训练过程中，特别关注那些关键点定位困难的样本

python复制class BalancedDataset(torch.utils.data.Dataset):
    def __init__(self, dataset, class_weights):
        self.dataset = dataset
        self.class_weights = class_weights
        # 为每个类别创建样本索引
        self.class_indices = [[] for _ in range(len(class_weights))]
        for idx, (_, label) in enumerate(dataset):
            cls = int(label[0])
            self.class_indices[cls].append(idx)
    
    def __getitem__(self, index):
        # 根据类别权重随机选择类别
        cls = random.choices(range(len(self.class_weights)), weights=self.class_weights)[0]
        # 从该类别中随机选择样本
        idx = random.choice(self.class_indices[cls])
        return self.dataset[idx]
    
    def __len__(self):
        return len(self.dataset)

6.2 关键点数量不一致的处理

不同类别的关键点数量不同会带来很多实现上的麻烦。我的经验是采用"最大公共维度"策略：

1. 按照最大关键点数量设计网络输出维度
2. 为每个类别维护一个关键点数量记录
3. 在训练和推理时，根据实际类别动态选择有效关键点
4. 对无效关键点位置，在损失计算时给予零权重

这种方法虽然会浪费一些计算资源，但实现起来最简单，且不会影响模型精度。

python复制def collate_fn(batch):
    """处理不同关键点数量的批数据"""
    images = []
    targets = []
    max_kpts = 0
    
    # 首先找出批数据中的最大关键点数量
    for img, target in batch:
        cls = int(target[0])
        num_kpts = int(target[1])
        if num_kpts > max_kpts:
            max_kpts = num_kpts
    
    # 统一填充关键点
    for img, target in batch:
        cls = int(target[0])
        num_kpts = int(target[1])
        box = target[2:6]
        kpts = target[6:6+2*num_kpts]
        
        # 填充无效关键点
        padded_kpts = np.zeros(2*max_kpts)
        padded_kpts[:2*num_kpts] = kpts
        
        # 创建新的目标向量
        new_target = np.concatenate([
            [cls, num_kpts],
            box,
            padded_kpts
        ])
        
        images.append(img)
        targets.append(torch.tensor(new_target))
    
    return torch.stack(images), torch.stack(targets)

7. 模型训练与调优经验

7.1 训练策略调整

多分类关键点检测模型的训练需要特别注意学习率调度和早停策略。我总结的有效做法包括：

1. 渐进式学习率：初始阶段用较小学习率(1e-4)稳定训练，后期逐步增大(3e-4)
2. 关键点损失预热：前几个epoch主要优化分类和bbox损失，逐渐引入关键点损失
3. 类别平衡采样：确保每个batch中包含所有类别的样本

python复制def train(model, train_loader, optimizer, epoch, warmup_epochs=3):
    model.train()
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()
        
        # 前向传播
        output = model(data)
        
        # 计算损失
        loss, loss_components = criterion(output, target)
        
        # 关键点损失预热
        if epoch < warmup_epochs:
            kpt_loss = loss_components[3]
            loss = loss - 0.5 * kpt_loss  # 减少关键点损失权重
        
        # 反向传播
        loss.backward()
        optimizer.step()
        
        # 学习率调整
        if epoch < 5:
            lr = 1e-4 + (3e-4 - 1e-4) * (epoch / 5)
            for param_group in optimizer.param_groups:
                param_group['lr'] = lr

7.2 模型评估指标

除了常规的mAP指标外，对于关键点检测还需要特别关注：

1. 关键点准确率(KPA)：关键点与标注点的平均距离
2. 类别敏感关键点误差：按类别分别统计关键点误差
3. 联合检测准确率：同时正确分类和定位关键点的比例

python复制def evaluate(model, val_loader, kpt_counts):
    model.eval()
    total_kpt_error = 0
    cls_kpt_errors = [0] * len(kpt_counts)
    cls_counts = [0] * len(kpt_counts)
    joint_correct = 0
    
    with torch.no_grad():
        for data, target in val_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            
            # 应用NMS
            detections = non_max_suppression_kpt(output, kpt_counts=kpt_counts)
            
            for det, gt in zip(detections, target):
                if len(det) == 0:
                    continue
                
                # 获取预测结果
                pred_cls = int(det[0, -1])
                pred_kpts = det[0, 6:6+2*kpt_counts[pred_cls]].cpu().numpy()
                
                # 获取真实标注
                gt_cls = int(gt[0])
                gt_kpts = gt[6:6+2*kpt_counts[gt_cls]].cpu().numpy()
                
                # 计算关键点误差
                error = np.mean(np.abs(pred_kpts - gt_kpts))
                total_kpt_error += error
                cls_kpt_errors[gt_cls] += error
                cls_counts[gt_cls] += 1
                
                # 联合准确率
                if pred_cls == gt_cls and error < 0.05:  # 误差小于5%认为正确
                    joint_correct += 1
    
    # 计算指标
    avg_kpt_error = total_kpt_error / sum(cls_counts)
    cls_avg_errors = [e/max(1,c) for e,c in zip(cls_kpt_errors, cls_counts)]
    joint_acc = joint_correct / max(1, sum(cls_counts))
    
    return {
        'avg_kpt_error': avg_kpt_error,
        'cls_kpt_errors': cls_avg_errors,
        'joint_acc': joint_acc
    }

8. 部署与性能优化

8.1 模型量化与加速

在实际部署中，模型大小和推理速度至关重要。我通常采用以下优化手段：

1. FP16量化：将模型权重转换为半精度浮点数
2. TensorRT优化：使用NVIDIA的TensorRT引擎加速推理
3. 关键点后处理优化：将关键点解码过程转移到GPU上执行

python复制def export_onnx(model, output_path, img_size=640):
    """导出为ONNX格式并优化"""
    model.eval()
    x = torch.randn(1, 3, img_size, img_size, requires_grad=True)
    
    # 导出原始模型
    torch.onnx.export(
        model,
        x,
        output_path,
        opset_version=12,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch'},
            'output': {0: 'batch'}
        }
    )
    
    # 使用ONNX Runtime优化
    sess_options = onnxruntime.SessionOptions()
    sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess_options.optimized_model_filepath = output_path.replace('.onnx', '_opt.onnx')
    
    # 量化到FP16
    onnx_model = onnx.load(output_path)
    onnx_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True)
    onnx.save(onnx_model, output_path.replace('.onnx', '_fp16.onnx'))

8.2 边缘设备部署技巧

在边缘设备(如Jetson系列)上部署时，还需要考虑：

1. 内存限制：可能需要降低输入分辨率或减少模型层数
2. 功耗优化：动态调整推理频率，空闲时降低功耗
3. 流水线处理：将关键点检测分成多个阶段，平衡延迟和吞吐量

python复制class EdgeInferencePipeline:
    def __init__(self, model_path, device='cuda', img_size=640):
        # 加载模型
        self.model = onnxruntime.InferenceSession(model_path)
        self.img_size = img_size
        self.device = device
        
        # 预热
        dummy_input = np.random.rand(1, 3, img_size, img_size).astype(np.float32)
        self.model.run(None, {'input': dummy_input})
    
    def preprocess(self, img):
        """图像预处理"""
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = img.transpose(2, 0, 1)  # HWC to CHW
        img = np.ascontiguousarray(img)
        img = img.astype(np.float32) / 255.0
        return np.expand_dims(img, 0)
    
    def postprocess(self, output, conf_thresh=0.3):
        """后处理"""
        # 应用NMS
        detections = non_max_suppression_kpt(
            torch.tensor(output[0]),
            conf_thres=conf_thresh
        )
        
        results = []
        for det in detections[0]:
            if det is None:
                continue
            cls = int(det[-1])
            kpts = det[6:6+2*self.kpt_counts[cls]].cpu().numpy()
            results.append({
                'class': cls,
                'bbox': det[:4].cpu().numpy(),
                'confidence': det[4].item(),
                'keypoints': kpts.reshape(-1, 2)
            })
        
        return results
    
    def process_frame(self, frame):
        """处理单帧"""
        # 预处理
        img = self.preprocess(frame)
        
        # 推理
        start = time.time()
        output = self.model.run(None, {'input': img})
        infer_time = time.time() - start
        
        # 后处理
        results = self.postprocess(output)
        
        return results, infer_time