Ubuntu下PyTorch如何进行GPU加速

在Ubuntu系统下，使用PyTorch进行GPU加速需要以下几个步骤：

1. 检查GPU兼容性： 确保你的GPU支持CUDA。你可以通过NVIDIA的官方网站查看CUDA兼容性列表。

2. 安装NVIDIA驱动： 如果你的系统还没有安装NVIDIA驱动，你需要先安装它。可以使用以下命令来安装：

   sudo apt update
   sudo apt install nvidia-driver-<version>
   

   其中<version>是你想要安装的驱动版本号。安装完成后，重启系统。

3. 安装CUDA Toolkit： 访问NVIDIA的CUDA Toolkit下载页面，选择适合你系统的版本进行下载和安装。对于Ubuntu，通常使用以下命令：

   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.4.2/local_installers/cuda-repo-ubuntu2004-11-4-local_11.4.2-1_amd64.deb
   sudo dpkg -i cuda-repo-ubuntu2004-11-4-local_11.4.2-1_amd64.deb
   sudo apt-key add /var/cuda-repo-ubuntu2004-11-4-local/7fa2af80.pub
   sudo apt update
   sudo apt install cuda
   

4. 安装cuDNN： cuDNN是用于深度神经网络的GPU加速库。你需要注册NVIDIA开发者账户并下载适合你CUDA版本的cuDNN。下载完成后，解压并将文件复制到CUDA目录：

   tar -xzvf cudnn-11.4-linux-x64-v8.2.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*
   

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

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

   然后运行source ~/.bashrc使更改生效。

6. 安装PyTorch： 根据你的CUDA版本选择合适的PyTorch安装命令。你可以在PyTorch官方网站找到相应的安装命令。例如，如果你使用的是CUDA 11.4，可以使用以下命令：

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

7. 验证安装： 运行以下Python代码来验证PyTorch是否能够检测到GPU：

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

如果以上步骤都正确执行，你应该能够在Ubuntu系统下使用PyTorch进行GPU加速了。