ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Ubuntu系统下YOLOv5环境搭建与配置指南

Ubuntu系统下YOLOv5环境搭建与配置指南

1. YOLO环境搭建全景指南

作为计算机视觉领域最流行的目标检测算法之一,YOLO(You Only Look Once)以其实时性和高准确率著称。但在实际应用中,环境搭建往往是新手面临的第一个挑战。本文将基于Ubuntu 20.04系统,详细解析从零开始搭建YOLOv5开发环境的完整流程,涵盖硬件准备、软件依赖、环境配置到最终验证的全套方案。

特别提示:建议使用NVIDIA显卡进行YOLO开发,CPU模式虽然可行但训练速度会显著降低。本文以RTX 3060显卡为例,其他NVIDIA显卡可参考类似配置。

1.1 硬件基础要求

YOLO环境对硬件有一定要求,以下是推荐配置:

  • GPU:NVIDIA显卡(建议RTX 2060及以上)
  • 内存:16GB及以上
  • 存储:至少50GB可用空间(用于存放数据集和模型)
  • CUDA兼容性:需确认显卡支持的CUDA版本

验证显卡兼容性的命令:

nvidia-smi

输出应包含显卡型号和驱动版本信息。如果没有输出,可能需要先安装NVIDIA驱动。

1.2 软件依赖全景图

完整的YOLO开发环境包含多个层级:

  1. 底层驱动:NVIDIA显卡驱动
  2. 计算平台:CUDA + cuDNN
  3. 深度学习框架:PyTorch
  4. YOLO实现:ultralytics/yolov5
  5. 辅助工具:Python环境、OpenCV等

2. 分步环境搭建实录

2.1 NVIDIA驱动安装

首先确保系统已安装合适的NVIDIA驱动。对于Ubuntu系统,推荐使用官方PPA:

sudo add-apt-repository ppa:graphics-drivers/ppa sudo apt update ubuntu-drivers devices # 查看推荐驱动版本 sudo apt install nvidia-driver-525 # 安装推荐版本

安装完成后重启系统,再次运行nvidia-smi确认驱动正常工作。

2.2 CUDA工具包安装

YOLOv5推荐使用CUDA 11.3,以下是安装步骤:

  1. 从NVIDIA官网下载对应版本的CUDA Toolkit
  2. 执行安装命令:
wget https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda_11.3.0_465.19.01_linux.run sudo sh cuda_11.3.0_465.19.01_linux.run

安装时注意:

  • 不要选择安装驱动(已单独安装)
  • 确保勾选了CUDA Toolkit
  • 记住安装路径(默认/usr/local/cuda-11.3)

安装完成后,将CUDA加入环境变量:

echo 'export PATH=/usr/local/cuda-11.3/bin:$PATH' >> ~/.bashrc echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc source ~/.bashrc

验证安装:

nvcc --version

2.3 cuDNN安装

cuDNN是NVIDIA提供的深度学习加速库,需单独安装:

  1. 从NVIDIA开发者网站下载对应CUDA 11.3的cuDNN
  2. 解压并复制文件:
tar -xzvf cudnn-11.3-linux-x64-v8.2.1.32.tgz sudo cp cuda/include/cudnn*.h /usr/local/cuda-11.3/include sudo cp cuda/lib64/libcudnn* /usr/local/cuda-11.3/lib64 sudo chmod a+r /usr/local/cuda-11.3/include/cudnn*.h /usr/local/cuda-11.3/lib64/libcudnn*

2.4 Python环境配置

推荐使用conda创建独立Python环境:

conda create -n yolo python=3.8 conda activate yolo

安装基础依赖:

pip install numpy opencv-python matplotlib seaborn tqdm pandas

2.5 PyTorch安装

根据CUDA版本安装对应的PyTorch:

pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio==0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html

验证PyTorch能否识别GPU:

import torch print(torch.cuda.is_available()) # 应输出True print(torch.version.cuda) # 应显示11.3

3. YOLOv5项目部署

3.1 获取源码

克隆官方YOLOv5仓库:

git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt # 安装项目依赖

3.2 环境验证

使用预训练模型进行快速验证:

python detect.py --weights yolov5s.pt --source data/images/

如果一切正常,将在runs/detect/exp目录下看到检测结果图像。

4. 常见问题与解决方案

4.1 CUDA版本不匹配

错误现象:运行时提示CUDA版本不兼容 解决方法:

  1. 确认nvcc --version和torch.version.cuda显示版本一致
  2. 如果不一致,重新安装对应版本的PyTorch

4.2 显卡内存不足

错误现象:训练时出现CUDA out of memory 解决方法:

  1. 减小batch-size参数
  2. 使用更小的模型(yolov5s)
  3. 尝试混合精度训练(添加--half参数)

4.3 依赖冲突

错误现象:各种奇怪的import错误 解决方法:

  1. 创建全新的conda环境
  2. 严格按照requirements.txt安装依赖
  3. 避免混用pip和conda安装

5. 高级配置技巧

5.1 Docker方式部署

对于需要环境隔离的场景,可以使用官方Docker镜像:

docker pull ultralytics/yolov5:latest docker run -it --gpus all ultralytics/yolov5:latest

5.2 自定义数据集准备

YOLO使用的数据集格式:

  • images/:存放图像文件
  • labels/:存放对应的txt标注文件
  • 每个txt文件格式:class_id x_center y_center width height(归一化坐标)

推荐使用LabelImg工具进行标注,保存为YOLO格式。

5.3 训练参数调优

关键训练参数说明:

  • --img:输入图像尺寸(建议640)
  • --batch:批次大小(根据显存调整)
  • --epochs:训练轮次(通常100-300)
  • --data:数据集配置文件路径
  • --cfg:模型配置文件路径
  • --weights:预训练权重路径

示例训练命令:

python train.py --img 640 --batch 16 --epochs 100 --data coco128.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt

6. 生产环境部署建议

6.1 模型导出

YOLOv5支持导出多种格式:

python export.py --weights yolov5s.pt --include torchscript onnx coreml

6.2 性能优化技巧

  1. TensorRT加速:
python export.py --weights yolov5s.pt --include engine --device 0
  1. ONNX Runtime推理:
import onnxruntime sess = onnxruntime.InferenceSession("yolov5s.onnx")

6.3 长期维护建议

  1. 定期更新代码库:
git pull pip install -r requirements.txt --upgrade
  1. 监控GPU使用情况:
watch -n 1 nvidia-smi

在实际项目中,我们发现保持环境独立(使用conda或docker)可以避免大多数依赖问题。对于团队开发,建议将环境配置过程脚本化,确保所有成员使用完全一致的环境配置。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进