YOLOv5环境搭建实战：对比Ubuntu 20.04下PyTorch的CUDA版与CPU-only版安装差异

圆山中庸

YOLOv5环境搭建实战：Ubuntu 20.04下PyTorch的CUDA与CPU-only版深度对比

在计算机视觉领域，YOLOv5以其卓越的实时目标检测性能广受欢迎。然而，许多开发者在环境搭建阶段就面临关键抉择：是选择GPU加速的CUDA版本，还是采用轻量级的CPU-only方案？本文将基于Ubuntu 20.04系统，从硬件准备到性能调优，全方位解析两种安装路径的核心差异。

1. 硬件准备与环境预检

1.1 GPU环境验证流程

对于计划使用CUDA加速的用户，首先需要确认NVIDIA显卡驱动状态。打开终端执行以下命令获取关键信息：

bash复制nvidia-smi

典型输出示例：

code复制+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.82.01    Driver Version: 470.82.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |
| 30%   45C    P8    15W / 180W |    200MiB /  8119MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

关键指标解读：

Driver Version：驱动版本需≥450.80.02
CUDA Version：显示显卡支持的最高CUDA版本
GPU-Util：验证显卡是否被系统识别

若未显示有效信息，需先安装NVIDIA驱动：

bash复制sudo ubuntu-drivers autoinstall

1.2 CPU环境特殊考量

纯CPU方案虽无需显卡驱动，但需注意：

内存容量≥8GB（处理高清图像建议16GB+）
建议使用SSD存储加速数据读取
检查Python版本（≥3.6.0）：

bash复制python3 --version

2. PyTorch安装路径分叉

2.1 CUDA版本精准安装

根据nvidia-smi显示的CUDA版本，选择对应PyTorch安装命令。以下是主流版本对照表：

CUDA版本	PyTorch安装命令（示例）
11.1	`pip install torch==1.10.1+cu111`
10.2	`pip install torch==1.10.1+cu102`
11.3	`pip install torch==1.13.1+cu113`

重要提示：PyTorch官网提供版本匹配工具，建议访问pytorch.org获取最新推荐命令。

2.2 CPU-only极简安装

对于无GPU设备，直接执行：

bash复制pip install torch==1.10.1+cpu torchvision==0.11.2+cpu

性能对比数据：

COCO数据集推理速度（640x640图像）
- RTX 3090：约2ms/帧
- i7-11800H（CPU）：约120ms/帧

3. 环境配置深度优化

3.1 CUDA生态完整搭建

完成PyTorch安装后，还需配置CUDA工具链：

下载对应版本CUDA Toolkit：

bash复制wget https://developer.download.nvidia.com/compute/cuda/11.4.2/local_installers/cuda_11.4.2_470.57.02_linux.run
sudo sh cuda_11.4.2_470.57.02_linux.run

配置环境变量（添加到~/.bashrc）：

bash复制export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

验证安装：

bash复制nvcc --version

3.2 CPU环境避坑指南

常见问题解决方案：

OpenMP冲突：设置线程数提升性能

python复制import os
os.environ["OMP_NUM_THREADS"] = "4"

内存不足：调整YOLOv5的--img-size参数
```
bash复制python detect.py --img-size 320
```

4. YOLOv5项目部署实战

4.1 源码获取与依赖安装

无论哪种环境，都需要以下步骤：

bash复制git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

国内用户建议：使用镜像源加速安装

阿里云：http://mirrors.aliyun.com/pypi/simple/
清华：https://pypi.tuna.tsinghua.edu.cn/simple

4.2 环境验证测试

运行官方测试脚本验证安装：

python复制import torch
from yolov5.models.experimental import attempt_load

# 设备检测
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')

# 加载示例模型
model = attempt_load('yolov5s.pt', map_location=device)
print('Environment check passed!')

5. 进阶调优策略

5.1 GPU环境性能榨取技巧

启用TensorRT加速：

bash复制python export.py --weights yolov5s.pt --include engine --device 0

使用半精度推理：

python复制model.half()  # 转换模型为FP16

5.2 CPU环境极限优化

启用OpenVINO加速：

bash复制python export.py --weights yolov5s.pt --include openvino

使用ONNX Runtime：

python复制import onnxruntime
sess = onnxruntime.InferenceSession('yolov5s.onnx')

实际项目中，我们发现在i9-12900K处理器上，通过OpenVINO优化可使推理速度提升40%。而RTX 3090配合TensorRT能达到比原生PyTorch高3倍的吞吐量。这种性能差异直接决定了方案选型——实时视频分析必须使用CUDA版本，而偶尔使用的离线图片处理可以考虑CPU方案降低成本。