如何在Ubuntu上利用GPU加速PyTorch

在Ubuntu上利用GPU加速PyTorch主要分为以下几个步骤：

1. 安装合适的GPU驱动

• NVIDIA显卡：

  • 打开终端，输入以下命令安装NVIDIA驱动：

    sudo add-apt-repository ppa:graphics-drivers/ppa
    sudo apt-get update
    sudo apt-get install nvidia-xxx
    

    将xxx替换为适合你显卡型号的版本号。

• AMD显卡：

  • 输入以下命令安装AMD驱动：

    sudo apt-get install mesa-vulkan-drivers
    sudo apt-get install amdgpu-drivers
    

• Intel显卡：

  • 输入以下命令安装Intel显卡驱动：

    sudo apt-get install intel-driver
    

2. 安装CUDA Toolkit

• 检查CUDA支持的版本： 使用nvidia-smi命令查看支持的CUDA版本。

• 下载并安装CUDA： 根据支持的CUDA版本，从NVIDIA官网下载并安装CUDA Toolkit。例如，对于CUDA 12.9，可以使用以下命令：

  sudo apt update
  sudo apt install cuda
  

• 配置环境变量： 在.bashrc文件中添加以下内容：

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

  然后执行以下命令使更改生效：

  source ~/.bashrc
  

3. 安装cuDNN

• 下载与CUDA版本兼容的cuDNN库，并按照NVIDIA官方指南进行安装。

4. 安装PyTorch

• 使用conda安装（推荐）： 创建一个新的conda环境并激活：

  conda create -n pytorch-gpu python=3.11
  conda activate pytorch-gpu
  

  然后安装PyTorch GPU版本：

  conda install pytorch torchvision torchaudio cudatoolkit=12.9 -c pytorch -c nvidia
  

  将12.9替换为你安装的CUDA版本。

• 使用pip安装：

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

  将cu129替换为你安装的CUDA版本。

5. 验证GPU加速

• 打开Python，输入以下命令检查PyTorch是否成功使用GPU：

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

  如果torch.cuda.is_available()返回True，说明PyTorch已成功使用GPU。

6. 优化GPU性能（可选）

• 使用nvidia-smi监控GPU使用情况。
• 在运行GPU密集型应用时，可以通过调整参数来提高GPU占用率。

通过以上步骤，你就可以在Ubuntu上利用GPU加速PyTorch了。如果在安装过程中遇到问题，建议查阅PyTorch的官方文档或相关社区论坛以获取帮助。