如何在Linux上安装PyTorch CUDA

在Linux上安装PyTorch CUDA需要遵循以下步骤：

1. 安装CUDA Toolkit

首先，你需要安装NVIDIA的CUDA Toolkit。以下是安装步骤：

方法一：使用官方安装脚本

下载CUDA Toolkit安装脚本：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-8-local/7fa2af80.pub
sudo apt update
sudo apt install cuda

设置环境变量：编辑~/.bashrc文件，添加以下行：

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

然后运行：

source ~/.bashrc

方法二：使用包管理器

如果你使用的是Ubuntu 20.04，可以直接使用包管理器安装：

sudo apt update
sudo apt install nvidia-cuda-toolkit

2. 安装cuDNN

cuDNN是用于深度神经网络的GPU加速库。你需要从NVIDIA官网下载并安装cuDNN。

注册NVIDIA开发者账号并登录。
下载cuDNN库（需要选择与CUDA版本匹配的cuDNN版本）。

解压并安装：

tar -xzvf cudnn-11.8-linux-x64-v8.5.0.96.tgz
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*

3. 安装PyTorch with CUDA支持

你可以使用pip来安装带有CUDA支持的PyTorch。首先，确保你已经安装了pip和虚拟环境（可选但推荐）。

创建虚拟环境（可选）

python3 -m venv pytorch_env
source pytorch_env/bin/activate

安装PyTorch

访问PyTorch官网，选择适合你系统的安装命令。例如，对于CUDA 11.8，可以使用以下命令：

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

4. 验证安装

验证PyTorch是否正确安装并检测到CUDA：

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

如果输出显示CUDA可用并且设备名称正确，说明安装成功。

注意事项

确保你的GPU驱动是最新的。
CUDA和cuDNN的版本必须与PyTorch兼容。
如果遇到问题，可以参考PyTorch官方文档或社区论坛寻求帮助。

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