共计 4044 个字符，预计需要花费 11 分钟才能阅读完成。

CenterNet 训练自定义数据集全流程解析

背景与痛点分析

CenterNet 作为 anchor-free 目标检测算法的代表，其核心优势在于：

• 将目标检测转化为关键点预测问题，简化了后处理流程
• 单阶段检测器中的精度与速度平衡较好
• 模型结构相对简单，便于定制化修改

但在实际使用自定义数据集时，开发者常遇到以下问题：

1. 标注格式不兼容（需要转换为 COCO/VOC 格式）
2. 数据增强策略与模型特性不匹配
3. 小样本场景下模型容易过拟合
4. 训练过程中出现 NaN 等异常值

技术实现方案

数据准备阶段

格式转换（以 VOC 转 COCO 为例）

# 环境要求：Python 3.6+, pycocotools
import json
import os
from xml.etree import ElementTree as ET

def voc_to_coco(voc_dir, output_json):
    categories = [{"id": 1, "name": "your_class"}]  # 修改为实际类别
    images = []
    annotations = []
    ann_id = 1

    for img_id, filename in enumerate(os.listdir(os.path.join(voc_dir, 'JPEGImages'))):
        # 解析 XML 标注文件
        tree = ET.parse(os.path.join(voc_dir, 'Annotations', filename.replace('.jpg', '.xml')))
        root = tree.getroot()

        # 添加图像信息
        img_info = {
            "id": img_id + 1,
            "file_name": filename,
            "width": int(root.find('size/width').text),
            "height": int(root.find('size/height').text)
        }
        images.append(img_info)

        # 处理每个目标框
        for obj in root.iter('object'):
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)
            width = xmax - xmin
            height = ymax - ymin

            annotations.append({
                "id": ann_id,
                "image_id": img_id + 1,
                "category_id": 1,  # 对应 categories 中的 id
                "bbox": [xmin, ymin, width, height],
                "area": width * height,
                "iscrowd": 0
            })
            ann_id += 1

    # 生成最终 COCO 格式
    coco_format = {
        "categories": categories,
        "images": images,
        "annotations": annotations
    }

    with open(output_json, 'w') as f:
        json.dump(coco_format, f)

数据增强策略

推荐使用 Albumentations 库的增强组合：

import albumentations as A

train_transform = A.Compose([A.RandomResizedCrop(512, 512),  # 输入尺寸需与 config 一致
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.HueSaturationValue(p=0.2),
], bbox_params=A.BboxParams(format='coco'))

模型配置关键点

修改 centernet_resnet50_voc.yml 配置文件时需特别注意：

1. input_size: [512, 512] 应与数据增强输出尺寸一致
2. num_classes: 1 修改为实际类别数
3. head_conv: 64 小数据集可降低为 32 减少参数量
4. flip_test: False 验证阶段关闭可提速

训练调优技巧

1. 学习率策略：
2. 初始 lr=1e-4（小数据集可降至 3e-5）

3. 使用 ReduceLROnPlateau 调度器

   scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=3)

4. 早停机制：

   if val_map > best_map:
       best_map = val_map
       torch.save(model.state_dict(), 'best.pth')
       early_stop_counter = 0
   else:
       early_stop_counter += 1
       if early_stop_counter >= 5:  # 连续 5 次未提升
           break

代码实现示例

自定义数据加载器

# PyTorch 1.7+ required
from torch.utils.data import Dataset
import cv2

class CenterNetDataset(Dataset):
    def __init__(self, coco_json, transform=None):
        with open(coco_json) as f:
            self.data = json.load(f)
        self.transform = transform
        self.img_dir = "path/to/images"

    def __len__(self):
        return len(self.data['images'])

    def __getitem__(self, idx):
        img_info = self.data['images'][idx]
        img = cv2.imread(os.path.join(self.img_dir, img_info['file_name']))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 获取对应标注
        anns = [a for a in self.data['annotations'] 
               if a['image_id'] == img_info['id']]
        bboxes = [a['bbox'] for a in anns]

        if self.transform:
            transformed = self.transform(
                image=img,
                bboxes=bboxes,
                class_labels=[1]*len(bboxes))  # 单类别示例
            img = transformed['image']
            bboxes = transformed['bboxes']

        # 转换为 CenterNet 需要的 heatmap 格式
        # 此处省略 heatmap 生成代码
        return img, heatmap

训练过程可视化

import matplotlib.pyplot as plt

def plot_loss(loss_log):
    plt.figure(figsize=(10,5))
    plt.plot(loss_log['train'], label='Train')
    plt.plot(loss_log['val'], label='Validation')
    plt.title('Training Loss')
    plt.legend()
    plt.savefig('loss_curve.png')

常见问题解决方案

NaN 值排查流程

1. 检查标注框坐标是否越界（xmin<0 或 >width）
2. 验证数据增强后 bbox 是否有效

   def check_bbox(bbox, img_size):
       x,y,w,h = bbox
       assert x >=0 and y >=0 and w>0 and h>0, "负值坐标"
       assert x+w <= img_size[0], "宽度越界"
       assert y+h <= img_size[1], "高度越界"

3. 降低初始学习率（从 1e- 5 开始尝试）

显存优化策略

1. 梯度累积技术

   for i, (img, target) in enumerate(train_loader):
       pred = model(img)
       loss = criterion(pred, target)
       loss = loss / 4  # 假设累积 4 次
       loss.backward()
   
       if (i+1) % 4 == 0:
           optimizer.step()
           optimizer.zero_grad()

2. 使用混合精度训练

   from torch.cuda.amp import autocast, GradScaler
   
   scaler = GradScaler()
   with autocast():
       pred = model(img)
       loss = criterion(pred, target)
   scaler.scale(loss).backward()
   scaler.step(optimizer)
   scaler.update()

小数据集过拟合应对

1. 冻结 backbone 部分层

   for name, param in model.backbone.named_parameters():
       if 'layer4' not in name:  # 仅训练最后几层
           param.requires_grad = False

2. 添加 CutMix 数据增强
3. 使用 Label Smoothing

性能验证指标

典型训练过程 mAP 变化示例：

总结与扩展

通过本文方案，我们实现了：

1. 自定义数据的高效格式转换
2. 针对小数据集的模型配置优化
3. 稳定的训练过程监控

对于其他检测任务，可考虑以下扩展：

• 修改 head 结构适应关键点检测
• 结合 DCNv2 提升不规则目标检测效果
• 部署时转换为 TensorRT 加速

建议在实际项目中先使用小规模数据验证流程，再逐步扩展到全量数据。遇到问题时，可优先检查数据质量与增强策略，再调整模型参数。