Ubuntu PyTorch GPU支持情况

在Ubuntu上配置PyTorch以支持GPU主要包括以下几个步骤：

1. 安装NVIDIA显卡驱动：

   • 使用ubuntu-drivers devices命令查看推荐的驱动版本，并使用sudo apt install nvidia-driver-<version>命令进行安装，最后重启系统。

2. 安装CUDA：

   • 根据PyTorch官网推荐的CUDA版本下载并安装CUDA。例如，当前最新版的PyTorch支持CUDA 11.6，因此需要安装CUDA 11.6。

3. 配置环境变量：

   • 编辑~/.bashrc文件，添加以下内容：

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

   • 执行source ~/.bashrc使配置生效。

4. 安装PyTorch GPU版本：

   • 使用conda创建并激活虚拟环境（推荐）：

     conda create -n dl-gpu python=3.8
     conda activate dl-gpu
     

   • 从PyTorch官网选择合适的PyTorch版本进行安装，例如：

     conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge
     

5. 验证安装：

   • 在Python交互模式中输入以下命令，如果返回True，则表示安装成功：

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

通过以上步骤，可以在Ubuntu上成功配置PyTorch以使用GPU加速功能。