Ubuntu 20.04安装ROCm后PyTorch无法识别AMD GPU求助

核心问题 #

你当前安装的PyTorch是CUDA版本（版本号显示2.0.0+cu117），而AMD GPU需要适配ROCm的PyTorch版本，这是导致检测不到GPU的根本原因。

解决步骤 #

1. 卸载现有CUDA版PyTorch #

• 若通过pip安装：

  pip uninstall -y torch torchvision torchaudio
  

• 若通过conda安装：

  conda uninstall -y torch torchvision torchaudio
  

2. 配置ROCm环境变量 #

确保ROCm的路径被系统识别，将以下内容添加到~/.bashrc（或~/.zshrc，取决于你使用的shell）：

export PATH=/opt/rocm/bin:$PATH
export LD_LIBRARY_PATH=/opt/rocm/lib:$LD_LIBRARY_PATH

执行命令使配置立即生效：

source ~/.bashrc

3. 安装适配ROCm 5.0.2的PyTorch版本 #

使用pip安装对应版本的ROCm版PyTorch：

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

4. 验证安装结果 #

重启终端后，运行以下命令验证：

python -c "import torch; print(torch.__version__); print('GPU可用:', torch.cuda.is_available())"

正常情况下，版本号会显示2.0.0+rocm5.0，且GPU可用会返回True。

额外检查项 #

• 确保当前用户属于video组，执行命令添加后重启系统：

  sudo usermod -aG video $USER
  

• 检查ROCm组件是否正常运行：

  hipcc --version
  

  若能输出HIP编译器版本，说明ROCm环境正常。

内容的提问来源于stack exchange，提问作者David Lee