1. 错误原因分析
- 根本问题:当前安装的PyTorch是CPU版本,无法调用GPU硬件加速。
- 当运行以下代码时会报错:
model = YOLO("yolov8n.pt").to("cuda") # 或 .cuda()
2. 解决方案步骤
步骤1:验证CUDA可用性
在Python中运行以下命令检查环境:
import torch
print(torch.__version__) # 查看PyTorch版本
print(torch.cuda.is_available()) # 必须返回True
如果输出为False
,说明PyTorch未启用CUDA支持。
步骤2:重新安装支持CUDA的PyTorch
根据您的CUDA版本(通过nvcc --version
查看),选择对应命令安装:
CUDA版本 | 安装命令 |
---|---|
CUDA 12.x | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 |
CUDA 11.8 | pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 |
CPU版本 | pip install torch torchvision torchaudio |
验证安装:
import torch
assert torch.cuda.is_available() # 不报错则成功
步骤3:重新安装Ultralytics YOLOv8
确保YOLO库与PyTorch兼容:
pip uninstall ultralytics -y
pip install ultralytics
3. 备用方案(如果无NVIDIA GPU)
如果设备确实没有GPU,强制使用CPU运行:
model = YOLO("yolov8n.pt").to("cpu") # 显式指定CPU
# 或者直接
model = YOLO("yolov8n.pt") # 默认自动选择设备
4. 常见问题排查
问题现象 | 解决方案 |
---|---|
nvcc --version 找不到 | 安装CUDA Toolkit:NVIDIA官网下载 |
版本冲突 | 创建干净的conda环境:conda create -n yolo python=3.8 conda activate yolo |
驱动过旧 | 更新GPU驱动:驱动下载 |
5. 完整正确代码示例
import torch
from ultralytics import YOLO# 确认CUDA可用
assert torch.cuda.is_available(), "CUDA不可用,请检查PyTorch安装!"# 加载模型并启用GPU
model = YOLO("yolov8n.pt").to("cuda")# 运行推理
results = model("image.jpg") # 自动使用GPU加速
6. 性能对比(GPU vs CPU)
设备 | 推理速度 (FPS) | 显存占用 |
---|---|---|
RTX 3090 | 120+ | ~2GB |
CPU (i9-12900K) | 8-12 | 无 |
通过正确配置CUDA环境,将获得10倍以上的速度提升!