Install PyTorch for ROCm#
Refer to this section for the recommended PyTorch via PIP installation method, as well as Docker-based installation.
AMD recommends the PIP install method to create a PyTorch environment when working with ROCm™ for machine learning development.
Install PyTorch via PIP
Enter the following command to unpack and begin set up.
sudo apt install python3-pip -y
Enter this command to update the pip wheel.
pip3 install --upgrade pip wheel
Enter this command to install Torch and Torchvision for ROCm AMD GPU support.
wget https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0.2/torch-2.1.2+rocm6.0-cp310-cp310-linux_x86_64.whl wget https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0.2/torchvision-0.16.1+rocm6.0-cp310-cp310-linux_x86_64.whl pip3 install --force-reinstall torch-2.1.2+rocm6.0-cp310-cp310-linux_x86_64.whl torchvision-0.16.1+rocm6.0-cp310-cp310-linux_x86_64.whl
This may take several minutes.
Important! AMD recommends proceeding with ROCm WHLs available at repo.radeon.com. The ROCm WHLs available at PyTorch.org are not tested extensively by AMD as the WHLs change regularly when the nightly builds are updated.
Using Docker provides portability, and access to a prebuilt Docker container that has been rigorously tested within AMD. Docker also cuts down compilation time, and should perform as expected without installation issues.
Prerequisites to install PyTorch using Docker
Docker for Ubuntu® must be installed.
To install Docker for Ubuntu, enter the following command:
sudo apt install docker.io
Use Docker image with pre-installed PyTorch
Follow these steps for installing using a Docker image.
Enter the following command to pull the public PyTorch Docker image.
sudo docker pull rocm/pytorch:rocm6.0.2_ubuntu22.04_py3.10_pytorch_2.1.2
Optional: You can also download a specific and supported configuration with different user-space ROCm versions, PyTorch versions, and supported operating systems.
Refer to hub.docker.com/r/rocm/pytorch to download the PyTorch Docker image.
Start a Docker container using the downloaded image.
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:rocm6.0.2_ubuntu22.04_py3.10_pytorch_2.1.2
This will automatically download the image if it does not exist on the host. You can also pass the
-vargument to mount any data directories from the host onto the container.
See PyTorch Installation for ROCm for more information.
Verify PyTorch installation#
Confirm if PyTorch is correctly installed.
Verify if Pytorch is installed and detecting the GPU compute device.
python3 -c 'import torch' 2> /dev/null && echo 'Success' || echo 'Failure'
Enter command to test if the GPU is available.
python3 -c 'import torch; print(torch.cuda.is_available())'
Enter command to display installed GPU device name.
python3 -c "import torch; print(f'device name :', torch.cuda.get_device_name(0))"
Expected result: Example: device name : Radeon RX 7900 XTX
device name : <Supported AMD GPU>
Enter command to display component information within the current PyTorch environment.
python3 -m torch.utils.collect_env
PyTorch version ROCM used to build PyTorch OS Is CUDA available GPU model and configuration HIP runtime version MIOpen runtime version
Environment set-up is complete, and the system is ready for use with PyTorch to work with machine learning models, and algorithms.