Ubuntu上如何利用GPU加速PyTorch

在Ubuntu上利用GPU加速PyTorch，你需要完成以下几个步骤：

1. 安装NVIDIA驱动： 确保你的GPU支持CUDA，并且已经安装了相应的NVIDIA驱动。你可以通过以下命令来检查是否安装了正确的驱动：

   nvidia-smi
   

   如果显示了GPU的信息，说明驱动已经安装。

2. 安装CUDA Toolkit： 根据你的GPU型号和PyTorch版本，选择合适的CUDA Toolkit版本进行安装。可以从NVIDIA官网下载CUDA Toolkit：https://developer.nvidia.com/cuda-toolkit

   安装CUDA Toolkit的命令示例（以CUDA 11.7为例）：

   wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run
   sudo sh cuda_11.7.0_515.43.04_linux.run
   

   安装过程中，确保选择了安装CUDA Toolkit和驱动程序。

3. 安装cuDNN： cuDNN是用于深度神经网络的GPU加速库。你需要注册NVIDIA开发者账号并下载与CUDA Toolkit版本相匹配的cuDNN库。下载地址：https://developer.nvidia.com/cudnn

   下载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*
   

4. 配置环境变量： 为了能够在任何地方使用CUDA命令和库，你需要将CUDA的路径添加到环境变量中。编辑~/.bashrc文件，添加以下行：

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

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

5. 安装PyTorch： 安装支持CUDA的PyTorch版本。你可以使用pip来安装：

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

   上面的命令会安装与CUDA 11.7兼容的PyTorch版本。

6. 验证安装： 安装完成后，你可以通过运行一个简单的PyTorch程序来验证GPU是否可用：

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

   如果输出为True，则表示PyTorch已经成功配置并可以使用GPU加速。

完成以上步骤后，你就可以在Ubuntu上利用GPU加速PyTorch了。记得在运行PyTorch程序时，使用.to('cuda')或者.cuda()方法将模型和数据移动到GPU上。