Ubuntu22.04安装AMD显卡驱动和pytorch

安装pip和vim

sudo apt-get -y install python3-pip vim

配置pip源

mkdir ~/.pip && vim ~/.pip/pip.conf

pip.conf

[global]
index-url = https://mirrors.cloud.tencent.com/pypi/simple
[install]
trusted-host = mirrors.cloud.tencent.com

安装显卡驱动 

$ cd ~
$ uname -srmv # rocm5.4.2需要Ubuntu内核5.15+

# Ubuntu v22.04
$ wget https://repo.radeon.com/amdgpu-install/5.4.2/ubuntu/jammy/amdgpu-install_5.4.50402-1_all.deb 

# Ubuntu v20.04
# wget https://repo.radeon.com/amdgpu-install/5.4.2/ubuntu/focal/amdgpu-install_5.4.50402-1_all.deb

$ sudo apt-get install ./amdgpu-install_5.4.50402-1_all.deb
$ sudo apt-get update
# docker安装需要修改参数，详情参考rocm安装文档
$ amdgpu-install -y --usecase=graphics,rocm

# 设置运行权限
$ ls -l /dev/dri/render*
$ sudo usermod -a -G render $LOGNAME
$ sudo usermod -a -G video $LOGNAME

# 重启
$ reboot

# 验证
$ dkms status
$ /opt/rocm-5.4.2/bin/rocminfo
# 或者
$ /opt/rocm-5.4.2/opencl/bin/clinfo
$ rocm-smi # 显示GPU活动信息

安装pytorch和测试

$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.4.2
$ python3
>> import torch
>> torch.cuda.is_available() # true为安装成功

参考资料

pytorch官网
AMD显卡驱动下载

AMD显卡驱动安装文档

ROCm安装文档