MMDetection实战:从零构建自定义数据集到YOLOX模型部署全流程

张开发
2026/6/8 13:47:38 15 分钟阅读
MMDetection实战:从零构建自定义数据集到YOLOX模型部署全流程
1. 认识MMDetection与YOLOX第一次接触MMDetection时我被这个工具箱的完整性震撼到了。作为OpenMMLab旗下最受欢迎的项目之一它几乎囊括了目标检测领域所有主流算法。而YOLOX作为其中的佼佼者在速度和精度之间取得了很好的平衡。MMDetection的模块化设计让使用者可以像搭积木一样自由组合各个组件。从数据加载、模型构建到训练验证每个环节都有清晰的接口。我特别喜欢它的配置文件系统不需要修改代码就能实现各种实验配置的切换。比如你想对比不同学习率的效果只需要复制几份配置文件就能快速验证。YOLOX相比前代YOLO系列有几个显著改进采用了更高效的解耦头设计引入Anchor-free机制简化了实现使用SimOTA标签分配策略提升精度添加了强大的数据增强组合这些改进使得YOLOX在保持实时性的同时mAP指标提升了3-5个百分点。我在实际项目中使用YOLOX-s小模型版本在1080Ti显卡上能达到40FPS的推理速度完全满足工业检测的需求。2. 准备自定义数据集很多教程都假设你使用标准COCO数据集但现实中我们往往要处理各种非标准格式。最近一个项目中客户提供的是VOC格式的XML标注文件下面分享我的处理经验。首先确保你的数据集目录结构清晰。我习惯这样组织custom_dataset/ ├── raw_images/ # 原始图片 ├── annotations/ # XML标注文件 ├── ImageSets/ # 数据集划分 └── converted/ # 转换后的COCO格式数据集划分是关键步骤。我写了个Python脚本自动完成这个工作import os import random from sklearn.model_selection import train_test_split def split_dataset(image_dir, val_ratio0.2, test_ratio0.1): all_files [f for f in os.listdir(image_dir) if f.endswith(.jpg)] # 先分出测试集 train_val, test train_test_split(all_files, test_sizetest_ratio) # 再分训练集和验证集 train, val train_test_split(train_val, test_sizeval_ratio/(1-test_ratio)) return train, val, test对于XML转COCO格式我推荐使用xmltodict库简化解析过程。转换时要注意几个细节类别ID要从1开始0保留给背景bbox格式是[x,y,width,height]每个标注需要计算area字段转换完成后建议用pycocotools验证json文件是否正确from pycocotools.coco import COCO coco COCO(converted/annotations/instances_train.json) print(f数据集包含 {len(coco.dataset[categories])} 个类别)3. 配置YOLOX模型克隆MMDetection仓库后配置文件位于configs/yolox/目录。以yolox_s为例我们需要修改几个关键部分首先是数据路径配置data_root data/custom_dataset/ train_ann_file annotations/instances_train.json train_data_prefix dict(imgtrain/)然后是类别定义这里最容易出错metainfo dict( classes(product_a, product_b, defect), # 你的实际类别 palette[(220,20,60), (119,11,32), (0,0,142)] # 每个类别对应的显示颜色 )模型结构调整建议对于小数据集可以减小deepen_factor和widen_factor调整img_scale匹配你的输入尺寸修改num_classes为你的实际类别数数据增强策略也很重要train_pipeline [ dict(typeMosaic, img_scaleimg_scale, pad_val114.0), dict( typeRandomAffine, scaling_ratio_range(0.5, 1.5), border(-img_scale[0]//2, -img_scale[1]//2)), dict(typeMixUp, img_scaleimg_scale, ratio_range(0.8, 1.6)), dict(typeHSVAugment, hgain0.015, sgain0.7, vgain0.4) ]4. 模型训练与调优启动训练的命令很简单python tools/train.py configs/yolox/yolox_s_8xb8-300e_coco.py但实际训练中我总结了几点经验学习率设置初始学习率根据batch size调整使用warmup避免初期不稳定optim_wrapper dict( optimizerdict(typeSGD, lr0.01, momentum0.9, weight_decay0.0005), paramwise_cfgdict(norm_decay_mult0., bias_decay_mult0.))早停策略val_evaluator dict( typeCocoMetric, ann_filedata_root annotations/instances_val.json, metricbbox, format_onlyFalse)常见问题解决OOM错误减小batch size或输入尺寸NaN损失检查数据标注是否有无效框指标不提升尝试关闭部分数据增强训练过程可视化很重要MMDetection支持Tensorboard和WandBtensorboard --logdirwork_dirs5. 模型部署实战训练完成后我们得到.pth权重文件。部署时需要考虑几点模型转换from mmdet.apis import init_detector model init_detector(config_file, checkpoint_file, devicecuda:0) torch.save(model.state_dict(), deploy_model.pth)推理脚本示例def inference(image_path): img mmcv.imread(image_path) result inference_detector(model, img) # 可视化结果 vis_image model.show_result( img, result, score_thr0.3, showFalse) cv2.imwrite(result.jpg, vis_image)性能优化技巧使用TensorRT加速开启cudnn benchmark对静态输入启用jit编译对于嵌入式设备部署我推荐先转换为ONNX格式python tools/deployment/pytorch2onnx.py \ configs/yolox/yolox_s_8xb8-300e_coco.py \ work_dirs/yolox_s_8xb8-300e_coco/epoch_300.pth \ --output-file yolox.onnx \ --input-img demo/demo.jpg \ --shape 640 6406. 实际应用中的技巧在多个工业项目中我积累了一些实用经验数据方面对于小目标适当增大输入尺寸类别不平衡时使用Class Aware Sampler模糊标注使用soft-nms处理模型方面尝试不同的neck结构PAFPN效果不错调整loss权重平衡定位和分类使用EMA模型提升稳定性部署优化半精度推理速度提升30%使用多线程预处理对固定摄像头场景可以做背景建模一个检测pipeline的典型耗时分布数据加载: 15ms 预处理: 10ms 模型推理: 25ms 后处理: 5ms 可视化: 8ms7. 进阶扩展方向当你掌握基础流程后可以尝试这些进阶玩法自定义模型组件MODELS.register_module() class MyAttentionHead(BaseModule): def __init__(self, in_channels, out_channels): super().__init__() self.conv ConvModule(in_channels, out_channels, 3) def forward(self, x): return self.conv(x)知识蒸馏distiller dict( typeDetectionDistiller, teacher_trainableFalse, components[ dict( student_moduleneck.out_convs[0], teacher_moduleneck.out_convs[0], losses[dict(typeKLDivLoss, loss_weight1.0)]) ])模型量化python tools/deployment/pytorch2torchscript.py \ --quantize \ --config configs/yolox/yolox_s_8xb8-300e_coco.py \ --checkpoint work_dirs/yolox_s_8xb8-300e_coco/epoch_300.pth多任务学习model dict( typeMultiTaskDetector, backbone..., neck..., heads[ dict(typeYOLOXHead, num_classes10), # 任务1 dict(typeFCOSHead, num_classes5) # 任务2 ])

更多文章