手把手教你用YOLOX训练自定义数据集：从VOC格式转换到模型测试（附完整代码）

YOLOX实战指南从零构建自定义目标检测模型1. 环境配置与项目初始化在开始YOLOX项目前确保你的开发环境满足以下基础要求操作系统推荐Ubuntu 18.04/20.04或Windows 10/11WSL2环境下Python版本3.7或更高CUDA工具包11.0及以上需与GPU驱动兼容cuDNN7.6.5或更高安装核心依赖项的最佳实践# 创建隔离的Python环境 python -m venv yolox-env source yolox-env/bin/activate # Linux/Mac # yolox-env\Scripts\activate # Windows # 安装PyTorch与torchvision根据CUDA版本选择 pip install torch1.7.1cu110 torchvision0.8.2cu110 -f https://download.pytorch.org/whl/torch_stable.html提示若遇到CUDA版本不匹配问题可通过nvcc --version和nvidia-smi命令交叉验证驱动版本项目初始化步骤克隆官方仓库并安装依赖git clone https://github.com/Megvii-BaseDetection/YOLOX.git cd YOLOX pip install -r requirements.txt pip install -v -e . # 可编辑模式安装验证基础功能python tools/demo.py image -n yolox-s -c yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result常见环境问题解决方案问题类型典型表现解决方法CUDA不匹配RuntimeError: CUDA unknown error检查torch与CUDA版本对应关系显存不足CUDA out of memory减小batch size或使用更小模型依赖冲突ImportError无法导入模块重建虚拟环境并按顺序安装2. 数据集准备与VOC格式转换2.1 自定义数据集结构规范标准VOC数据集目录结构要求VOCdevkit/ └── VOC2007/ ├── Annotations/ # XML标注文件 ├── JPEGImages/ # 原始图像 └── ImageSets/ └── Main/ # 数据集划分文件 ├── train.txt ├── val.txt └── test.txt2.2 标注格式转换实战将常见标注格式转换为VOC的标准流程import xml.etree.ElementTree as ET from PIL import Image def json_to_voc(json_path, output_dir): 将LabelMe JSON标注转换为VOC XML格式 with open(json_path) as f: data json.load(f) img_path os.path.join(output_dir, JPEGImages, data[imagePath]) img Image.open(img_path) width, height img.size annotation ET.Element(annotation) ET.SubElement(annotation, filename).text data[imagePath] size ET.SubElement(annotation, size) ET.SubElement(size, width).text str(width) ET.SubElement(size, height).text str(height) ET.SubElement(size, depth).text 3 for shape in data[shapes]: obj ET.SubElement(annotation, object) ET.SubElement(obj, name).text shape[label] bndbox ET.SubElement(obj, bndbox) points np.array(shape[points]) ET.SubElement(bndbox, xmin).text str(int(points[:,0].min())) ET.SubElement(bndbox, ymin).text str(int(points[:,1].min())) ET.SubElement(bndbox, xmax).text str(int(points[:,0].max())) ET.SubElement(bndbox, ymax).text str(int(points[:,1].max())) tree ET.ElementTree(annotation) xml_path os.path.join(output_dir, Annotations, os.path.splitext(data[imagePath])[0] .xml) tree.write(xml_path)数据集划分的推荐比例训练集70%验证集15%测试集15%使用scikit-learn实现自动划分from sklearn.model_selection import train_test_split all_files [f.split(.)[0] for f in os.listdir(Annotations)] train_val, test train_test_split(all_files, test_size0.15, random_state42) train, val train_test_split(train_val, test_size0.1765, random_state42) # 0.1765≈15%/85% with open(ImageSets/Main/train.txt, w) as f: f.write(\n.join(train))3. 模型配置与训练调优3.1 关键配置文件修改需要调整的核心配置文件exps/example/yolox_voc/yolox_voc_s.py- 主实验配置yolox/data/datasets/voc_classes.py- 类别定义yolox/exp/yolox_base.py- 基础训练参数典型配置修改示例# yolox_voc_s.py self.num_classes 5 # 修改为实际类别数 self.data_dir data/VOCdevkit self.train_ann train.txt self.val_ann val.txt # voc_classes.py VOC_CLASSES ( person, car, dog, cat, bicycle )3.2 训练启动与参数优化单GPU训练命令python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 16 --fp16 -c yolox_s.pth关键训练参数说明参数作用推荐值-dGPU数量根据设备调整-b总batch size8-64视显存而定--fp16混合精度训练建议开启-c预训练权重官方提供或自定义--resume恢复训练配合checkpoint使用学习率调整策略修改exp文件self.basic_lr_per_img 0.01 / 64.0 # 基础学习率 self.max_epoch 300 # 总训练轮次 self.no_aug_epochs 20 # 最后N轮关闭数据增强 self.warmup_epochs 5 # 热身阶段4. 模型测试与性能优化4.1 单图推理与批量测试基础测试命令python tools/demo.py image -f exps/example/yolox_voc/yolox_voc_s.py \ -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth \ --path test_images/ --conf 0.3 --nms 0.5 \ --tsize 640 --save_result --device gpu关键参数调整建议置信度阈值(--conf)根据实际需求在0.1-0.5之间调整NMS阈值(--nms)通常0.45-0.65效果最佳测试尺寸(--tsize)保持与训练一致默认6404.2 模型性能评估指标使用COCO评估标准python tools/eval.py -f exps/example/yolox_voc/yolox_voc_s.py \ -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth \ -b 8 -d 1 --conf 0.001 --fp16 --test典型评估结果解读指标含义优秀值域AP平均精度(IOU0.5:0.95)50%AP50IOU0.5时的AP75%AP75IOU0.75时的AP60%FPS推理速度50(1080Ti)4.3 模型导出与部署ONNX导出命令python tools/export_onnx.py --output-name yolox_s.onnx \ -f exps/example/yolox_voc/yolox_voc_s.py \ -c YOLOX_outputs/yolox_voc_s/best_ckpt.pthTensorRT优化建议# trt_infer.py import tensorrt as trt logger trt.Logger(trt.Logger.INFO) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(yolox_s.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.max_workspace_size 1 30 # 1GB serialized_engine builder.build_serialized_network(network, config) with open(yolox_s.engine, wb) as f: f.write(serialized_engine)5. 实战技巧与问题排查5.1 训练过程常见问题问题1Loss震荡不收敛检查学习率是否过大验证数据标注质量尝试减小batch size问题2显存不足(OOM)# 修改exp文件中的以下参数 self.input_size (640, 640) # 可减小为(416,416) self.random_size (14, 26) # 调整多尺度范围 self.mosaic_prob 0.5 # 降低数据增强强度问题3验证指标异常检查验证集标注是否存在漏标确认eval_interval设置合理建议2-5个epoch验证数据增强是否过于激进5.2 性能优化技巧混合精度训练加速python tools/train.py ... --fp16 --occ多GPU数据并行python -m torch.distributed.launch --nproc_per_node4 tools/train.py ... -d 4缓存机制加速# 在exp文件中添加 self.enable_cache True self.cache_type ram # 或disk自定义数据增强self.mosaic_prob 1.0 self.mixup_prob 0.5 self.hsv_prob 1.0 self.flip_prob 0.5 self.degrees 10.0 # 旋转角度范围 self.translate 0.1 # 平移比例在实际项目中我们发现合理调整mosaic和mixup的概率能显著提升小目标检测性能。对于300epoch以上的训练建议在最后20-30epoch关闭mosaic增强以获得更稳定的收敛。