AMD显卡(RX6750gre)+ubuntu22.04基于docker部署ROCm环境(支持CUDA)

写在前面

个人感觉,到2025年3月8日,AMD显卡的深度学习支持已经做得相当给力了.
我尝试了三种在AMD显卡上部署支持CUDA的ROCm环境的方式:

1. 使用windows操作系统上的基于docker的ubuntu20.04容器环境
2. 使用windows操作系统上的wsl2(linux子系统)环境
3. 使用ubuntu22.04操作系统上的基于docker的ubuntu22.04.5容器(使用最新的官方镜像)
   前两种完全失败,因为windows系统完全不支持AMD显卡的ROCm驱动,具体表现为
   amdgpu kfd模块无法安装.
   所以,一定要使用直接安装在物理机上的ubuntu系统.如果你的电脑上只有windows系统,建议压缩任意磁盘空间后,安装ubuntu系统.

详细系统相关信息

• CPU: AMD Ryzen5 5600
• GPU:AMD Radeon rx 6750 gre 10GB
• OS: Ubuntu 22.04.5 LTS

操作过程

第一步,在你的ubuntu系统上安装amd的ROCm驱动

官方下载地址:
https://rocm.docs.amd.com/projects/install-on-linux/en/latest/install/install-methods/amdgpu-installer/amdgpu-installer-ubuntu.html

1. 应用于Ubuntu22.04(jammy),下载官方驱动下载软件’amdgpu’的命令如下:

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

2. 随后,使用’amdgpu’下载ROCm组件:

sudo amdgpu-install -y --usecase=rocm
sudo amdgpu-install -y --usecase=openclsdk
sudo amdgpu-install -y --usecase=hiplibsdk

参考了AMD显卡配置深度学习环境（ROCm-pytorch）https://zhuanlan.zhihu.com/p/654453422

3. 最后,使用以下命令确保你已经安装了ROCm:

rocm-smi

该命令用于查看显卡使用情况.

第二步,使用docker来pull官方镜像

下载docker步骤略.
官方pull和运行镜像的教程:
https://rocm.docs.amd.com/projects/install-on-linux/en/develop/install/3rd-party/pytorch-install.html

1. pull镜像

sudo docker pull rocm/pytorch:latest

2. 运行镜像

sudo docker run -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
--device=/dev/kfd --device=/dev/dri --group-add video \
--ipc=host --shm-size 8G rocm/pytorch:latest

3. 验证容器(镜像)是否已在运行

sudo docker ps

第三步,在容器中验证pytorch的cuda是否可用

在py_3.10的conda环境下执行以下代码(容器自带anaconda):

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

若输出为True,则环境配置成功.

总结

和以前相比,AMD显卡配置深度学习环境变得更加方便简单了.
但目前还不清楚AMD显卡运行深度学习时的GPU利用率是否较高.
此外,使用vscode+相关插件(Docker,Dev Containers,Remote - SSH, Remote Explorer等) 能更加方便的管理docker容器并进行开发.