当前位置：首页 > news >正文

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

news 2026/8/10 19:36:39

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 python=3.8 conda activate yolov5

安装PyTorch时最容易踩坑。建议先去PyTorch官网用他们的配置生成器，根据你的CUDA版本生成安装命令。比如我的RTX 3080是CUDA 11.3：

pip install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.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_id=0）：

0 0.5 0.5 0.5 0.5

推荐用LabelImg标注，但要注意两点：

保存格式选YOLO而非PascalVOC
类名文件要保存为dataset.yaml里的names顺序

转换已有数据集时，这个Python代码片段能把COCO格式转YOLO：

from 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(imgIds=img_id) anns = coco.loadAnns(ann_ids) with open(f'labels/{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可以尝试1280，3060建议640
--batch：尽可能大但不要触发OOM。可用--batch-size自动寻找最大值
--epochs：简单数据集50-100足够，复杂场景建议300+

监控训练进度时，别只看loss曲线。我习惯同时关注这些指标：

mAP@0.5:IOU：常规检测精度
mAP@0.5:0.95：严格指标
各类别的precision/recall：发现样本不均衡

遇到显存不足时，试试这些方法：

减小--batch-size
使用--multi-scale开启多尺度训练
添加--adam改用Adam优化器
启用梯度累积：--accumulate 2

4. 模型部署：从检测到落地应用

训练完成后，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', path='best.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), swapRB=True) net.setInput(blob) outs = net.forward()

最后提醒几个部署时的坑：