ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

YOLOv5环境搭建与自定义数据集训练全流程实战指南

YOLOv5环境搭建与自定义数据集训练全流程实战指南 1. 项目概述与核心价值最近在折腾一个目标检测的小项目核心需求是在Ubuntu 20.04系统上从零开始搭建YOLOv5的运行环境然后准备自己的图片数据完成标注和训练最终得到一个能识别特定目标的模型。这听起来像是深度学习入门的一个经典“Hello World”项目但实际操作起来从系统环境到模型训练每一步都可能藏着不少坑。我把自己这次完整走通的流程和踩过的坑记录下来希望能给同样想入门YOLOv5实战的朋友一个清晰的参考。整个过程会涉及Linux基础操作、Python虚拟环境管理、PyTorch框架安装、YOLOv5源码获取、数据标注工具使用以及训练调参内容比较干建议跟着步骤一步步来。2. 环境搭建从裸机到可运行的YOLOv5环境搭建是第一步也是最容易出问题的一步。一个干净、隔离的Python环境是后续所有工作的基础。2.1 系统准备与基础依赖安装我的实验环境是一台安装了Ubuntu 20.04.6 LTS的服务器没有图形界面全程通过SSH操作。如果你用的是桌面版大部分命令是通用的。首先更新系统包列表并升级现有软件包是个好习惯sudo apt update sudo apt upgrade -y接下来安装一些编译和运行所需的底层依赖库。这些库是PyTorch、OpenCV等Python包在Linux系统上正常编译和运行的基础sudo apt install -y python3-pip python3-dev build-essential libssl-dev libffi-dev sudo apt install -y libjpeg-dev libpng-dev libtiff-dev sudo apt install -y libavcodec-dev libavformat-dev libswscale-dev libv4l-dev sudo apt install -y libxvidcore-dev libx264-dev sudo apt install -y libgtk-3-dev sudo apt install -y libatlas-base-dev gfortran sudo apt install -y libhdf5-serial-dev sudo apt install -y git wget curl vim这里重点解释几个包python3-dev包含Python3的头文件和静态库编译某些Python扩展模块如PyTorch的部分C扩展时必需。libgtk-3-devOpenCV的imshow等GUI功能依赖的图形库。即使在无头服务器上如果不安装import cv2时可能会报错尽管我们可能不用GUI功能但为了兼容性最好装上。libatlas-base-dev gfortran提供基础的线性代数运算库是NumPy、SciPy等科学计算包的底层依赖。注意如果是在全新的Ubuntu 20.04上操作python3和pip3默认已经安装。但系统自带的Python版本是3.8这对于YOLOv5是完全兼容的无需升级到更高版本避免不必要的兼容性问题。2.2 创建与管理Python虚拟环境强烈建议使用虚拟环境它可以为项目创建一个独立的Python包安装空间避免与系统Python或其他项目发生冲突。我选择venv它是Python3内置的模块无需额外安装。# 创建一个名为‘yolov5_env’的虚拟环境你可以换成任何你喜欢的名字 python3 -m venv yolov5_env # 激活虚拟环境 source yolov5_env/bin/activate激活后你的命令行提示符前通常会显示环境名(yolov5_env)表示后续的所有pip install操作都只影响这个环境。实操心得我习惯把虚拟环境创建在项目根目录之外比如在~/envs/目录下统一管理。这样即使误删了项目文件夹环境配置还在。激活环境后可以先升级一下pip本身避免因版本过旧导致安装失败pip install --upgrade pip。2.3 安装PyTorch与TorchVision这是核心步骤。YOLOv5基于PyTorch框架因此必须先正确安装PyTorch及其配套的torchvision。安装命令需要根据你的硬件是否有CUDA显卡去PyTorch官网获取。1. 确认CUDA版本仅限有NVIDIA显卡的用户如果你有NVIDIA显卡并打算使用GPU加速训练这能极大缩短训练时间需要先安装NVIDIA驱动和CUDA工具包。可以通过以下命令查看驱动和CUDA版本nvidia-smi输出顶部会显示CUDA Version例如“CUDA Version: 11.7”。这个版本号决定了你应该安装哪个版本的PyTorch。2. 安装PyTorch前往 PyTorch官网 利用其配置工具生成安装命令。以我的环境CUDA 11.7为例生成的命令如下# 使用国内镜像源加速下载 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117如果没有GPU或者只想用CPU运行训练速度会非常慢仅建议用于验证流程则使用CPU版本pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu3. 验证安装安装完成后在Python交互环境中验证import torch print(torch.__version__) # 输出PyTorch版本如 2.0.0cu117 print(torch.cuda.is_available()) # 输出 True 表示GPU可用False则不可用 import torchvision print(torchvision.__version__)踩坑记录最常遇到的问题就是PyTorch版本与CUDA版本不匹配或者PyTorch的CUDA版本与系统实际安装的CUDA版本不匹配。这会导致torch.cuda.is_available()返回False。务必严格按照官网根据你的CUDA版本生成的命令来安装。如果不确定先装CPU版本跑通流程再解决GPU问题。2.4 安装YOLOv5及其他依赖现在可以克隆YOLOv5的官方仓库并安装其要求的依赖包了。# 克隆仓库如果慢可以考虑使用Gitee镜像 git clone https://github.com/ultralytics/yolov5.git cd yolov5 # 安装requirements.txt中列出的所有依赖 pip install -r requirements.txt这个requirements.txt文件包含了OpenCV-Python、Matplotlib、Pandas、Seaborn等数十个包。安装过程可能需要几分钟。注意事项requirements.txt里默认的opencv-python在某些系统上可能安装失败或缺少某些功能。如果遇到问题可以尝试安装opencv-python-headless无GUI依赖更适合服务器pip uninstall opencv-python -y pip install opencv-python-headless至此环境搭建完成。你可以运行一个简单的检测命令来测试环境是否正常python detect.py --source data/images/bus.jpg --weights yolov5s.pt这条命令会下载预训练的yolov5s.pt模型权重并对示例图片bus.jpg进行检测。如果一切正常会在runs/detect/exp目录下生成带有检测框的结果图片。3. 数据集准备从原始图片到YOLO格式模型训练的好坏七分靠数据。准备一个高质量、格式正确的数据集至关重要。YOLOv5要求的数据格式是一种特定的TXT标注文件。3.1 数据收集与目录结构规划假设我们要训练一个识别“猫”和“狗”的模型。收集图片尽可能收集多样化的图片不同角度、光照、背景、姿态数量越多越好。对于入门每个类别有几百张图片就可以开始看到效果。图片格式支持JPG、PNG等。规划目录我推荐以下目录结构清晰且符合YOLOv5的默认数据加载逻辑。datasets/ └── my_cat_dog/ ├── images/ │ ├── train/ # 存放训练集图片 │ └── val/ # 存放验证集图片 └── labels/ ├── train/ # 存放训练集标签文件 (.txt) └── val/ # 存放验证集标签文件 (.txt)你需要手动将收集到的所有图片放入images/train/然后按一定比例如8:2拆分一部分到images/val/作为验证集。或者写一个简单的Python脚本自动拆分。3.2 数据标注使用LabelImg工具我们需要为每张图片中的“猫”和“狗”画上边界框并打上标签。LabelImg是一个图形化的标注工具。安装LabelImg# 确保在虚拟环境中 pip install labelImg # 安装后直接在终端输入命令启动 labelImg标注流程打开LabelImg点击“Open Dir”选择datasets/my_cat_dog/images/train/目录。在右侧将标注格式设置为“YOLO”非常重要默认可能是PascalVOC。使用快捷键w调出画框工具框住目标物体。在弹出的对话框中输入类别名称如“cat”。软件会自动在labels/train/目录下生成同名的TXT文件。按d下一张重复直到所有训练集图片标注完成。同理标注验证集图片images/val/下的图片。生成的TXT标签文件解读每一行代表图片中的一个物体格式为class_id x_center y_center width heightclass_id: 类别索引从0开始。例如我们定义0cat, 1dog。x_center, y_center: 边界框中心的x和y坐标已归一化即除以图片宽度和高度。width, height: 边界框的宽度和高度已归一化。例如0 0.5 0.5 0.3 0.4表示一个类别为“猫”的物体其边界框中心位于图片正中央宽度占图片宽的30%高度占图片高的40%。实操心得与避坑指南标注一致性框要尽可能紧贴物体边缘避免包含太多背景。对于同一类物体框的大小和位置要尽量一致。验证集独立验证集的图片绝对不能出现在训练集中且最好由不同的人或在不同时间标注以模拟真实场景的泛化能力。标签文件对应确保每张图片如cat_001.jpg都对应一个同名的标签文件cat_001.txt即使某张图片中没有目标物体也需要一个空的TXT文件。类别文件在数据集根目录my_cat_dog/下创建一个data.yaml文件这是YOLOv5读取数据集的配置文件。内容如下# data.yaml path: ../datasets/my_cat_dog # 数据集根目录的相对路径相对于yolov5项目目录 train: images/train # 训练集图片路径相对于path val: images/val # 验证集图片路径相对于path # 类别数量和名称 nc: 2 # number of classes names: [cat, dog] # class names这个文件是连接数据和训练脚本的桥梁路径一定要写对。4. 模型训练配置、启动与监控数据准备好后就可以开始训练模型了。YOLOv5提供了非常方便的训练脚本。4.1 训练参数解析与配置训练的核心命令是train.py。我们需要理解几个关键参数python train.py \ --img 640 \ # 训练图片尺寸必须是32的倍数640是常用尺寸 --batch 16 \ # 批次大小根据GPU内存调整。内存不足就调小。 --epochs 100 \ # 训练轮数。对于小数据集100-300轮常见。 --data ./datasets/my_cat_dog/data.yaml \ # 上一步创建的配置文件路径 --weights yolov5s.pt \ # 预训练权重。从YOLOv5s模型开始微调收敛更快。 --project runs/train \ # 训练结果保存的根目录 --name cat_dog_exp1 # 本次实验的名称会在project下生成同名文件夹--weights yolov5s.pt: 使用在COCO数据集上预训练好的小模型yolov5s权重进行迁移学习。这是训练自己数据集的标准做法比从零训练快得多效果也好。YOLOv5还提供了n,m,l,x等更大更慢的模型可根据需求选择。--batch: 批次大小。这是影响训练速度和GPU内存占用的关键参数。如果训练时出现“CUDA out of memory”错误首先尝试减小--batch如改为8、4或者减小--img如改为416。--epochs: 训练总轮数。不是越多越好需要观察验证集指标防止过拟合。4.2 启动训练与过程监控在YOLOv5项目根目录下运行上述配置好的命令。训练开始后终端会输出每一轮epoch的损失loss和性能指标mAP。更重要的监控方式是使用TensorBoardYOLOv5在训练时会自动记录日志。在另一个终端窗口激活同一个虚拟环境进入项目目录启动TensorBoardtensorboard --logdir runs/train然后在浏览器中打开http://localhost:6006如果服务器无图形界面需做SSH端口转发。在这里你可以看到损失曲线train/box_loss,train/obj_loss,train/cls_loss应随着训练持续下降val/下的对应损失也应下降并最终趋于平稳。如果验证损失中途开始上升可能是过拟合。性能指标重点关注metrics/mAP_0.5和metrics/mAP_0.5:0.95。mAP是衡量检测精度的重要指标值在0到1之间越高越好。模型权重分布可以查看模型参数的直方图分布。常见问题与排查训练Loss为NaN或突然变得巨大通常是学习率--lr设置过高。YOLOv5有自适应学习率调度一般不用改。如果出现此问题可以尝试在命令中添加--lr 0.01默认是0.01来显式设置一个更小的初始学习率。mAP始终为0首要检查数据标注格式是否正确data.yaml中的路径和类别名是否正确验证集是否有标签文件检查标签确认TXT标签文件中的类别ID是否从0开始连续编号且与data.yaml中的names列表顺序一致。检查数据运行python train.py --data ... --weights yolov5s.pt --epochs 1只训练1个epoch然后查看生成的runs/train/exp*/train_batch*.jpg图片。这些图片显示了训练时数据增强后的效果并且标注框会被画上去。如果框的位置完全错误或没有框说明数据加载或标注格式有问题。GPU内存不足OOM如前所述降低--batch-size和--img-size。也可以尝试使用更小的模型如--weights yolov5n.pt。4.3 训练结果分析与模型选择训练完成后所有结果都保存在runs/train/cat_dog_exp1/目录下。其中最重要的文件是weights/best.pt在验证集上表现最好的模型权重。weights/last.pt最后一轮训练的模型权重。results.png汇总了所有损失和指标曲线的图片。confusion_matrix.png混淆矩阵查看模型在各类别上的混淆情况。如何选择模型通常选择best.pt。你可以使用验证集或自己准备的测试图片进行快速验证python detect.py --source ../datasets/my_cat_dog/images/val/ --weights runs/train/cat_dog_exp1/weights/best.pt --conf 0.5--conf 0.5表示置信度阈值高于0.5的检测框才会被显示。可以根据实际效果调整这个值。5. 模型导出与部署前测试训练出的.pt文件是PyTorch模型要在某些特定环境如某些嵌入式设备、移动端或需要更高推理速度的场景部署可能需要转换成其他格式。5.1 模型导出为ONNX格式ONNX是一种开放的模型交换格式被很多推理引擎支持。python export.py --weights runs/train/cat_dog_exp1/weights/best.pt --include onnx导出的best.onnx文件可以在ONNX Runtime、OpenCV DNN等框架中加载和推理。导出时可以指定输入图片尺寸--img 640也可以使用动态尺寸--dynamic。5.2 使用导出的模型进行推理为了验证导出是否正确可以用YOLOv5自带的代码加载ONNX模型进行推理需要安装onnxruntimeimport cv2 import numpy as np import onnxruntime as ort # 加载ONNX模型 session ort.InferenceSession(best.onnx, providers[CPUExecutionProvider]) # 使用CPU # 如果是GPU可以使用 [CUDAExecutionProvider, CPUExecutionProvider] # 准备输入数据 (示例) img cv2.imread(test.jpg) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_resized cv2.resize(img_rgb, (640, 640)) input_data img_resized.transpose(2, 0, 1).astype(np.float32) / 255.0 # HWC to CHW, 归一化 input_data np.expand_dims(input_data, axis0) # 添加batch维度 # 推理 input_name session.get_inputs()[0].name outputs session.run(None, {input_name: input_data}) # 处理输出 (YOLOv5 ONNX模型的输出格式需要根据版本解析此处为简化示例) # 通常输出是一个列表包含检测框、置信度、类别等 print(outputs[0].shape)部署注意事项不同部署后端如TensorRT, OpenVINO, CoreML, TFLite对模型有不同要求。YOLOv5的export.py脚本支持--include torchscript, onnx, engine, coreml, tflite等多种格式。选择哪种格式取决于你的目标部署平台。转换后务必在目标平台上进行严格的精度和速度测试因为转换过程可能引入微小的数值误差。6. 超参数调优与模型改进思路第一次训练得到基础模型后如果想进一步提升性能可以从以下几个方面入手6.1 数据层面改进数据增强Data AugmentationYOLOv5默认开启了强大的数据增强Mosaic MixUp等这些在train.py中通过--hyp参数指定超参数文件来控制。对于小数据集增强尤其重要。你可以尝试修改data/hyps/hyp.scratch-low.yaml文件针对小数据集中的增强参数例如调整色彩空间扰动hsv_h, hsv_s, hsv_v的强度或增加旋转、缩放、裁剪的幅度。复制一份自定义然后在训练时指定--hyp your_custom_hyp.yaml。数据质量重新审视标注数据修正不准确的框框太大、太小、位置偏、错误或遗漏的标签。脏数据对模型性能的损害是致命的。数据平衡检查各个类别的图片数量是否均衡。如果“猫”有1000张“狗”只有100张模型会严重偏向“猫”。解决方法包括为少数类别收集更多数据或在数据加载时进行过采样oversampling。6.2 模型与训练超参数调优更换模型尺度如果yolov5s精度不够但速度要求宽松可以尝试更大的模型如--weights yolov5m.pt或yolov5l.pt。更大的模型容量更高但需要更多数据和更长的训练时间也更容易过拟合。调整学习率虽然自适应调度器很好用但有时手动调整会有奇效。如果模型收敛很慢可以尝试稍微增大学习率如--lr 0.02如果训练不稳定loss震荡则减小学习率如--lr 0.001。可以使用--lr-scheduler参数更换调度策略。早停Early StoppingYOLOv5内置了早停机制--patience参数。例如--patience 50表示在验证集指标连续50轮没有提升时自动停止训练并恢复best.pt。这能有效防止过拟合节省时间。6.3 利用TensorBoard进行深度分析不要只看最终的results.png。深入分析TensorBoard中的图表对比train/box_loss和val/box_loss的曲线。如果训练损失持续下降而验证损失在某个点后开始上升这是典型的过拟合信号。需要增加数据增强、使用更小的模型或增加正则化如权重衰减--weight-decay。观察metrics/precision和metrics/recall。高精度低召回率说明模型很保守只检测它非常确信的目标漏检多低精度高召回率说明模型激进误检多。调整推理时的置信度阈值--conf可以在这两者之间取得平衡。整个流程走下来从系统环境到训练出可用的模型其实是一个不断遇到问题、排查问题、解决问题的过程。最花时间的往往不是敲命令而是数据准备和调参优化。对于新手来说我的建议是先严格按照标准流程跑通一个完整的小例子比如用公开的数据集如VOC或COCO的一小部分先走一遍理解每个环节的输出和意义然后再处理自己的数据这样会顺利很多。
返回列表