关键在于数据、标注、配置三者对齐;推荐PyTorch+torchvision,用COCO格式组织数据,自定义Dataset同步变换图像与bbox,微调Faster R-CNN时替换box_predictor并注意类别ID从1开始,用原生训练循环控制loss与评估。

直接在自定义数据集上训练目标检测模型,关键不是换框架,而是把数据、标注、配置三者对齐。PyTorch(搭配torchvision或Detectron2)和TensorFlow(搭配TFOD API)都支持,但PyTorch生态目前更灵活、文档更贴近实战,推荐从 torchvision.models.detection 的预训练Faster R-CNN或RetinaNet入手。
数据准备:按规范组织图像和标注
检测任务需要图像 + 对应的边界框(bbox)+ 类别标签。不建议用XML或JSON手写——容易出错。推荐统一转成COCO格式(json)或Pascal VOC格式(每图一个XML),其中COCO更通用。
- 图像存放在 images/ 文件夹,命名简洁(如
001.jpg) - 标注用工具生成:LabelImg(VOC)、CVAT 或 Roboflow(导出COCO),确保每个bbox有
[x_min, y_min, x_max, y_max]和class_id - 划分 train/val 子集,生成对应
train.json和val.json(COCO)或ImageSets/Main/train.txt(VOC)
数据加载:写好Dataset类并做必要变换
torchvision 的 torch.utils.data.Dataset 必须返回 image(PIL或Tensor)和 target(字典,含 boxes、labels、可选 image_id 和 area)。注意三点:
-
boxes必须是 float32 Tensor,shape为(N, 4),且坐标不能越界(如 x_max > image_width) - 所有变换(Resize、ToTensor、Normalize)要同步作用于图像和 boxes —— 用
torchvision.transforms.v2(新版)或自定义函数,避免用老版transforms导致bbox错位 - 类别 ID 从 1 开始(0 留给背景),例如你的数据只有“car”和“person”,label 应为
[1, 2],不是[0, 1]
模型微调:加载预训练权重,替换分类头
以 Faster R-CNN 为例,不从头训练:
立即学习“Python免费学习笔记(深入)”;
- 用
torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")加载COCO预训练权重 - 修改
model.roi_heads.box_predictor:新类别数 = N(你的类别数)+ 1(背景),传入FastRCNNPredictor(1024, N+1) - 保持 backbone 和 neck 冻结前几层(可选),只训 head;训练稳定后再解冻微调
训练与验证:用原生PyTorch循环,别绕远路
不用封装库,自己写训练循环更可控:
- 损失函数自动计算:Faster R-CNN 的
model(images, targets)返回 dict(含loss_classifier、loss_box_reg等) - 优化器用
torch.optim.SGD(带 momentum=0.9)或 AdamW;学习率从 0.005 起步,配合torch.optim.lr_scheduler.StepLR - 验证时用
model.eval()+torch.no_grad(),用 COCOEval 或 simple metric(如 mAP@0.5)评估 - 每轮保存
model.state_dict(),别存整个 model 对象(体积大、兼容性差)
基本上就这些。难点不在代码量,而在数据格式校验和训练细节的耐心调试——比如 bbox 坐标是否归一化、类别ID是否对齐、GPU显存是否够用(batch_size=2 往往就够)。跑通第一个 epoch 有 loss 下降,后面就是调参和迭代的事了。










