如何在Ubuntu上安装PyTorch CUDA

在Ubuntu上安装PyTorch CUDA的步骤如下：

1. 安装NVIDIA驱动

首先，确保你的系统已经安装了NVIDIA显卡驱动。你可以通过以下命令检查是否已安装驱动：

nvidia-smi

如果显示了显卡信息，说明驱动已安装。如果没有，你需要先安装驱动。

安装NVIDIA驱动

1. 添加NVIDIA PPA：

   sudo add-apt-repository ppa:graphics-drivers/ppa
   sudo apt update
   

2. 安装推荐的驱动版本（例如470）：

   sudo apt install nvidia-driver-470
   

3. 重启系统：

   sudo reboot
   

2. 安装CUDA Toolkit

访问NVIDIA CUDA Toolkit下载页面，选择适合你系统的版本并下载安装脚本。

使用命令行安装CUDA Toolkit

1. 下载CUDA Toolkit安装脚本（例如cuda_11.7.0_515.43.04_linux.run）。
2. 赋予安装脚本执行权限：

   chmod +x cuda_11.7.0_515.43.04_linux.run
   

3. 运行安装脚本：

   sudo ./cuda_11.7.0_515.43.04_linux.run
   

4. 按照提示完成安装。在安装过程中，确保选择安装CUDA Toolkit和驱动程序。

3. 配置环境变量

编辑~/.bashrc文件，添加以下内容：

export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

然后，使环境变量生效：

source ~/.bashrc

4. 安装cuDNN

访问NVIDIA cuDNN下载页面，下载与CUDA Toolkit版本匹配的cuDNN库。

使用命令行安装cuDNN

1. 解压下载的文件：

   tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.tgz
   

2. 将cuDNN文件复制到CUDA目录：

   sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
   sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
   sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
   

5. 验证安装

验证CUDA和cuDNN是否正确安装：

nvidia-smi
nvcc -V

如果显示了CUDA版本信息，说明CUDA安装成功。

6. 安装PyTorch with CUDA支持

使用pip安装PyTorch，确保选择与CUDA版本匹配的PyTorch版本。例如，如果你安装的是CUDA 11.7，可以使用以下命令：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117

7. 验证PyTorch with CUDA

编写一个简单的Python脚本来验证PyTorch是否能够使用CUDA：

import torch

print(torch.cuda.is_available())
print(torch.cuda.current_device())
print(torch.cuda.get_device_name(torch.cuda.current_device()))

如果输出显示CUDA可用，并且显示了GPU名称，说明PyTorch with CUDA安装成功。

通过以上步骤，你应该能够在Ubuntu上成功安装并配置PyTorch CUDA。