资讯动态

Ubuntu 20.04下YOLOv5环境搭建与自定义数据集训练全流程指南

发布时间:2026/8/22 5:14:28 来源:尧图企业网站定制
1. 项目概述从零到一构建YOLOv5实战环境搞深度学习的朋友尤其是刚入坑目标检测的十有八九都听说过YOLOv5。它凭借其出色的速度与精度平衡以及相对友好的工程化实现成为了很多项目落地的首选。但说实话第一次在Ubuntu上搭环境、训自己的数据这个过程里的小坑小洼可不少。今天我就以Ubuntu 20.04 LTS这个经典的、稳定的系统为基础带大家走一遍完整的流程从系统准备、环境搭建到准备自己的数据集、修改配置文件最后完成训练和简单测试。整个过程我会尽量把每一步的原理和“为什么这么做”讲清楚特别是那些官方文档一笔带过但实际操作中能让你卡半天的细节。无论你是用物理机、虚拟机还是云服务器这套流程的核心都是相通的。2. 核心思路与准备工作2.1 为什么选择Ubuntu 20.04 YOLOv5这个组合首先得说说选型。Ubuntu 20.04 LTS是长期支持版本官方维护到2025年这意味着系统底层稳定软件源丰富社区支持好遇到问题容易搜到解决方案。对于深度学习开发来说稳定压倒一切你肯定不想在训练到一半的时候因为系统更新导致CUDA驱动出问题。YOLOv5虽然并非YOLO原作者作品但它用PyTorch框架重写代码结构清晰文档相对完善而且活跃的社区持续在优化和添加新特性。相比于YOLOv3/v4的Darknet框架PyTorch对研究者更友好调试、修改模型结构都更方便。它的设计目标就是让工业部署和学术研究都能快速上手。所以这个组合可以看作是在“系统稳定性”和“算法工程友好度”之间取了一个很好的平衡点。2.2 硬件与基础软件清单在开始之前请确保你的环境满足以下最低要求我强烈建议硬件配置越高越好这能极大节省你的时间。硬件CPU至少4核。用于数据加载和预处理。内存至少8GB推荐16GB以上。加载大型数据集如COCO时很吃内存。GPU这是核心。必须是一块支持CUDA的NVIDIA显卡。显存至少4GB如GTX 1650用于训练自己的数据集6GB或以上如RTX 2060, RTX 3060体验会好很多。显存大小直接决定了你能使用的批处理大小batch size影响训练速度。存储至少50GB可用空间。系统、环境、代码、数据集和训练生成的权重文件都会占用空间。软件操作系统Ubuntu 20.04.6 LTS桌面版或服务器版均可。本文以桌面版为例操作有图形界面更直观。NVIDIA显卡驱动必须安装版本需要与后续的CUDA Toolkit匹配。CUDA ToolkitNVIDIA的并行计算平台深度学习框架的底层依赖。YOLOv5官方推荐CUDA 11.3及以上我们选择广泛兼容的CUDA 11.8。cuDNNNVIDIA深度神经网络加速库CUDA的扩展。Python3.8或3.9。Ubuntu 20.04默认可能安装3.8这是一个稳妥的选择。PyTorch深度学习框架。需要安装与CUDA版本对应的PyTorch。其他依赖通过pip安装的Python包如opencv-python,matplotlib,pycocotools等。注意整个安装流程遵循“驱动 - CUDA - cuDNN - PyTorch”的依赖顺序顺序错了会导致各种诡异问题。3. 深度学习环境搭建全流程3.1 系统更新与基础工具安装首先打开终端CtrlAltT让我们把系统更新到最新状态并安装一些必要的编译工具和包管理工具。sudo apt update sudo apt upgrade -y sudo apt install -y build-essential cmake unzip pkg-config 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 python3-dev python3-pip sudo apt install -y git wget curl这里安装的build-essential、cmake是编译工具libjpeg-dev等是图像处理库的开发文件libgtk-3-dev是OpenCV GUI组件可能需要的前端依赖python3-dev和python3-pip则是Python开发环境的基础。一次装好避免后续编译各种库时缺这少那。3.2 NVIDIA驱动与CUDA安装这是最关键也最容易出错的一步。我推荐使用Ubuntu官方仓库的ubuntu-drivers工具来安装驱动相对省心。首先添加显卡驱动PPA并安装工具sudo add-apt-repository ppa:graphics-drivers/ppa sudo apt update sudo apt install ubuntu-drivers-common检测并安装推荐驱动ubuntu-drivers devices这个命令会列出所有可用的驱动版本并标出一个“recommended”的版本。例如输出可能显示driver : nvidia-driver-550 - third-party free recommended。直接安装这个推荐版本sudo apt install nvidia-driver-550安装完成后必须重启系统。sudo reboot验证驱动安装 重启后在终端输入nvidia-smi如果安装成功你会看到一个表格显示GPU型号、驱动版本、CUDA版本这里是驱动内嵌的CUDA版本如12.4它决定了最高可支持的工具包版本以及GPU使用情况。记下右上角显示的“CUDA Version”例如“12.4”。安装CUDA Toolkit 11.8 虽然驱动支持12.4但为了与PyTorch稳定版本最佳匹配我们安装CUDA 11.8。访问 NVIDIA CUDA Toolkit Archive 选择“CUDA Toolkit 11.8.0”。 根据你的系统选择对应的安装方式。这里我们用runfile本地安装因为它更灵活可以只安装Toolkit而不覆盖现有驱动。wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run sudo sh cuda_11.8.0_520.61.05_linux.run运行安装程序后会有一个文本界面。非常重要使用方向键和回车键进行操作。首先按空格键取消勾选“Driver”因为我们已经安装了更新的驱动。确保“CUDA Toolkit 11.8”是选中的。在“Options”里可以自定义安装路径但建议保持默认。然后选择“Install”开始安装。配置环境变量 安装完成后需要将CUDA的路径添加到系统环境变量中。echo export PATH/usr/local/cuda-11.8/bin${PATH::${PATH}} ~/.bashrc echo export LD_LIBRARY_PATH/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH::${LD_LIBRARY_PATH}} ~/.bashrc source ~/.bashrc验证CUDA安装nvcc --version此时应显示nvcc: NVIDIA (R) Cuda compiler driver版本为11.8。3.3 安装cuDNNcuDNN需要从NVIDIA开发者网站下载需要注册账号免费。登录后找到对应CUDA 11.x的cuDNN版本例如cuDNN for Linux x86_64 (Tar)。下载后解压并复制文件到CUDA目录# 假设下载的文件为 cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-11.8/include sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-11.8/lib64 sudo chmod ar /usr/local/cuda-11.8/include/cudnn*.h /usr/local/cuda-11.8/lib64/libcudnn*3.4 创建Python虚拟环境并安装PyTorch强烈建议使用虚拟环境如conda或venv来管理项目依赖避免包冲突。这里我们用venv它是Python标准库的一部分。# 创建名为yolov5的虚拟环境 python3 -m venv yolov5_env # 激活虚拟环境 source yolov5_env/bin/activate激活后你的命令行提示符前会出现(yolov5_env)字样。接下来安装PyTorch。前往 PyTorch官网 选择对应配置Stable (1.13.1) Linux Pip Python CUDA 11.8。官网会给出安装命令类似于pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118注意这里的cu118就对应CUDA 11.8。安装完成后验证PyTorch是否能识别GPUpython3 -c import torch; print(torch.__version__); print(torch.cuda.is_available())如果输出类似1.13.1cu118和True恭喜你PyTorch GPU环境配置成功。3.5 克隆YOLOv5仓库并安装剩余依赖现在可以拉取YOLOv5的代码了。git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txtrequirements.txt里包含了opencv-python,pillow,matplotlib,pandas,seaborn,scipy,tqdm,pycocotools等所有必要依赖。安装过程可能需要几分钟。实操心得安装pycocotools时如果报错通常是缺少Python开发文件。可以尝试先安装sudo apt install python3.8-dev根据你的Python版本调整然后再pip install pycocotools。另外网络问题可能导致某些包下载慢或失败可以考虑配置pip国内镜像源。至此一个完整的YOLOv5开发环境就搭建好了。你可以运行一个简单的测试来验证一切正常python detect.py --source data/images/bus.jpg --weights yolov5s.pt这条命令会下载预训练的yolov5s.pt模型小型模型并对示例图片bus.jpg进行推理。如果看到终端输出检测结果并在runs/detect/exp目录下生成了带标注框的图片说明环境完全没问题。4. 准备与制作自定义数据集环境搭好了接下来就是喂给模型“食物”——数据。用公开数据集练手可以但最终总要解决自己的问题。制作一个格式正确、标注规范的数据集是训练成功的一半。4.1 数据收集与目录结构规划假设我们要做一个“安全帽检测”项目。你需要收集包含“戴安全帽”和“未戴安全帽”人员的图片。数据可以从网上公开数据集获取或者自己拍摄。数量上每个类别至少几百张越多越好且要保证多样性不同场景、光照、角度、遮挡。在YOLOv5项目根目录外建立一个清晰的数据集文件夹结构我习惯这样组织custom_dataset/ ├── images/ │ ├── train/ # 训练集图片 │ │ ├── 001.jpg │ │ ├── 002.jpg │ │ └── ... │ └── val/ # 验证集图片 │ ├── 101.jpg │ ├── 102.jpg │ └── ... └── labels/ ├── train/ # 训练集标签文件与images/train/一一对应 │ ├── 001.txt │ ├── 002.txt │ └── ... └── val/ # 验证集标签文件 ├── 101.txt ├── 102.txt └── ...关键点images和labels下的train、val子目录必须严格对应即images/train/001.jpg的标签文件是labels/train/001.txt。通常按照8:2或9:1的比例随机划分训练集和验证集。4.2 数据标注与YOLO格式解析你需要一个标注工具如LabelImg、CVAT或Roboflow。这里以LabelImg为例它可以直接生成YOLO格式的标签。安装LabelImgpip install labelImg labelImg标注流程打开软件设置图片目录为custom_dataset/images/train/。更改保存目录为custom_dataset/labels/train/。在菜单栏选择标注格式为YOLO。创建标签类别例如helmet戴安全帽和no_helmet未戴安全帽。对每张图片画框并选择对应类别。标注完成后每个.txt标签文件的内容格式如下object-class x_center y_center width heightobject-class: 物体的类别索引从0开始。例如helmet对应0no_helmet对应1。x_center y_center width height: 物体边界框的中心点x、y坐标以及框的宽度和高度。这些值都是相对于图片宽度和高度的归一化值范围0-1。例如一张640x480的图片上一个框的左上角在(100, 120)右下角在(300, 360)那么宽度 300 - 100 200高度 360 - 120 240中心点x (100 300) / 2 / 640 400 / 2 / 640 0.3125中心点y (120 360) / 2 / 480 480 / 2 / 480 0.5归一化宽度 200 / 640 0.3125归一化高度 240 / 480 0.5 标签行就是0 0.3125 0.5 0.3125 0.5。注意事项标注质量至关重要。框要紧密贴合物体不要包含太多背景对于被遮挡的物体尽量根据可见部分标注类别要统一。错误的标注是导致模型loss降不下来、性能差的最常见原因之一。4.3 创建数据集配置文件YOLOv5需要一个.yaml文件来告诉它数据集在哪里、有哪些类别。在YOLOv5项目根目录下的data/文件夹里我们创建一个helmet.yaml文件。# helmet.yaml # 训练和验证图像的路径相对于YOLOv5根目录 train: ../custom_dataset/images/train/ val: ../custom_dataset/images/val/ # 类别数量 nc: 2 # 类别名称列表 names: [helmet, no_helmet]路径可以是绝对路径也可以是像上面这样相对于YOLOv5根目录的路径。nc必须和你定义的类别数完全一致。5. 模型训练与调参实战5.1 选择与修改模型配置文件YOLOv5提供了不同大小的模型yolov5s.pt小、yolov5m.pt中、yolov5l.pt大、yolov5x.pt超大。s最快但精度最低x最慢但精度最高。对于新手和中等规模数据集从yolov5s开始是个好选择。我们需要修改模型配置文件以适应自定义的类别数。配置文件在models/目录下。复制一份yolov5s.yaml并重命名cp models/yolov5s.yaml models/helmet_yolov5s.yaml然后编辑models/helmet_yolov5s.yaml只修改一个关键参数# parameters nc: 2 # 将原来的80改为你的类别数2 ... # 后面的anchors和backbone结构保持不变5.2 启动训练与核心参数解析万事俱备开始训练。使用train.py脚本。一个最基本的训练命令如下python train.py --img 640 --batch 16 --epochs 100 --data data/helmet.yaml --cfg models/helmet_yolov5s.yaml --weights yolov5s.pt --name helmet_exp1让我拆解一下这些核心参数--img 640输入图像的大小会被自动缩放到这个尺寸。YOLOv5训练时使用多尺度训练但这是基础尺寸。更大的尺寸如1280可能提升精度但显著增加显存消耗和训练时间。--batch 16批处理大小。一次迭代送入模型的图片数量。这个值受GPU显存限制。如果出现“CUDA out of memory”错误首先降低--batch如改为8或4。它直接影响训练速度和梯度稳定性。--epochs 100训练轮数。整个数据集被完整遍历一次称为一个epoch。需要多少轮取决于数据集大小和复杂度通常从100开始观察loss曲线决定是否增加。--data data/helmet.yaml指向我们创建的数据集配置文件。--cfg models/helmet_yolov5s.yaml指向我们修改后的模型结构配置文件。--weights yolov5s.pt加载预训练权重。强烈建议使用预训练权重它是在COCO等大型数据集上训练过的包含了通用的特征提取能力能极大加速收敛并提升最终性能即迁移学习。--name helmet_exp1本次实验的名称。所有输出日志、权重、图表都会保存在runs/train/helmet_exp1目录下。训练开始后终端会输出每个epoch的训练损失和验证损失以及mAP等指标。同时TensorBoard日志也会生成你可以用以下命令实时查看更丰富的可视化信息tensorboard --logdir runs/train然后在浏览器打开http://localhost:6006可以看到损失曲线、精度曲线、权重分布等这是监控训练进程、诊断问题的重要工具。5.3 训练过程监控与关键指标解读训练过程中要重点关注以下几个指标损失函数Losstrain/box_loss,train/obj_loss,train/cls_loss分别代表边界框回归损失、目标置信度损失和分类损失。一个健康的训练过程这些损失应该随着epoch增加而稳步下降最终趋于平缓。如果损失剧烈震荡或迟迟不降可能是学习率太高、批处理大小太小或数据有问题。val/下的损失验证集上的损失。理想情况下它应该随训练损失同步下降且最终值略高于训练损失。如果验证损失很早就开始上升而训练损失持续下降这是典型的过拟合现象说明模型只记住了训练集没有泛化能力。精度指标Metricsmetrics/mAP_0.5在IoU交并比阈值为0.5时的平均精度mean Average Precision。这是目标检测的核心指标值越接近1越好。metrics/mAP_0.5:0.95在IoU阈值从0.5到0.95步长0.05区间内的平均mAP是更严格的指标。metrics/precision,metrics/recall精确率和召回率。需要根据你的应用场景权衡。例如安全监控中可能希望召回率高不漏掉任何未戴安全帽的人即使精确率稍低有一些误报。学习率Learning Rate YOLOv5默认使用余弦退火等动态调整学习率的策略。在TensorBoard中可以看到学习率的变化曲线。如果损失下降很慢可以尝试增大初始学习率通过--lr0参数如果损失震荡可以尝试减小它。5.4 超参数调优浅谈YOLOv5的hyp.scratch.yaml文件在data/目录下定义了一组默认的超参数如学习率、动量、权重衰减、数据增强参数等。对于新手我建议先使用默认超参数完成一次完整的训练得到一个基准模型。如果对基准模型的效果不满意再考虑调优。调优不是玄学要有针对性模型收敛慢可以适当增大--lr0初始学习率或检查数据标注质量。模型过拟合可以增加数据增强的强度在hyp文件中调整hsv_h,hsv_s,hsv_v,degrees,translate,scale,shear等参数或者使用更小的模型如从yolov5l换到yolov5m或者添加正则化增大权重衰减weight_decay。模型欠拟合可以减小数据增强强度使用更大的模型或者训练更多轮次。实操心得调参时一次只改变一个变量并记录每次实验的配置和结果。使用--name参数给每次实验起不同的名字方便对比。盲目同时调整多个参数会让你无法判断到底是哪个改动起了作用。6. 模型验证、测试与部署准备6.1 验证最佳模型训练结束后在runs/train/helmet_exp1/weights/目录下你会看到两个关键的权重文件best.pt在验证集上表现最好的权重根据metrics/mAP_0.5:0.95或你指定的指标。last.pt最后一个epoch训练完的权重。通常我们使用best.pt进行后续的验证和测试。使用val.py脚本在验证集上评估模型python val.py --weights runs/train/helmet_exp1/weights/best.pt --data data/helmet.yaml --img 640这个命令会输出详细的评估结果包括各个类别的精确率、召回率、mAP等并生成一个包含预测框的图片样本在runs/val/exp目录下方便你直观查看模型在验证集上的表现。6.2 使用模型进行推理用训练好的模型对新图片或视频进行检测使用detect.py脚本# 检测单张图片 python detect.py --weights runs/train/helmet_exp1/weights/best.pt --source path/to/your/test_image.jpg --conf 0.25 # 检测一个目录下的所有图片 python detect.py --weights runs/train/helmet_exp1/weights/best.pt --source path/to/test/folder/ --conf 0.25 # 检测视频 python detect.py --weights runs/train/helmet_exp1/weights/best.pt --source path/to/your/video.mp4 --conf 0.25--conf 0.25置信度阈值。只有预测框的置信度高于此值的才会被显示。你可以根据需求调整调高会减少误报但可能漏检调低则相反。结果会保存在runs/detect/exp目录下。6.3 模型导出与部署考虑为了将模型部署到生产环境如服务器、边缘设备通常需要将PyTorch模型.pt转换成更高效的格式。YOLOv5提供了方便的导出脚本python export.py --weights runs/train/helmet_exp1/weights/best.pt --include onnx engine--include onnx导出为ONNX格式。ONNX是一种开放的模型交换格式被许多推理引擎如OpenVINO, TensorRT支持。--include engine需要配合TensorRT直接导出为TensorRT引擎文件.engine这通常在NVIDIA Jetson等边缘设备上使用。导出的ONNX模型可以被C、Python等多种语言加载便于集成到各种应用系统中。7. 常见问题与故障排除实录在这一部分我汇总了搭建和训练过程中最常遇到的几个“坑”及其解决方案。7.1 环境搭建类问题问题1运行nvidia-smi提示“Command not found”或没有输出GPU信息。原因NVIDIA驱动未正确安装或未加载。解决确认是否已按照步骤安装驱动并重启。使用ubuntu-drivers devices再次检查推荐驱动并安装。使用dkms status检查DKMS模块状态或尝试sudo modprobe nvidia手动加载驱动。如果是云服务器请确认实例类型是否配备了GPU如NVIDIA T4, V100等。问题2PyTorch安装后torch.cuda.is_available()返回False。原因PyTorch版本与CUDA版本不匹配或者CUDA/cuDNN安装有问题。解决在Python中执行import torch; print(torch.version.cuda)查看PyTorch编译时使用的CUDA版本。确保其与你安装的CUDA Toolkit版本nvcc --version一致或兼容。严格按照驱动-CUDA-cuDNN-PyTorch的顺序安装并确保环境变量PATH,LD_LIBRARY_PATH设置正确。尝试完全卸载PyTorch后从PyTorch官网获取与你的CUDA版本精确对应的安装命令重装。问题3安装requirements.txt时pycocotools编译失败。原因缺少系统依赖或Python头文件。解决# 安装系统依赖 sudo apt install -y python3-dev # 或者指定Python版本如python3.8-dev # 然后尝试单独安装 pip install pycocotools如果还不行可以尝试安装预编译的wheel文件或者使用pip install githttps://github.com/philferriere/cocoapi.git#eggpycocotoolssubdirectoryPythonAPI从源码安装。7.2 训练过程类问题问题4训练一开始就报错“CUDA out of memory”。原因批处理大小--batch或图像尺寸--img设置过大超出了GPU显存容量。解决首要降低--batch将其减半如16-88-4直到不报错。如果--batch已经降到1还不行尝试减小--img如640-512。注意这会改变输入分辨率可能影响模型精度。使用更小的模型例如从yolov5l.yaml换到yolov5s.yaml。在train.py中尝试使用--device 0如果有多卡来指定单卡或检查是否有其他进程占用了显存。问题5训练Loss不下降或者mAP始终为0。原因这是最复杂的问题可能原因很多。排查步骤检查数据这是最常见的原因。运行python train.py --data your_data.yaml --img 64 --batch 8 --epochs 1进行一个极短的训练然后查看TensorBoard中的“Images”标签页或查看runs/train/exp下的训练样本图片。确认数据增强后的图片和标签框是否对齐、类别是否正确。标签文件格式归一化坐标是否正确检查配置文件确认data.yaml中的nc类别数和names列表是否与你的数据集完全匹配。确认models/*.yaml中的nc是否已修改。检查学习率默认学习率可能对你的数据集不合适。尝试使用--lr0参数微调例如设为0.01或0.001。检查预训练权重确保--weights参数正确指向了预训练模型如yolov5s.pt并且下载完整。可以尝试从官方仓库重新下载。简化问题用极小的子数据集如10张图和很少的轮数如5个epoch跑一下看loss是否有任何变化。如果连小数据都学不会那问题肯定出在数据或配置上。问题6验证集Loss远高于训练集Loss且差距越来越大过拟合。原因模型过于复杂记住了训练集的噪声而无法泛化到新数据。解决增加数据收集更多、更多样化的训练数据。增强数据增强在hyp.yaml中适当增大hsv_h,hsv_s,hsv_v色彩抖动degrees旋转translate平移scale缩放shear剪切等参数。使用正则化在hyp.yaml中增大weight_decay参数权重衰减。早停Early Stopping监控验证集loss当其在连续多个epoch不再下降时手动停止训练。换用更小的模型。7.3 推理与导出类问题问题7导出的ONNX模型在OpenVINO或TensorRT中推理速度慢或出错。原因导出设置或后续优化步骤有问题。解决确保导出时指定了正确的操作集版本例如--opset 12。对于TensorRT最好使用TensorRT提供的trtexec工具或Python API对ONNX模型进行进一步的优化和序列化生成.engine文件才能获得最佳性能。检查目标部署平台支持的算子是否与模型兼容。YOLOv5的某些层如Focus在v6.0以后已被替换可能需要特殊处理或插件支持。整个流程走下来从裸机系统到训练出自己的目标检测模型涉及的系统、软件、算法知识面很广。最关键的是保持耐心遇到问题善用搜索引擎错误信息直接复制去搜多查阅官方Issue和社区讨论。每一次排错的过程都是对这套技术栈理解加深的过程。我的建议是先严格按照流程跑通一遍获得一个能工作的模型建立信心。然后再去深入研究数据增强、模型结构、超参数调优等更深入的课题逐步迭代优化你的模型。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价