如何在Ubuntu上安装PyTorch GPU版本 - 问答

在Ubuntu上安装PyTorch GPU版本需要以下几个步骤：

安装NVIDIA驱动：确保你的系统已经安装了与你的GPU兼容的NVIDIA驱动。你可以通过以下命令来检查是否已经安装了NVIDIA驱动：
```
nvidia-smi
```
如果显示了GPU信息，说明驱动已经安装。如果没有，你需要先安装驱动。
添加NVIDIA CUDA Toolkit： PyTorch GPU版本依赖于CUDA Toolkit。你可以从NVIDIA官网下载并安装适合你系统的CUDA Toolkit。以下是安装CUDA Toolkit的步骤：
- 访问NVIDIA CUDA Toolkit下载页面。
- 选择适合你Ubuntu版本的CUDA Toolkit进行下载。
- 下载完成后，按照安装向导进行安装。例如，对于CUDA 11.7，你可以使用以下命令：
```
sudo dpkg -i cuda-repo-ubuntu2004-11-7-local_11.7.0-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-7-local/7fa2af80.pub
sudo apt-get update
sudo apt-get install cuda
```
安装cuDNN： cuDNN是用于深度神经网络的GPU加速库。你需要从NVIDIA官网下载并安装与你的CUDA版本兼容的cuDNN库。以下是安装cuDNN的步骤：
- 访问NVIDIA cuDNN下载页面。
- 注册并登录NVIDIA开发者账号。
- 下载适合你CUDA版本的cuDNN库。
- 解压下载的文件并将文件复制到CUDA目录中。例如：
```
tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.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*
```
设置环境变量：为了确保PyTorch能够找到CUDA和cuDNN库，你需要设置一些环境变量。你可以在~/.bashrc文件中添加以下内容：
```
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
```
然后运行以下命令使环境变量生效：
```
source ~/.bashrc
```
安装PyTorch GPU版本：你可以使用pip来安装PyTorch GPU版本。首先，确保你已经安装了pip。然后，运行以下命令来安装PyTorch：
```
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
```
这里的cu117表示CUDA 11.7版本。如果你使用的是其他版本的CUDA，请相应地更改URL。

完成以上步骤后，你应该能够在Ubuntu上成功安装并运行PyTorch GPU版本。你可以通过以下命令来验证安装是否成功：

import torch
print(torch.cuda.is_available())

如果输出为True，说明PyTorch GPU版本已经成功安装并可以使用。

0 赞

0 踩