CenterNet实战指南：如何高效训练自定义数据集并优化检测性能

1次阅读

没有评论

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

CenterNet 是一种基于关键点检测的目标检测算法，相比传统方法具有结构简单、无需 anchor 等优势。但在实际应用中，训练自定义数据集时往往会遇到以下典型问题：

小目标漏检：由于 heatmap 分辨率限制（通常为输入图像的 1 /4），小目标容易在降采样过程中丢失
目标尺度变化大：自定义数据集中物体尺寸分布与 COCO 等标准数据集差异明显
训练不稳定：offset 分支容易产生梯度爆炸，导致检测框位置不准

COCO 预训练模型与自定义数据集的主要参数差异体现在：

输入分辨率：COCO 通常使用 512×512，但自定义数据集可能需要根据目标大小调整
heatmap 标准差：控制关键点响应的扩散范围，小目标需要更小的 σ 值
损失权重：自定义数据集中 offset_loss 可能需要更高权重来提升定位精度

关键参数调整公式：

$$
\sigma = \frac{r}{6}, \quad r=\text{目标半径（像素）}
$$

import json
from PIL import Image
import os

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

    for img_id, img_file in enumerate(os.listdir(os.path.join(voc_dir, "JPEGImages"))):
        img_path = os.path.join(voc_dir, "JPEGImages", img_file)
        w, h = Image.open(img_path).size

        images.append({
            "id": img_id,
            "file_name": img_file,
            "width": w,
            "height": h
        })

        # 解析对应的 xml 标注文件（略）# 将 bbox 转换为 COCO 格式[x,y,width,height]

    with open(output_json, 'w') as f:
        json.dump({"images":images, "annotations":annotations, "categories":categories}, f)

对每个关键点生成高斯分布响应图：

$$
G(x,y) = \exp(-\frac{(x-\tilde{x})^2+(y-\tilde{y})^2}{2\sigma^2})
$$

Offset 计算确保亚像素级精度：

$$
\Delta x = \frac{x}{R} – \left\lfloor \frac{x}{R} \right\rfloor,\quad \Delta y = \frac{y}{R} – \left\lfloor \frac{y}{R} \right\rfloor
$$

其中 R = 4 是降采样倍数

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, targets in dataloader:
    optimizer.zero_grad()

    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()