AMD显卡(RX6750gre)+ubuntu22.04基于docker部署ROCm环境(支持CUDA)
写在前面
个人感觉,到2025年3月8日,AMD显卡的深度学习支持已经做得相当给力了.
我尝试了三种在AMD显卡上部署支持CUDA的ROCm环境的方式:
- 使用windows操作系统上的基于docker的ubuntu20.04容器环境
- 使用windows操作系统上的wsl2(linux子系统)环境
- 使用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
- 应用于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
- 随后,使用’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
- 最后,使用以下命令确保你已经安装了ROCm:
rocm-smi
该命令用于查看显卡使用情况.
第二步,使用docker来pull官方镜像
下载docker步骤略.
官方pull和运行镜像的教程:
https://rocm.docs.amd.com/projects/install-on-linux/en/develop/install/3rd-party/pytorch-install.html
- pull镜像
sudo docker pull rocm/pytorch:latest
- 运行镜像
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
- 验证容器(镜像)是否已在运行
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容器并进行开发.
