Windows WSL环境下安装 pytorch +ROCM 支持AMD显卡

官方文档：Install PyTorch for ROCm — Use ROCm on Radeon GPUs

一、操作系统及驱动

windows 下安装WSL 环境( windows subsystem for Linux), 安装ubuntu 22.04环境。

安装 rocm 软件包：

sudo apt update
wget https://repo.radeon.com/amdgpu-install/6.1.3/ubuntu/jammy/amdgpu-install_6.1.60103-1_all.deb
sudo apt install ./amdgpu-install_6.1.60103-1_all.deb





sudo amdgpu-install --list-usecase


amdgpu-install -y --usecase=wsl,rocm --no-dkms

运行rocminfo 确认

rocminfo

Expected result:

[...]
*******
Agent 2
*******
  Name:                    gfx1100
  Marketing Name:          Radeon RX 7900 XTX
  Vendor Name:             AMD
  [...]
[...]

 

二、安装anaconda

安装好后，创建pytorch环境，注意，只支持python 3.10

conda create -n pytorch python=3.10

三、安装pytorch+rocm

pip3 install --upgrade pip wheel  #更新pip /wheel工具

安装pytorch +rocm 

wget https://repo.radeon.com/rocm/manylinux/rocm-rel-6.1.3/torch-2.1.2%2Brocm6.1.3-cp310-cp310-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-6.1.3/torchvision-0.16.1%2Brocm6.1.3-cp310-cp310-linux_x86_64.whl
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-6.1.3/pytorch_triton_rocm-2.1.0%2Brocm6.1.3.4d510c3a44-cp310-cp310-linux_x86_64.whl
pip3 uninstall torch torchvision pytorch-triton-rocm numpy
pip3 install torch-2.1.2+rocm6.1.3-cp310-cp310-linux_x86_64.whl torchvision-0.16.1+rocm6.1.3-cp310-cp310-linux_x86_64.whl pytorch_triton_rocm-2.1.0+rocm6.1.3.4d510c3a44-cp310-cp310-linux_x86_64.whl

 替换运行库：

location=`pip show torch | grep Location | awk -F ": " '{print $2}'`
cd ${location}/torch/lib/
rm libhsa-runtime64.so*
cp /opt/rocm/lib/libhsa-runtime64.so.1.2 libhsa-runtime64.so

 anaconda中会遇到gcc不兼容，更新下：

conda install -c conda-forge gcc=12.1.0