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

1次阅读

没有评论

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

CenterNet 是一种基于关键点检测的目标检测算法，其核心思想是将目标检测问题转化为关键点预测问题。相比传统的 Faster R-CNN、YOLO 等算法，CenterNet 具有以下优势：

结构简单：无需设计 anchor boxes，减少了超参数调优的复杂度
检测精度高：通过预测目标中心点和尺寸，避免了 anchor 匹配带来的误差
推理速度快：单阶段检测架构，适合实时应用场景

在训练自定义数据集时，初学者常遇到以下问题：

标注格式转换困难：现有数据集可能是 VOC 或 COCO 格式，与 CenterNet 要求格式不匹配
小目标检测效果差：默认配置对小型物体敏感度不足
训练不稳定：容易出现 NaN 损失或梯度爆炸
模型收敛慢：需要合理设置学习率和损失权重
推理速度不理想：在边缘设备上部署时帧率较低

CenterNet 支持 COCO 和 PASCAL VOC 格式，推荐使用 COCO 格式。转换脚本示例：

import json
from pycocotools.coco import COCO

# VOC 转 COCO 格式示例
def voc_to_coco(voc_annotations, output_path):
    coco_data = {"images": [],
        "annotations": [],
        "categories": [{"id": 1, "name": "your_class"}]
    }

    # 具体转换逻辑
    for img_id, anns in enumerate(voc_annotations):
        coco_data["images"].append({
            "id": img_id,
            "file_name": anns["filename"],
            "width": anns["width"],
            "height": anns["height"]
        })

        for box in anns["bboxes"]:
            coco_data["annotations"].append({"id": len(coco_data["annotations"]),
                "image_id": img_id,
                "category_id": 1,
                "bbox": [box["x"], box["y"], box["w"], box["h"]],
                "area": box["w"] * box["h"],
                "iscrowd": 0
            })

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

关键参数说明（以 DLA-34 backbone 为例）：

# configs/your_dataset.py
model = dict(
    type='CenterNet',
    backbone=dict(
        type='DLA',
        depth=34,
        heads=dict(
            hm=80,  # 热图通道数（对应类别数）wh=2,   # 宽高预测
            reg=2   # 中心点偏移
        ),
        head_conv=64,  # 头部卷积通道数
        hm_weight=1.0, # 热图损失权重
        wh_weight=0.1  # 尺寸损失权重
    )
)

完整训练代码示例：

import torch
from torch.utils.data import DataLoader
from models.centernet import CenterNet
from datasets import YourDataset

# 1. 数据加载
train_dataset = YourDataset(
    img_dir='data/train',
    ann_file='data/annotations/train.json',
    input_size=512  # 输入尺寸
)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# 2. 模型初始化
model = CenterNet(backbone='dla34', num_classes=80)
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 3. 训练循环
for epoch in range(100):
    model.train()
    for batch in train_loader:
        images = batch['image'].cuda()
        targets = {k: v.cuda() for k, v in batch['target'].items()}

        optimizer.zero_grad()
        losses = model(images, targets)
        loss = sum(losses.values())
        loss.backward()
        optimizer.step()

        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

NaN 损失问题 ：
检查学习率是否过高（建议初始值 1e-4）
添加梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)
低 mAP 问题 ：
验证标注是否正确（使用可视化工具检查）
调整 hm_weight 和 wh_weight 比例
增加数据增强（随机翻转、颜色抖动）
显存不足 ：
减小 batch_size

使用混合精度训练：

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    outputs = model(inputs)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

小目标检测差 ：
增大输入分辨率（如从 512×512 调整为 1024×1024）
使用 FPN 结构增强特征融合
推理速度慢 ：
使用 TensorRT 加速
量化模型：torch.quantization.quantize_dynamic(model, {torch.nn.Conv2d}, dtype=torch.qint8)