rocm/pytorch
This repo hosts the docker images for PyTorch releases with ROCm backend support.
(1) Supported GPUs:
The docker images hosted in this registry will run on gfx906 (Vega20-type GPU - MI50, MI60), gfx908 (MI100) and gfx90a (MI210/MI250/MI250x).
Starting from ROCm5.0, the "latest" and "*_staging" tag docker images also contain full Navi21 (gfx1030) support for AMD Radeon™ Pro W6800 and AMD Radeon™ Pro V620.
Starting from ROCm5.7, the "latest" tag docker images also contain Navi31 (gfx1100) support for AMD Radeon RX 7900 XTX and AMD Radeon PRO W7900.
(2) Install Docker Environment
To configure docker environment for ROCm, please refer to the following installation guide: https://github.com/RadeonOpenCompute/ROCm-docker/blob/master/quick-start.md
(3) Run the rocm/pytorch Docker Image:
To facilitate commands to pull and run the latest PyTorch docker container, add the following alias to your .profile or .bashrc:
alias drun='sudo docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G -v $HOME/dockerx:/dockerx -w /dockerx'
For example, to get the latest PyTorch on ROCm, run the following command:
drun rocm/pytorch
or
drun rocm/pytorch:latest
docker pull rocm/pytorch