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

1. 从零开始搭建mmdetection训练环境

第一次接触mmdetection时，我被它强大的目标检测能力所吸引，但也被复杂的配置过程劝退过几次。经过多个项目的实战，我总结出了一套最稳定的环境搭建方法，特别适合新手快速上手。

mmdetection作为OpenMMLab家族的重要成员，对Python和PyTorch版本有严格要求。我强烈建议使用Python 3.8 + PyTorch 1.9的组合，这个组合在多个项目中验证过最稳定。安装时最容易踩的坑是CUDA版本不匹配，这里分享一个检查清单：

# 检查NVIDIA驱动版本 nvidia-smi # 查看CUDA版本 nvcc --version # 验证PyTorch能否调用GPU python -c "import torch; print(torch.cuda.is_available())"

安装mmdetection时，我习惯先创建干净的conda环境：

conda create -n mmdet python=3.8 -y conda activate mmdet pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html pip install mmcv-full -f https://download.mmcv.ai/mmcv/dist/cu111/torch1.9.0/index.html git clone https://github.com/open-mmlab/mmdetection.git cd mmdetection pip install -v -e .

实测发现，使用国内镜像源可以大幅加速安装过程。在pip安装命令后加上-i https://pypi.tuna.tsinghua.edu.cn/simple能节省大量时间。安装完成后，建议运行demo测试：

from mmdet.apis import init_detector, inference_detector config = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py' checkpoint = 'checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth' model = init_detector(config, checkpoint, device='cuda:0') inference_detector(model, 'demo/demo.jpg')

2. 自定义数据集准备与处理

在工业质检项目中，我处理过各种奇怪的数据集格式，最终发现COCO格式是兼容性最好的选择。制作自定义数据集时，建议使用labelme标注后通过脚本转换，这里分享我常用的转换模板：

import json import os from PIL import Image def labelme2coco(labelme_json, output_json): coco = {"images": [], "annotations": [], "categories": []} # 类别定义 for i, cls in enumerate(classes): coco["categories"].append({"id": i+1, "name": cls}) # 图片和标注处理 ann_id = 1 for img_id, json_file in enumerate(labelme_json): with open(json_file) as f: data = json.load(f) # 处理图片信息 img_path = os.path.join(img_dir, data["imagePath"]) width, height = Image.open(img_path).size coco["images"].append({ "id": img_id, "file_name": data["imagePath"], "width": width, "height": height }) # 处理标注信息 for shape in data["shapes"]: points = np.array(shape["points"]) x_min, y_min = points.min(axis=0) x_max, y_max = points.max(axis=0) coco["annotations"].append({ "id": ann_id, "image_id": img_id, "category_id": classes.index(shape["label"])+1, "bbox": [x_min, y_min, x_max-x_min, y_max-y_min], "area": (x_max-x_min)*(y_max-y_min), "iscrowd": 0 }) ann_id += 1 with open(output_json, 'w') as f: json.dump(coco, f)

数据集目录结构应该这样组织：

coco-data/ ├── annotations │ ├── instances_train.json │ └── instances_val.json └── images ├── train └── val

特别提醒：标注文件中的category_id必须从1开始，0在mmdetection中被保留为背景类。遇到过因为id从0开始导致训练报错的案例，调试了整整一天才发现问题。

3. 配置文件深度定制技巧

选择模型就像选工具，不同任务需要不同的"螺丝刀"。对于小目标检测，我推荐Cascade RCNN；对实时性要求高的场景，YOLOv3是不错的选择；而Deformable DETR在处理不规则物体时表现突出。

修改配置文件时，这三个地方必须检查：

1. mmdet/datasets/coco.py中的CLASSES
2. mmdet/evaluation/functional/class_names.py中的coco_classes
3. 模型配置文件中的num_classes和数据路径

以Faster RCNN为例，关键配置修改如下：

# 修改模型类别数 model = dict( roi_head=dict( bbox_head=dict(num_classes=8))) # 修改数据集配置 data = dict( train=dict( ann_file='data/coco/annotations/instances_train.json', img_prefix='data/coco/images/train/'), val=dict( ann_file='data/coco/annotations/instances_val.json', img_prefix='data/coco/images/val/'), test=dict( ann_file='data/coco/annotations/instances_val.json', img_prefix='data/coco/images/val/'))

训练参数调优经验：

• 学习率：batch_size=2时，lr=0.002比较稳妥
• 训练轮次：小数据集建议50-100epoch，大数据集可以适当减少
• 数据增强：RandomFlip和PhotoMetricDistortion对提升模型鲁棒性很有效

4. 训练监控与问题排查

训练开始时，我习惯先用小批量数据测试流程是否通畅：

python tools/train.py configs/faster_rcnn.py --work-dir work_dirs --validate --gpus 1

常见的训练问题及解决方案：

1. Loss不下降：检查学习率是否过小，数据标注是否正确
2. 显存不足：减小batch_size，开启梯度累积
3. 验证指标异常：确认验证集标注是否完整

训练过程中，TensorBoard是最佳观察窗口：

tensorboard --logdir work_dirs --port 6006

重点关注这些指标变化：

• loss_rpn_cls：建议值0.0-0.3
• loss_rpn_bbox：建议值0.0-0.2
• mAP@0.5：达到0.7以上说明模型表现良好

5. 预测结果可视化优化

预测结果可视化是展示成果的重要环节，我经常需要调整这些参数：

# 修改检测框颜色 def _draw_instances(self, image, instances, colors=[(255,255,255)]*80): # 修改置信度显示格式 labels = [f'{self.classes[class_id]} {score:.2f}' for class_id, score in zip(instances.labels, instances.scores)]

实际项目中，我总结出这些可视化技巧：

• 不同类别使用对比色，方便区分
• 置信度保留两位小数即可
• 大图适当调大字体，小图减小线宽
• 添加推理时间显示，方便性能评估

完整的预测脚本示例：

from mmdet.apis import init_detector, show_result_pyplot config = 'configs/faster_rcnn.py' checkpoint = 'work_dirs/latest.pth' model = init_detector(config, checkpoint) result = inference_detector(model, 'test.jpg') show_result_pyplot(model, 'test.jpg', result, score_thr=0.3)

在部署到生产环境时，建议将模型导出为ONNX格式：

torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})