使用 Conda 安装 PyTorch GPU版本

0. 环境

• Windows 11
• Python 3.12.4
• GPU: NVIDIA GeForce RTX 4060
• CUDA 12.4
• 曾经使用pip安装过PyTorch

1. 创建虚拟环境

本部即将使用Conda创建一个虚拟环境，名为pytorch-gpu

conda create -n pytorch-gpu python=3.12.4
conda activate pytorch-gpu

2. 安装PyTorch GPU版本

请访问这个链接，选择合适的配置，然后复制安装命令。

你可以使用以下命令查看CUDA版本：

nvcc --version

此时会显示如下信息：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Mar_28_02:30:10_Pacific_Daylight_Time_2024
Cuda compilation tools, release 12.4, V12.4.131
Build cuda_12.4.r12.4/compiler.34097967_0

可以看到我的CUDA版本是12.4，所以我选择了CUDA 12.4版本的PyTorch。

conda install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

3. 测试

import torch
print(torch.cuda.is_available())

如果输出True，则说明安装成功。

此时我遇到的问题

经过以上测试，我发现输出为False，经过我的排查，在尝试以下命令后，发现异常

import torch
print(torch.__version__)

正常情况下，输出结果应该是我选择的2.5.1，但是实际上输出的是2.4.1。我回想起之前使用pip安装过PyTorch，因此我卸载了用pip安装的PyTorch。

4. 卸载pip安装的PyTorch

pip uninstall torch

后重新运行测试代码torch.cuda.is_available()，发现输出为True，说明安装成功。

此时我遇到的问题

我在运行YOLO时，发现一个奇怪的报错：

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend.

这时候我意识到，torchvision也是用pip安装的，所以我卸载了torchvision。

5. 卸载pip安装的torchvision 与 torchaudio （如有）

pip uninstall torchvision
pip uninstall torchaudio

经过以上内容，就成功安装了PyTorch GPU版本。

使用torch.cuda.is_available()也会返回True，YOLO也能正常训练。