让YOLO飞起来：从CPU到GPU的配置指南

最近在配置YOLO（You Only Look Once）进行物体检测和图像分割任务时，发现默认安装的情况下，YOLO使用的是CPU进行计算。
这对于需要处理大量图像或实时检测的任务来说，效率明显不足。
本文将详细介绍如何将YOLO从CPU模式切换到GPU模式，显著提升运行效率。
1. 配置步骤

1.1. 检查当前PyTorch是否支持GPU

首先需要确认当前安装的PyTorch是否支持GPU。打开Python环境，运行以下代码：

1. import torch
2. print(f"PyTorch版本: {torch.__version__}")
3. print(f"CUDA是否可用: {torch.cuda.is_available()}")
4. print(f"当前设备: {torch.cuda.current_device() if torch.cuda.is_available() else 'CPU'}")
5. 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，需要先卸载相关库：

1. pip uninstall torch torchvision torchaudio

复制代码

1.3. 查看本机GPU情况（Windows 11系统）

在Windows 11系统中，可以通过以下方式查看GPU信息：

• 按Win + X键，选择"任务管理器"
  
• 切换到"性能"选项卡
  
• 查看GPU信息，确认GPU型号和CUDA支持情况
  

或者使用命令行：

1. nvidia-smi

复制代码

这将显示NVIDIA GPU的详细信息，包括CUDA版本。
我的电脑显示信息如下：

1. Sat Sep 27 17:35:25 2025
2. +-----------------------------------------------------------------------------------------+
3. | NVIDIA-SMI 556.12                 Driver Version: 556.12         CUDA Version: 12.5     |
4. |-----------------------------------------+------------------------+----------------------+
5. | GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
6. | Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
7. |                                         |                        |               MIG M. |
8. |=========================================+========================+======================|
9. |   0  NVIDIA GeForce RTX 2060      WDDM  |   00000000:01:00.0  On |                  N/A |
10. | N/A   35C    P8             14W /   80W |     937MiB /   6144MiB |     10%      Default |
11. |                                         |                        |                  N/A |
12. +-----------------------------------------+------------------------+----------------------+

复制代码

1.4. 安装匹配的GPU版本PyTorch

从上面的命令显示结果来看，我的CUDA Version是 12.5。
所以应该使用如下命令安装：

1. pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu125

复制代码

但是，目前似乎没有提供cu125的版本，上面的命令会报错，于是安装了cu121版本。

1. pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

复制代码

至此，安装成功。
1.5. 验证GPU使用效果

安装成功后，运行验证代码：

1. import torch
2. from ultralytics import YOLO
4. # 检查GPU是否可用
5. print(f"CUDA是否可用: {torch.cuda.is_available()}")
6. print(f"设备名称: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else '无GPU设备'}")
8. # 加载YOLO模型并指定使用GPU
9. model = YOLO('yolov11n.pt')  # 以YOLOv8n为例
11. 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，尽快替换吧！

来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！