CenterNet实战：从零开始训练自定义数据集的完整指南

1次阅读

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

CenterNet 作为单阶段目标检测算法的代表，其核心优势在于简洁的 anchor-free 设计和高效的训练推理速度。相比传统检测模型，它通过预测目标中心点、尺寸和偏移量来检测物体，避免了复杂的 anchor 匹配过程。但在实际应用中，开发者常遇到以下问题：

自定义数据集标注格式与模型要求不匹配（如 COCO/VOC 转 CenterNet 格式）
数据增强策略选择不当导致模型欠拟合或过拟合
训练过程中出现 NaN 损失等异常情况
小数据集场景下模型泛化能力不足

CenterNet 默认支持 COCO 格式，但实际业务数据往往是 VOC 或其他格式。这里提供通用转换方案：

格式转换：将 VOC XML 转换为 COCO JSON 格式

import xml.etree.ElementTree as ET
import json
import os

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

    for xml_file in os.listdir(os.path.join(voc_dir, 'Annotations')):
        tree = ET.parse(os.path.join(voc_dir, 'Annotations', xml_file))
        root = tree.getroot()

        # 处理图片信息
        img_info = {'id': len(images) + 1,
            'file_name': root.find('filename').text,
            '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)

            annotations.append({
                'id': ann_id,
                'image_id': img_info['id'],
                'category_id': 1,  # 对应 categories 中的 id
                'bbox': [xmin, ymin, xmax - xmin, ymax - ymin],
                'area': (xmax - xmin) * (ymax - ymin),
                'iscrowd': 0
            })
            ann_id += 1

    coco_dict = {
        'images': images,
        'annotations': annotations,
        'categories': categories
    }

    with open(output_path, 'w') as f:
        json.dump(coco_dict, f)

数据增强：推荐使用 Albumentations 库实现定制化增强

import albumentations as A

# 典型增强组合
train_transform = A.Compose([A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(p=0.2),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, p=0.5),
    A.RandomResizedCrop(height=512, width=512, scale=(0.8, 1.2), ratio=(0.8, 1.2)),
], bbox_params=A.BboxParams(format='coco'))

关键参数解析（基于官方 DLA-34 backbone 配置）：

{'heads': {'hm': 80, 'wh': 2, 'reg': 2},  # hm 对应类别数
  'head_conv': 256,  # 头部卷积通道数，影响模型容量
  'input_res': 512,  # 输入分辨率，需与预处理保持一致
  'num_stacks': 2,   # 堆叠沙漏网络数量
  'dense_wh': False, # 是否预测密集宽高
  'modulate_hm': True  # 使用高斯热图调制
}

学习率策略：
初始学习率建议 3e-4（Adam 优化器）
使用余弦退火调度器：

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)

早停机制：当验证集 mAP 连续 3 个 epoch 不提升时终止训练

from torch.utils.data import Dataset
import cv2
import numpy as np

class CenterNetDataset(Dataset):
    def __init__(self, coco_json, img_dir, transform=None):
        with open(coco_json) as f:
            self.coco = json.load(f)
        self.img_dir = img_dir
        self.transform = transform

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

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

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

        # 应用增强
        if self.transform:
            transformed = self.transform(image=img, bboxes=bboxes)
            img = transformed['image']
            bboxes = transformed['bboxes']

        # 生成热图目标
        heatmap = np.zeros((1, img.shape[0]//4, img.shape[1]//4))  # 下采样 4 倍
        # ... 热图生成逻辑（详见官方实现）return {'image': img, 'heatmap': heatmap, 'reg': reg_target}

import matplotlib.pyplot as plt

def plot_loss(loss_log):
    plt.figure(figsize=(12, 4))
    plt.subplot(121)
    plt.plot(loss_log['train_loss'], label='Train')
    plt.plot(loss_log['val_loss'], label='Val')
    plt.legend()

    plt.subplot(122)
    plt.plot(loss_log['val_mAP'])
    plt.title('Validation mAP')
    plt.show()

NaN 损失排查：
检查标注框坐标是否越界（xmax > width 等）
验证数据增强是否产生无效框（面积 <=0）
添加梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), 10)
显存优化：
使用 torch.utils.data.DataLoader 的pin_memory和 num_workers 参数
梯度累积技巧（batch_size= 4 时等效于真实 batch_size=8）：

for i, data in enumerate(dataloader):
    loss = model(data)
    loss = loss / 2  # 假设累积 2 次
    loss.backward()

    if (i+1) % 2 == 0:
        optimizer.step()
        optimizer.zero_grad()

小数据集应对：
冻结 backbone 前几层：for param in model.backbone[:10].parameters(): param.requires_grad = False
使用 MixUp 增强：

def mixup_data(x, y, alpha=0.4):
    lam = np.random.beta(alpha, alpha)
    batch_size = x.size(0)
    index = torch.randperm(batch_size)
    mixed_x = lam * x + (1 - lam) * x[index]
    return mixed_x, y, y[index], lam

典型训练曲线（自定义行人检测数据集）：