Ubuntu上如何安装PyTorch的CUDA版本

Ubuntu上安装PyTorch CUDA版本的步骤

1. 安装前准备

在安装PyTorch CUDA版本前，需完成以下基础准备：

• 更新系统包：运行sudo apt update && sudo apt upgrade -y，确保系统软件包为最新版本。
• 安装Python及pip：Ubuntu 22.04及以上默认包含Python 3.10+，通过sudo apt install python3 python3-pip python3-venv -y安装pip（若未安装）。
• 创建虚拟环境（可选但推荐）：避免依赖冲突，运行python3 -m venv pytorch_env创建环境，再通过source pytorch_env/bin/activate激活。

2. 安装NVIDIA驱动（关键前提）

PyTorch的CUDA版本需依赖系统中的NVIDIA驱动。需确保：

• 显卡支持CUDA：确认显卡型号在NVIDIA CUDA兼容列表中。
• 安装/更新驱动：运行sudo ubuntu-drivers autoinstall自动安装适合显卡的最新驱动；或通过sudo apt install nvidia-driver-XX（XX为驱动版本号，如535）手动安装。安装后重启系统，通过nvidia-smi命令验证驱动是否正常（显示显卡型号及CUDA版本）。

3. 通过pip安装PyTorch CUDA版本

pip是PyTorch推荐的安装方式，步骤如下：

• 选择CUDA版本：根据显卡驱动支持的CUDA版本（通过nvidia-smi查看），选择对应的PyTorch包。例如：
  • 若驱动支持CUDA 12.1，安装命令为：
    pip3 install torch==2.1.0+cu121 torchvision==0.16.0+cu121 torchaudio==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121
  • 若驱动支持CUDA 11.8，安装命令为：
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118（默认安装最新兼容版本）。
• 验证安装：在Python中运行以下代码，检查CUDA是否可用：

  import torch
  print("PyTorch版本:", torch.__version__)
  print("CUDA是否可用:", torch.cuda.is_available())  # 应返回True
  print("当前设备:", torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
  

4. 通过conda安装PyTorch CUDA版本（替代方案）

若使用conda管理环境，步骤如下：

• 安装conda：从Miniconda官网下载并安装Miniconda（轻量级版）。
• 创建conda环境：运行conda create -n pytorch_env python=3.9创建名为pytorch_env的环境（Python版本建议3.8及以上），再通过conda activate pytorch_env激活。
• 安装PyTorch：根据CUDA版本选择命令，例如：
  • CUDA 11.8：conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
  • CUDA 12.1：conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
• 验证安装：同pip安装后的验证步骤。

注意事项

• 版本兼容性：PyTorch版本与CUDA版本需严格匹配（如PyTorch 2.1.0支持CUDA 12.1，PyTorch 2.6.0支持CUDA 12.6），具体版本对应关系可参考PyTorch官网。
• 驱动更新：若nvidia-smi显示的CUDA版本低于PyTorch要求的版本，需先更新驱动（驱动版本需≥PyTorch要求的CUDA版本）。
• 网络问题：若安装速度慢，可使用国内镜像源（如清华源），例如pip命令添加--index-url https://pypi.tuna.tsinghua.edu.cn/simple。