【YOLOv5 v6.1】从零到一：手把手实战自定义数据集训练与部署避坑指南

1. 环境准备从零搭建YOLOv5训练环境第一次接触YOLOv5时我最头疼的就是环境配置。这里分享一个经过多次验证的稳定方案适用于大多数NVIDIA显卡设备。首先需要安装Anaconda这是管理Python环境的利器。我习惯用Miniconda体积更小但功能完整wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh创建专属环境时特别注意Python版本要与PyTorch匹配。实测v6.1版本在Python 3.8上最稳定conda create -n yolov5 python3.8 conda activate yolov5安装PyTorch时最容易踩坑。建议先去PyTorch官网用他们的配置生成器根据你的CUDA版本生成安装命令。比如我的RTX 3080是CUDA 11.3pip install torch1.10.1cu113 torchvision0.11.2cu113 torchaudio0.10.1 -f https://download.pytorch.org/whl/cu113/torch_stable.html最后克隆YOLOv5仓库并安装依赖。注意要指定v6.1版本因为main分支可能已经更新git clone -b v6.1 https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt验证安装是否成功时别急着跑训练先用这个小技巧测试import torch print(torch.cuda.is_available()) # 应该返回True print(torch.cuda.device_count()) # 显示可用GPU数量2. 数据准备制作高质量自定义数据集去年给工厂做零件检测项目时我深刻体会到数据质量决定模型上限。先说数据集结构YOLOv5需要这样的目录custom_dataset/ ├── images/ │ ├── train/ # 训练图片 │ └── val/ # 验证图片 └── labels/ ├── train/ # 训练标签 └── val/ # 验证标签标签格式是YOLO特色的归一化坐标每行一个对象格式为class_id x_center y_center width height比如标注一个位于图片正中央、占画面1/4大小的苹果class_id00 0.5 0.5 0.5 0.5推荐用LabelImg标注但要注意两点保存格式选YOLO而非PascalVOC类名文件要保存为dataset.yaml里的names顺序转换已有数据集时这个Python代码片段能把COCO格式转YOLOfrom pycocotools.coco import COCO import os coco COCO(annotations/instances_train2017.json) cat_ids coco.getCatIds() img_ids coco.getImgIds() for img_id in img_ids: img_info coco.loadImgs(img_id)[0] ann_ids coco.getAnnIds(imgIdsimg_id) anns coco.loadAnns(ann_ids) with open(flabels/{img_info[file_name].replace(.jpg,.txt)}, w) as f: for ann in anns: x,y,w,h ann[bbox] x_center (x w/2) / img_info[width] y_center (y h/2) / img_info[height] width w / img_info[width] height h / img_info[height] f.write(f{ann[category_id]-1} {x_center} {y_center} {width} {height}\n)3. 模型训练参数调优与性能监控启动训练前必须配置好dataset.yaml文件。这个文件就像菜谱告诉模型去哪找数据train: ../custom_dataset/images/train val: ../custom_dataset/images/val nc: 3 # 类别数 names: [apple, orange, banana] # 类别名称开始训练的命令看似简单但参数组合大有学问python train.py --img 640 --batch 16 --epochs 100 --data dataset.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt几个关键参数经验值--img分辨率越大精度越高但显存占用呈平方增长。RTX 3090可以尝试12803060建议640--batch尽可能大但不要触发OOM。可用--batch-size自动寻找最大值--epochs简单数据集50-100足够复杂场景建议300监控训练进度时别只看loss曲线。我习惯同时关注这些指标mAP0.5:IOU常规检测精度mAP0.5:0.95严格指标各类别的precision/recall发现样本不均衡遇到显存不足时试试这些方法减小--batch-size使用--multi-scale开启多尺度训练添加--adam改用Adam优化器启用梯度累积--accumulate 24. 模型部署从检测到落地应用训练完成后best.pt就是你的成果。用detect.py测试时这些参数最实用python detect.py --weights runs/train/exp/weights/best.pt --source test_images/ --conf 0.5 --iou 0.45 --imgsz 640其中--conf是置信度阈值根据场景调整安防监控0.7-0.9减少误报工业质检0.3-0.6避免漏检部署到生产环境时我推荐用TorchScript导出import torch model torch.hub.load(ultralytics/yolov5, custom, pathbest.pt) model.eval() traced_model torch.jit.trace(model, torch.randn(1, 3, 640, 640)) traced_model.save(yolov5_custom.pt)在OpenCV中调用导出的模型import cv2 net cv2.dnn.readNetFromTorch(yolov5_custom.pt) blob cv2.dnn.blobFromImage(img, 1/255.0, (640,640), swapRBTrue) net.setInput(blob) outs net.forward()最后提醒几个部署时的坑输入图像预处理必须和训练时一致归一化、BGR/RGB不同框架的NMS实现可能有差异TensorRT加速时需要重新校准anchor