最近在配置YOLO（You Only Look Once）進(jìn)行物體檢測和圖像分割任務(wù)時，發(fā)現(xiàn)默認(rèn)安裝的情況下，YOLO使用的是CPU進(jìn)行計算。

這對于需要處理大量圖像或?qū)崟r檢測的任務(wù)來說，效率明顯不足。

本文將詳細(xì)介紹如何將YOLO從CPU模式切換到GPU模式，顯著提升運(yùn)行效率。

1. 配置步驟

1.1. 檢查當(dāng)前PyTorch是否支持GPU

首先需要確認(rèn)當(dāng)前安裝的PyTorch是否支持GPU。打開Python環(huán)境，運(yùn)行以下代碼：

import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA是否可用: {torch.cuda.is_available()}")
print(f"當(dāng)前設(shè)備: {torch.cuda.current_device() if torch.cuda.is_available() else 'CPU'}")
print(f"設(shè)備名稱: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else '無GPU設(shè)備'}")

如果輸出顯示CUDA是否可用: False，說明需要重新安裝支持GPU的PyTorch版本。

我在默認(rèn)安裝 YOLO 之后，顯示的就是False。

1.2. 卸載現(xiàn)有的torch庫

如果當(dāng)前PyTorch不支持GPU，需要先卸載相關(guān)庫：

pip uninstall torch torchvision torchaudio

1.3. 查看本機(jī)GPU情況（Windows 11系統(tǒng)）

在Windows 11系統(tǒng)中，可以通過以下方式查看GPU信息：

1. 按Win + X鍵，選擇"任務(wù)管理器"
2. 切換到"性能"選項卡
3. 查看GPU信息，確認(rèn)GPU型號和CUDA支持情況

或者使用命令行：

nvidia-smi

這將顯示NVIDIA GPU的詳細(xì)信息，包括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

從上面的命令顯示結(jié)果來看，我的CUDA Version是 12.5。

所以應(yīng)該使用如下命令安裝：

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使用效果

安裝成功后，運(yùn)行驗證代碼：

import torch
from ultralytics import YOLO

# 檢查GPU是否可用
print(f"CUDA是否可用: {torch.cuda.is_available()}")
print(f"設(shè)備名稱: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else '無GPU設(shè)備'}")

# 加載YOLO模型并指定使用GPU
model = YOLO('yolov11n.pt')  # 以YOLOv8n為例

results = model('path/to/your/test.mp4')

2. 性能對比

完成配置后，你會注意到：

• 訓(xùn)練速度：GPU訓(xùn)練通常比CPU快很多
• 推理速度：實時檢測的幀率大幅提升
• 批量處理：GPU可以并行處理更多圖像

在我的電腦上，換成GPU之后，那個test.mp4的處理速度從44秒多降到7秒多，大約快了6倍多。

我的顯卡很一般，好的顯卡效果更明顯。

3. 常見問題解決

1. CUDA版本不匹配：確保安裝的PyTorch版本與系統(tǒng)CUDA版本兼容
2. 內(nèi)存不足：如果遇到GPU內(nèi)存不足，可以減小批量大?。╞atch size）
3. 驅(qū)動問題：確保安裝了最新的NVIDIA顯卡驅(qū)動

4. 總結(jié)

通過將YOLO從CPU遷移到GPU，你可以顯著提升模型訓(xùn)練和推理的效率。

這一簡單的配置調(diào)整將為你的計算機(jī)視覺項目帶來質(zhì)的飛躍。

如果電腦有GPU，盡快替換吧！