Ubuntu中PyTorch GPU加速设置

在Ubuntu系统中配置PyTorch以使用GPU加速涉及几个关键步骤。以下是详细的指南：

准备工作

1. 安装NVIDIA GPU驱动程序：

   • 使用命令 lspci | grep -i nvidia 查看显卡型号。
   • 根据显卡型号从NVIDIA官方网站下载并安装相应的驱动程序。

2. 安装CUDA Toolkit：

   • 访问NVIDIA CUDA Toolkit下载页面，选择适合Ubuntu版本的CUDA Toolkit进行下载。
   • 安装CUDA Toolkit：

     sudo apt update
     sudo apt install -y nvidia-cuda-toolkit
     

   • 验证CUDA安装：

     nvcc --version
     

3. 安装cuDNN库：

   • 注册并登录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*
     

安装PyTorch

1. 安装依赖项：

   sudo apt update
   sudo apt install -y python3-pip python3-dev
   

2. 安装PyTorch GPU版本：

   • 访问PyTorch官网，获取适合系统的安装命令。以下是一个示例命令，用于安装PyTorch GPU版本：

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

   • 确保使用与CUDA版本兼容的PyTorch版本。

验证安装

1. 检查CUDA版本：

   nvcc --version
   

2. 设置环境变量：

   • 为了使PyTorch能够使用GPU加速，需要设置一些环境变量：

     export PATH=/usr/local/cuda/bin:$PATH
     export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
     

   • 使环境变量生效：

     source ~/.bashrc
     

3. 验证GPU加速：

   import torch
   print(torch.__version__)
   print("CUDA is available. You can use GPU.") if torch.cuda.is_available() else print("CUDA is not available. Using CPU.")
   

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

注意事项

• 上述步骤可能会随着时间的推移而变化，因此建议查看PyTorch官方网站上的最新安装指南。

通过以上步骤，您应该能够在Ubuntu系统上成功配置PyTorch以使用GPU加速。现在，您可以使用PyTorch进行深度学习模型的训练，享受GPU加速带来的速度提升。