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.

Check Pytorch.org for latest PIP install instructions and availability. See Compatibility matrices for support information.

Install PyTorch via PIP

  1. Enter the following command to unpack and begin set up.

    sudo apt install python3-pip -y

  2. Enter this command to update the pip wheel.

    pip3 install --upgrade pip wheel

  3. 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.

Next, verify your PyTorch installation.

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.

  1. 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.

  2. 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 -v argument to mount any data directories from the host onto the container.

Next, verify your PyTorch installation.

See PyTorch Installation for ROCm for more information.

Verify PyTorch installation#

Confirm if PyTorch is correctly installed.

  1. Verify if Pytorch is installed and detecting the GPU compute device.

    python3 -c 'import torch' 2> /dev/null && echo 'Success' || echo 'Failure'

    Expected result:

    Success

  2. Enter command to test if the GPU is available.

    python3 -c 'import torch; print(torch.cuda.is_available())'

    Expected result:

    True

  3. Enter command to display installed GPU device name.

    python3 -c "import torch; print(f'device name [0]:', torch.cuda.get_device_name(0))"

    Expected result: Example: device name [0]: Radeon RX 7900 XTX

    device name [0]: <Supported AMD GPU>

  4. Enter command to display component information within the current PyTorch environment.

    python3 -m torch.utils.collect_env

    Expected result:

    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.