PyTorch在Ubuntu上的GPU加速方法

在Ubuntu上使用PyTorch进行GPU加速，主要依赖于CUDA工具包。以下是详细的步骤：

安装CUDA

1. 检查CUDA兼容性：首先，确保你的Ubuntu系统支持CUDA。PyTorch支持的CUDA最低版本为18.04，因此建议使用Ubuntu 18.04或更高版本。
2. 安装CUDA Toolkit：可以从NVIDIA官网下载并安装适合你系统的CUDA Toolkit。安装命令如下：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget http://developer.download.nvidia.com/compute/cuda/11.0.3/local_installers/cuda-repo-ubuntu1804-11-0-local_11.0.3-450.51.06-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu1804-11-0-local_11.0.3-450.51.06-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu1804-11-0-local/7fa2af80.pub
sudo apt-get update
sudo apt-get -y install cuda[11.0-450.51.06-1~ubuntu18.04]

安装cuDNN

1. 下载cuDNN：访问NVIDIA cuDNN官方网站，下载与你的CUDA版本兼容的cuDNN库。
2. 安装cuDNN：解压下载的cuDNN文件，并将其复制到CUDA的安装目录。例如：

tar -xzvf cudnn-11.0-linux-x64-v8.0.4.30.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

1. 创建虚拟环境（可选但推荐）：使用conda创建一个新的虚拟环境，以便在其中安装和管理PyTorch版本。

conda create -n pytorch_env python=3.8
conda activate pytorch_env

2. 安装PyTorch：根据你的CUDA版本，从PyTorch官网获取相应的安装命令。例如，对于CUDA 11.0，可以使用以下命令：

conda install pytorch torchvision torchaudio cudatoolkit=11.0 -c pytorch

或者，如果你更喜欢使用pip，可以安装CPU版本，然后通过指定CUDA路径来使用GPU：

pip install torch torchvision torchaudio
export CUDA_HOME=/usr/local/cuda
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

验证安装

安装完成后，可以通过以下代码验证PyTorch是否正确安装并能够使用GPU：

import torch
print(torch.__version__)
print(torch.cuda.is_available())  # 应该返回True，如果GPU可用