如何用pip安装支持CUDA的PyTorch版本？

Question

**问题：如何用pip安装支持CUDA的PyTorch版本？** 在使用PyTorch进行深度学习开发时，常常需要安装支持CUDA的版本以利用GPU加速计算。然而，很多用户在使用`pip`安装时并不清楚如何确保安装的是CUDA支持版本。那么，如何通过`pip`正确安装与当前环境匹配、且支持CUDA的PyTorch版本？

Answer 1

1. 理解PyTorch与CUDA的关系

PyTorch 是一个广泛使用的深度学习框架，支持 GPU 加速的版本通常依赖于 CUDA（Compute Unified Device Architecture）。CUDA 是 NVIDIA 提供的并行计算平台和编程模型。

在安装 PyTorch 时，若希望使用 GPU 进行训练或推理，则必须确保所安装的 PyTorch 版本与当前系统的 CUDA 驱动兼容。否则，即使安装成功，也无法利用 GPU 加速。

2. 检查系统环境中的CUDA驱动版本

在安装之前，首先需要确认你的系统中是否已经安装了合适的 NVIDIA 显卡驱动，并且支持某个版本的 CUDA。

• Windows用户：打开命令提示符，运行 nvidia-smi 命令查看驱动版本和支持的CUDA最高版本。
• Linux用户：同样可以使用 nvidia-smi 查看相关信息。

# 示例输出
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03    Driver Version: 535.54.03    CUDA Version: 12.2     |
+-----------------------------------------------------------------------------+

从上述输出可以看到当前支持的 CUDA 版本为 12.2。

3. 确定所需的PyTorch版本与CUDA版本对应关系

PyTorch 官方提供了不同版本与其对应的 CUDA 支持情况。可以通过访问 PyTorch官方安装页面 获取最新的信息。

PyTorch版本	CUDA版本	安装命令示例
2.0.1	CUDA 11.8	pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 --extra-index-url https://download.pytorch.org/whl/cu118
1.13.1	CUDA 11.7	pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
1.12.1	CUDA 11.6	pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

4. 使用pip安装指定版本的PyTorch（含CUDA支持）

一旦确定了所需的 PyTorch 和 CUDA 对应关系，就可以使用 pip 安装包含 CUDA 支持的版本。

以 PyTorch 2.0.1 和 CUDA 11.8 为例：

pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 --extra-index-url https://download.pytorch.org/whl/cu118

该命令会从 PyTorch 的镜像源下载并安装支持 CUDA 11.8 的版本。

5. 验证安装是否成功及是否启用CUDA

安装完成后，建议验证 PyTorch 是否能够正确识别并使用 GPU。

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

如果输出如下内容，则说明安装成功并且 GPU 可用：

2.0.1+cu118
True
GeForce RTX 3090

6. 常见问题与解决方案

1. 问题1：安装后无法检测到GPU？
   可能原因：
   • NVIDIA驱动未正确安装或版本过低；
   • 安装的PyTorch版本与CUDA版本不匹配；
   • 系统未重启导致显卡驱动未生效。
2. 问题2：pip找不到指定版本？
   解决方法：
   • 检查版本拼写是否正确；
   • 尝试添加额外索引：--extra-index-url https://download.pytorch.org/whl/cu118；
   • 考虑使用 conda 安装更稳定的版本。

7. 自动化选择推荐版本的脚本流程图

为了简化用户的操作，可以编写一个简单的 shell 脚本或 Python 脚本来自动判断推荐的 PyTorch 安装命令。

#!/bin/bash
# 示例脚本片段
cuda_version=$(nvidia-smi | grep "CUDA Version" | awk '{print $NF}' | sed 's/\.$//')
echo "Detected CUDA version: $cuda_version"
# 根据版本映射输出推荐安装命令...

以下是一个 mermaid 流程图展示如何根据检测结果推荐安装命令：

graph TD
    A[开始] --> B{检测CUDA版本}
    B --> C{CUDA >= 11.8?}
    C -->|是| D[推荐PyTorch 2.0.1 +cu118]
    C -->|否| E{CUDA == 11.7?}
    E --> F[推荐PyTorch 1.13.1 +cu117]
    E --> G[其他情况]
    G --> H[提示用户手动查询]
    D --> I[结束]
    F --> I
    H --> I