最近在配置YOLO
(You Only Look Once)进行物体检测和图像分割任务时,发现默认安装的情况下,YOLO
使用的是CPU进行计算。
这对于需要处理大量图像或实时检测的任务来说,效率明显不足。
本文将详细介绍如何将YOLO
从CPU
模式切换到GPU
模式,显著提升运行效率。
1. 配置步骤
1.1. 检查当前PyTorch是否支持GPU
首先需要确认当前安装的PyTorch是否支持GPU。打开Python环境,运行以下代码:
import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")
print(f"当前设备: {torch.cuda.current_device() if torch.cuda.is_available() else 'CPU'}")
print(f"设备名称: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else '无GPU设备'}")
如果输出显示CUDA是否可用: False
,说明需要重新安装支持GPU的PyTorch版本。
我在默认安装 YOLO 之后,显示的就是False
。
1.2. 卸载现有的torch库
如果当前PyTorch不支持GPU,需要先卸载相关库:
pip uninstall torch torchvision torchaudio
1.3. 查看本机GPU情况(Windows 11系统)
在Windows 11系统中,可以通过以下方式查看GPU信息:
- 按
Win + X
键,选择"任务管理器" - 切换到"性能"选项卡
- 查看GPU信息,确认GPU型号和CUDA支持情况
或者使用命令行:
nvidia-smi
这将显示NVIDIA GPU的详细信息,包括CUDA版本。
我的电脑显示信息如下:
Sat Sep 27 17:35:25 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 556.12 Driver Version: 556.12 CUDA Version: 12.5 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Driver-Model | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 2060 WDDM | 00000000:01:00.0 On | N/A |
| N/A 35C P8 14W / 80W | 937MiB / 6144MiB | 10% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
1.4. 安装匹配的GPU版本PyTorch
从上面的命令显示结果来看,我的CUDA Version
是 12.5
。
所以应该使用如下命令安装:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu125
但是,目前似乎没有提供cu125
的版本,上面的命令会报错,于是安装了cu121
版本。
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
至此,安装成功。
1.5. 验证GPU使用效果
安装成功后,运行验证代码:
import torch
from ultralytics import YOLO# 检查GPU是否可用
print(f"CUDA是否可用: {torch.cuda.is_available()}")
print(f"设备名称: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else '无GPU设备'}")# 加载YOLO模型并指定使用GPU
model = YOLO('yolov11n.pt') # 以YOLOv8n为例results = model('path/to/your/test.mp4')
2. 性能对比
完成配置后,你会注意到:
- 训练速度:GPU训练通常比CPU快很多
- 推理速度:实时检测的帧率大幅提升
- 批量处理:GPU可以并行处理更多图像
在我的电脑上,换成GPU
之后,那个test.mp4
的处理速度从44秒
多降到7秒
多,大约快了6倍
多。
我的显卡很一般,好的显卡效果更明显。
3. 常见问题解决
- CUDA版本不匹配:确保安装的PyTorch版本与系统CUDA版本兼容
- 内存不足:如果遇到GPU内存不足,可以减小批量大小(batch size)
- 驱动问题:确保安装了最新的NVIDIA显卡驱动
4. 总结
通过将YOLO
从CPU
迁移到GPU
,你可以显著提升模型训练和推理的效率。
这一简单的配置调整将为你的计算机视觉项目带来质的飞跃。
如果电脑有GPU,尽快替换吧!