YOLOv5从入门到部署:手把手教你完成自定义数据集训练与模型优化

YOLOv5从入门到部署:手把手教你完成自定义数据集训练与模型优化 1. 为什么选择YOLOv5做目标检测第一次接触YOLOv5是在2020年当时需要做一个工业质检项目。对比了当时主流的Faster R-CNN、SSD等方案后最终选择了YOLOv5。原因很简单它够快、够准、够简单。现在三年过去了YOLOv5依然是工业界落地最多的目标检测框架之一。YOLOv5最大的优势在于它的工程化程度。相比学术论文里的模型YOLOv5从数据准备到模型部署的全流程都做了深度优化。举个例子同样的RTX 3060显卡YOLOv5s模型训练COCO数据集只需要2天而其他框架可能需要3-4天。在部署阶段YOLOv5提供的export.py能一键转换成ONNX/TensorRT格式这对嵌入式设备特别友好。不过要注意的是YOLOv5有多个版本v6.0/v6.1/v7.0等不同版本之间有些API不兼容。我建议新手直接从最新的v7.0开始因为它的文档最全社区支持也最好。另外YOLOv5有四种预训练模型YOLOv5nNano2.6M参数YOLOv5sSmall7.2M参数YOLOv5mMedium21.2M参数YOLOv5lLarge46.5M参数对于大多数应用场景YOLOv5s已经足够用了。只有在需要检测特别小的物体时比如遥感图像中的车辆才需要考虑更大的模型。2. 环境搭建与数据准备2.1 快速搭建开发环境我习惯用conda管理Python环境这样可以避免包冲突。以下是具体步骤conda create -n yolov5 python3.8 conda activate yolov5 pip install torch1.12.1cu113 torchvision0.13.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113 pip install yolov5这里有个坑要注意PyTorch的CUDA版本必须和本地显卡驱动匹配。可以用nvidia-smi查看驱动支持的CUDA版本。如果装错了版本训练时会报CUDA out of memory的错误。2.2 数据标注实战技巧LabelImg是最常用的标注工具但我更推荐使用CVATComputer Vision Annotation Tool。它不仅支持本地使用还能团队协作。安装方法docker run -d -p 8080:8080 --name cvat openvino/cvat_server标注时要注意几个细节标注框要紧贴物体边缘但不要重叠对于遮挡物体按可见部分标注每个类别至少准备200张图片实际项目建议500保持长宽比避免极端尺寸如1000x10标注完成后数据目录应该这样组织dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/3. 模型训练与调优3.1 训练参数详解YOLOv5的train.py有30多个参数但新手只需要关注这几个python train.py \ --weights yolov5s.pt \ --data data/custom.yaml \ --epochs 100 \ --batch-size 16 \ --img 640 \ --device 0 \ --hyp data/hyps/hyp.scratch-low.yaml重点参数说明--img 640输入图像尺寸越大精度越高但速度越慢--batch-size根据显存调整RTX 3090可以设到32--hyp超参数文件控制学习率、数据增强等训练过程中可以用TensorBoard监控tensorboard --logdir runs/train3.2 数据增强策略YOLOv5默认会启用Mosaic增强这对小数据集特别有效。如果想自定义增强策略修改data/hyps/hyp.scratch-low.yamlhsv_h: 0.015 # 色调增强幅度 hsv_s: 0.7 # 饱和度增强幅度 hsv_v: 0.4 # 明度增强幅度 degrees: 10.0 # 旋转角度范围 translate: 0.1 # 平移范围 scale: 0.5 # 缩放范围有个实用技巧当发现模型对某些角度识别不好时可以增大degrees值如果颜色变化影响识别就降低hsv_h。4. 模型评估与部署4.1 关键指标解读训练完成后在runs/train/exp目录下会生成results.pngmAP0.5IoU阈值为0.5时的平均精度mAP0.5:0.95IoU阈值从0.5到0.95的平均精度precision查准率预测正确的比例recall查全率实际被检出的比例一般来说mAP0.5 0.9 说明模型优秀mAP0.5 0.5 需要重新调整数据或参数4.2 模型部署实战部署到生产环境推荐用TensorRT加速python export.py \ --weights runs/train/exp/weights/best.pt \ --include engine \ --device 0 \ --half部署时常见问题解决方案如果推理速度慢尝试启用--half使用FP16精度出现Unable to load model检查CUDA/cuDNN版本内存不足时减小--img-size对于视频流处理可以用这个模板代码import cv2 from yolov5 import YOLOv5 model YOLOv5(yolov5s.engine) cap cv2.VideoCapture(0) while True: ret, frame cap.read() results model.predict(frame) rendered results.render()[0] cv2.imshow(YOLOv5, rendered) if cv2.waitKey(1) ord(q): break5. 进阶优化技巧5.1 模型剪枝与量化当需要在边缘设备部署时可以对模型进行优化python -m yolov5.export \ --weights yolov5s.pt \ --include onnx \ --dynamic \ --simplify然后用ONNX Runtime量化python -m onnxruntime.tools.convert_onnx_models_to_ort yolov5s.onnx实测在Jetson Nano上量化后的模型推理速度能提升2-3倍。5.2 解决常见问题问题1过拟合症状训练集mAP很高但验证集很低 解决方法增加数据增强减小模型尺寸换yolov5n添加Dropout层问题2漏检症状某些类别recall很低 解决方法检查标注质量增加该类别样本数量调整anchor boxes问题3误检症状出现大量错误预测 解决方法提高confidence阈值增加负样本不包含目标的图像使用更复杂的模型yolov5m/l