如何告诉PyTorch使用哪个CUDA版本？

Question

我在 Ubuntu 16.04 上安装了两个版本的 CUDA：9.0 和 10.1。 它们分别位于路径 /usr/local/cuda-9.0 和 /usr/local/10.1。 如果我通过 pip 安装 PyTorch 1.6.0（需要 CUDA 10.1）（pip install torch==1.6.0），它将使用版本 9.0，因此检测不到 GPU。我已经将 LD_LIBRARY_PATH 更改为 "/usr/local/cuda-10.1/lib64:/usr/local/cuda-10.1/cuda/extras/CUPTI/lib64"，但是 PyTorch 仍然使用 CUDA 9.0。 如何告诉 PyTorch 使用 CUDA 10.1？

- infochip

---

PyTorch在通过pip或conda安装时不使用系统cuda，它会随附自己的cuda版本。您如何确定您的PyTorch正在使用cuda 9.0？您拥有哪种型号的GPU？ - jodag

我使用 torch._C._cuda_getDriverVersion() 检查了 CUDA 版本，结果为 9000。如果 PyTorch 自带 CUDA，为什么会使用不兼容的版本呢？ - infochip

我的GPU是一款NVIDIA GeForce GTX 1080。 - infochip

torch._C._cuda_getDriverVersion() 不是 pytorch 使用的 cuda 版本，它是 GPU 驱动程序支持的最新版本的 cuda（应该与 nvidia-smi 报告的版本相同）。它返回的值意味着您的驱动程序已过时。您需要更新显卡驱动程序以使用 cuda 10.1。实际上被 pytorch 使用的 cuda 版本可以通过 torch.version.cuda 查询（假设确实正在使用）。有关使用 cuda 支持安装 pytorch 的系统要求的更多信息，请参见此答案。 - jodag

我猜您在安装CUDA 9.0工具包时选择了“兼容驱动程序”选项。这很可能是一个错误，因为它可能会安装随9.0工具包分发的旧驱动程序。 - jodag

感谢您的澄清 - 那确实是问题所在。 - infochip

Answer 1

可在 torch 稳定版发布页面 找到使用不同 CUDA 版本构建的预编译torch wheels。例如，您可以像这样安装使用 CUDA v11.1 构建的 torch v1.9.0：

pip install --upgrade torch==1.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

但并非所有的组合都可用。

- Mehraban