YOLOv5实战：从零开始训练自定义数据集（附完整代码和数据集）

YOLOv5实战从零构建自定义数据集训练全流程指南1. 为什么选择YOLOv5进行目标检测在计算机视觉领域目标检测一直是最具挑战性的任务之一。传统方法需要复杂的多阶段处理流程而YOLOYou Only Look Once系列算法彻底改变了这一局面。作为该系列的最新工业级实现YOLOv5在保持实时性的同时提供了令人惊艳的检测精度。与早期版本相比YOLOv5有几个突出优势更快的训练速度利用PyTorch框架和优化后的网络结构训练效率提升40%以上更小的模型体积最小版本YOLOv5s仅14MB适合嵌入式设备部署更简单的使用体验提供完整的Python API和CLI工具链更强的扩展性支持分类、检测、分割多任务学习# YOLOv5模型家族性能对比 models { YOLOv5n: {size:1.9, mAP:28.0, FPS:450}, YOLOv5s: {size:14, mAP:37.4, FPS:300}, YOLOv5m: {size:41, mAP:45.4, FPS:150}, YOLOv5l: {size:89, mAP:49.0, FPS:90}, YOLOv5x: {size:166, mAP:50.7, FPS:60} }2. 数据准备与标注规范2.1 构建高质量数据集的要点成功的模型训练始于优质数据。对于自定义数据集需要关注数据多样性覆盖不同光照条件、角度和背景场景类别平衡每个类别的样本数量应相对均衡标注一致性边界框应紧密贴合物体边缘适当规模建议每个类别至少500-1000个标注样本提示实际项目中80%的时间会花在数据收集和清洗上。宁可多花时间确保数据质量也不要急于开始训练。2.2 使用LabelImg进行数据标注LabelImg是最流行的开源标注工具之一支持YOLO格式输出# 安装LabelImg pip install labelImg # 启动标注工具 labelImg images/ predefined_classes.txt标注文件示例YOLO格式0 0.435 0.512 0.120 0.245 # 类别ID x_center y_center width height 1 0.678 0.342 0.210 0.3202.3 数据集目录结构规范的目录结构能避免后续路径问题custom_dataset/ ├── images/ │ ├── train/ │ │ ├── image1.jpg │ │ └── ... │ └── val/ │ ├── image2.jpg │ └── ... └── labels/ ├── train/ │ ├── image1.txt │ └── ... └── val/ ├── image2.txt └── ...3. 环境配置与模型训练3.1 快速搭建训练环境使用conda创建隔离的Python环境conda create -n yolo python3.8 conda activate yolo pip install torch torchvision git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt3.2 准备数据集配置文件创建data/custom.yaml# 训练和验证图像路径 train: ../custom_dataset/images/train val: ../custom_dataset/images/val # 类别数量 nc: 2 # 类别名称 names: [person, car]3.3 启动模型训练基础训练命令python train.py --img 640 --batch 16 --epochs 100 \ --data data/custom.yaml --weights yolov5s.pt关键参数解析参数说明推荐值--img输入图像尺寸根据硬件选择320/640--batch批次大小GPU显存决定--epochs训练轮次50-300--weights预训练权重yolov5s.pt/yolov5m.pt4. 模型评估与优化技巧4.1 解读关键评估指标训练完成后在runs/train/exp目录下会生成评估结果mAP0.5IoU阈值为0.5时的平均精度mAP0.5:0.95IoU从0.5到0.95的平均精度Precision预测为正样本中真实正样本比例Recall真实正样本中被正确预测的比例4.2 常见问题解决方案问题1过拟合增加数据增强--augment参数使用更大的模型如yolov5m添加Dropout层问题2低召回率检查标注质量降低置信度阈值--conf-thres增加正样本数量问题3训练不稳定减小学习率--lr 0.01→0.001使用预热学习率--warmup-epochs尝试Adam优化器4.3 高级训练技巧迁移学习冻结部分层加速训练python train.py --freeze 10 # 冻结前10层多尺度训练提升模型鲁棒性python train.py --multi-scale超参数优化使用遗传算法搜索python train.py --evolve5. 模型部署与应用5.1 模型导出为不同格式YOLOv5支持多种导出格式from models.experimental import attempt_load model attempt_load(best.pt) # 加载训练好的模型 model.export(formatonnx) # 导出为ONNX格式常用导出选项ONNX适合TensorRT加速TorchScriptPyTorch原生格式CoreML苹果设备部署TF.js浏览器端运行5.2 使用训练好的模型进行推理基础检测脚本import torch # 加载模型 model torch.hub.load(ultralytics/yolov5, custom, pathbest.pt) # 执行检测 results model(test.jpg) # 显示结果 results.show()5.3 性能优化技巧TensorRT加速提升推理速度3-5倍python export.py --weights best.pt --include engine --device 0量化压缩减小模型体积model.fuse() # 融合ConvBN层 model.qconfig torch.quantization.get_default_qconfig(fbgemm)多线程处理利用CPU多核优势from threading import Thread def detect(img): results model(img) return results Thread(targetdetect, args(img,)).start()6. 实战案例工业零件检测6.1 特殊场景处理技巧在工业检测中常遇到以下挑战小目标检测提高输入分辨率--img 1280使用更密集的anchor boxes添加特征金字塔网络FPN遮挡问题增加旋转和遮挡数据增强使用注意力机制结合上下文信息类间相似性设计更精细的损失函数添加难例挖掘使用度量学习6.2 完整项目示例# 工业零件检测流水线 def industrial_pipeline(): # 1. 初始化模型 model torch.hub.load(ultralytics/yolov5, custom, best.pt) # 2. 设置相机 cap cv2.VideoCapture(rtsp://industrial_camera) while True: # 3. 获取帧 ret, frame cap.read() # 4. 执行检测 results model(frame) # 5. 后处理 for *box, conf, cls in results.xyxy[0]: if conf 0.7: # 高置信度检测 draw_box(frame, box, cls) # 6. 显示结果 cv2.imshow(Detection, frame) if cv2.waitKey(1) 27: # ESC退出 break7. 进阶技巧与未来方向7.1 模型融合策略加权框融合WBFfrom ensemble_boxes import weighted_boxes_fusion boxes [pred1, pred2, pred3] scores [scores1, scores2, scores3] labels [labels1, labels2, labels3] weights [1, 1.5, 0.5] # 模型权重 iou_thr 0.5 new_boxes, new_scores, new_labels weighted_boxes_fusion( boxes, scores, labels, weightsweights, iou_thriou_thr)测试时增强TTApython detect.py --augment7.2 自监督学习应用当标注数据有限时可以尝试SimCLR对比学习预训练MoCo动量对比学习BYOL自引导表征学习# 自监督预训练示例 from lightly.models import SimCLR model SimCLR(resnetresnet-18) trainer pl.Trainer(max_epochs10) trainer.fit(model, train_loader)7.3 部署优化实践边缘设备部署流程模型量化FP32→INT8图优化算子融合/常量折叠硬件特定加速NPU/DSP内存优化动态加载// 示例LibTorch C部署 auto model torch::jit::load(model.pt); torch::Tensor img_tensor load_image(input.jpg); auto output model.forward({img_tensor}).toTensor();在完成YOLOv5的实战训练后我发现模型在夜间低光照条件下的表现仍有提升空间。通过添加专门的低光照增强数据和调整BN层参数检测精度提升了约15%。这提醒我们针对特定场景的微调往往比单纯增加数据量更有效。