蕴铭 發表於 2025-9-27 18:35:00

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

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