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 python3.8 -y conda activate mmdet pip install torch1.9.0cu111 torchvision0.10.0cu111 -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, devicecuda: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: i1, 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(axis0) x_max, y_max points.max(axis0) 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在处理不规则物体时表现突出。修改配置文件时这三个地方必须检查mmdet/datasets/coco.py中的CLASSESmmdet/evaluation/functional/class_names.py中的coco_classes模型配置文件中的num_classes和数据路径以Faster RCNN为例关键配置修改如下# 修改模型类别数 model dict( roi_headdict( bbox_headdict(num_classes8))) # 修改数据集配置 data dict( traindict( ann_filedata/coco/annotations/instances_train.json, img_prefixdata/coco/images/train/), valdict( ann_filedata/coco/annotations/instances_val.json, img_prefixdata/coco/images/val/), testdict( ann_filedata/coco/annotations/instances_val.json, img_prefixdata/coco/images/val/))训练参数调优经验学习率batch_size2时lr0.002比较稳妥训练轮次小数据集建议50-100epoch大数据集可以适当减少数据增强RandomFlip和PhotoMetricDistortion对提升模型鲁棒性很有效4. 训练监控与问题排查训练开始时我习惯先用小批量数据测试流程是否通畅python tools/train.py configs/faster_rcnn.py --work-dir work_dirs --validate --gpus 1常见的训练问题及解决方案Loss不下降检查学习率是否过小数据标注是否正确显存不足减小batch_size开启梯度累积验证指标异常确认验证集标注是否完整训练过程中TensorBoard是最佳观察窗口tensorboard --logdir work_dirs --port 6006重点关注这些指标变化loss_rpn_cls建议值0.0-0.3loss_rpn_bbox建议值0.0-0.2mAP0.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_thr0.3)在部署到生产环境时建议将模型导出为ONNX格式torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})