Ubuntu上如何安装PyTorch的GPU版本

在Ubuntu上安装PyTorch GPU版，可按以下步骤进行：

1. 安装NVIDIA驱动：用nvidia-smi命令查看显卡信息，确定驱动版本。若未安装或需升级，可通过sudo ubuntu-drivers autoinstall安装推荐驱动。
2. 安装CUDA Toolkit：根据GPU型号和驱动版本，从NVIDIA官网下载对应的CUDA Toolkit安装包，如wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.90.07_linux.run，然后执行安装命令sudo sh cuda_12.4.0_550.90.07_linux.run，安装过程中可勾选安装CUDA Toolkit和驱动程序。
3. 安装cuDNN库：注册NVIDIA开发者账户，从cuDNN下载页面获取适合CUDA版本的cuDNN，下载后解压，将文件复制到CUDA目录，如tar -xzvf cudnn-12.4-linux-x64-v8.9.2.26.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*设置权限。
4. 设置环境变量：在~/.bashrc文件中添加export PATH=/usr/local/cuda-12.4/bin:$PATH和export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64:$LD_LIBRARY_PATH，然后运行source ~/.bashrc使设置生效。
5. 安装PyTorch GPU版：通过pip安装，以CUDA 12.4为例，命令为pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu124。

安装完成后，可通过运行import torch，再执行print(torch.__version__)和print(torch.cuda.is_available())来验证安装，若输出显示CUDA版本号且torch.cuda.is_available()返回True，则安装成功。