Install TensorFlow for ROCm

This page guides you through installing TensorFlow with ROCm support on AMD Instinct GPUs running Linux. It applies to supported AMD GPUs and platforms.

Device family

AMD Instinct™

Instinct GPU

AMD Instinct MI355X (gfx950)AMD Instinct MI350X (gfx950)AMD Instinct MI350P (gfx950)AMD Instinct MI325X (gfx942)AMD Instinct MI300X (gfx942)AMD Instinct MI300A (gfx942)AMD Instinct MI250X (gfx90a)AMD Instinct MI250 (gfx90a)AMD Instinct MI210 (gfx90a)

ROCm version

10.0.0

7.14.1

TensorFlow version

2.21

2.20

2.19.1

Installation method

Docker

pip

Prerequisites

• Ensure your system has the AMD GPU Driver (amdgpu) installed. See the ROCm compatibility matrix for driver support information. For installation instructions, see the AMD GPU Driver documentation.

• Ensure your system has the AMD GPU Driver (amdgpu) installed. See the ROCm compatibility matrix for driver support information. For installation instructions, see the AMD GPU Driver documentation.

• Ensure the host system has Docker Engine installed.

• Ensure your system has a supported Python version installed and accessible: 3.12

• Complete the ROCm Core SDK installation prerequisites for installing via pip. See Prerequisites (Install ROCm 10.0.0) for instructions.

• Complete the ROCm Core SDK installation prerequisites for installing via pip. See Prerequisites (Install ROCm 7.14.1) for instructions.

• Complete the ROCm Core SDK installation prerequisites for installing via pip. See Prerequisites (Install ROCm 7.14.0) for instructions.

Get started

1. Pull the ROCm TensorFlow 2.21 Docker image.

   docker pull rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.21
   

1. Pull the ROCm TensorFlow 2.20 Docker image.

   docker pull rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.20
   

1. Pull the ROCm TensorFlow 2.19.1 Docker image.

   Ubuntu 24.04

   docker pull rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.19.1
   

   Ubuntu 22.04

   docker pull rocm/tensorflow:rocm10.0-ubuntu22.04-py3.12-tf2.19.1
   

2. Start the Docker container.

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.21 \
      bash
   

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.20 \
      bash
   

   Ubuntu 24.04

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm10.0-ubuntu24.04-py3.12-tf2.19.1 \
      bash
   

   Ubuntu 22.04

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm10.0-ubuntu22.04-py3.12-tf2.19.1 \
      bash
   

Get started

1. Pull the ROCm TensorFlow 2.21 Docker image.

   docker pull rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.21
   

1. Pull the ROCm TensorFlow 2.20 Docker image.

   docker pull rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.20
   

1. Pull the ROCm TensorFlow 2.19.1 Docker image.

   Ubuntu 24.04

   docker pull rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.19.1
   

   Ubuntu 22.04

   docker pull rocm/tensorflow:rocm7.14.1-ubuntu22.04-py3.12-tf2.19.1
   

2. Start the Docker container.

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.21 \
      bash
   

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.20 \
      bash
   

   Ubuntu 24.04

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm7.14.1-ubuntu24.04-py3.12-tf2.19.1 \
      bash
   

   Ubuntu 22.04

   docker run -it --rm \
      --device /dev/kfd \
      --device /dev/dri \
      --network=host \
      --ipc=host \
      --group-add=video \
      --cap-add=SYS_PTRACE \
      --security-opt seccomp=unconfined \
      rocm/tensorflow:rocm7.14.1-ubuntu22.04-py3.12-tf2.19.1 \
      bash
   

Install TensorFlow using pip

For prerequisite steps and post-installation recommendations, see the ROCm installation instructions.

1. Set up your Python virtual environment.

   python3.12 -m venv .venv
   

2. Activate your Python virtual environment.

   source .venv/bin/activate
   

3. If you don’t have an existing ROCm installation, install ROCm using the following command; see the ROCm 10.0.0 installation instructions for complete guidance. Otherwise, proceed to installing TensorFlow packages.

   python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \
       "rocm[libraries,device-gfx950]==10.0.0"
   

   python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \
       "rocm[libraries,device-gfx942]==10.0.0"
   

   python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \
       "rocm[libraries,device-gfx90a]==10.0.0"
   

4. Install the ROCm-enabled TensorFlow libraries.

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.21.0+rocm10.0.0"
   

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.20.0+rocm10.0.0"
   

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.19.1+rocm10.0.0"
   

5. Update LD_LIBRARY_PATH as a workaround so TensorFlow can discover ROCm libraries and system dependencies.

   export LD_LIBRARY_PATH=$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib/rocm_sysdeps/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_libraries/lib:$LD_LIBRARY_PATH
   

5. Verify your TensorFlow installation.

   python -c "import tensorflow as tf; print('TensorFlow version: ', tf.__version__); print('GPUs:', tf.config.list_physical_devices('GPU'))"
   

Install TensorFlow using pip

For prerequisite steps and post-installation recommendations, see the ROCm installation instructions.

1. Set up your Python virtual environment.

   python3.12 -m venv .venv
   

2. Activate your Python virtual environment.

   source .venv/bin/activate
   

3. If you don’t have an existing ROCm installation, install ROCm using the following command; see the ROCm 7.14.1 installation instructions for complete guidance. Otherwise, proceed to installing TensorFlow packages.

   python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
       "rocm[libraries,device-gfx950]==7.14.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
       "rocm[libraries,device-gfx942]==7.14.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \
       "rocm[libraries,device-gfx90a]==7.14.1"
   

4. Install the ROCm-enabled TensorFlow libraries.

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.21.0+rocm7.14.1"
   

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.20.0+rocm7.14.1"
   

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.19.1+rocm7.14.1"
   

5. Update LD_LIBRARY_PATH as a workaround so TensorFlow can discover ROCm libraries and system dependencies.

   export LD_LIBRARY_PATH=$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib/rocm_sysdeps/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_libraries/lib:$LD_LIBRARY_PATH
   

6. Verify your TensorFlow installation.

   python -c "import tensorflow as tf; print('TensorFlow version: ', tf.__version__); print('GPUs:', tf.config.list_physical_devices('GPU'))"
   

Known issues

After installing rocm_tensorflow using pip, attempting to run TensorFlow can result in multiple ImportError.

ImportError: libhipsparse.so.4
...
ImportError: librocm_sysdeps_asm.so.1
...

As a workaround, update LD_LIBRARY_PATH to link to the required ROCm libraries and system dependencies in your installation path:

export LD_LIBRARY_PATH=$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_core/lib/rocm_sysdeps/lib:$VIRTUAL_ENV/lib/python3.12/site-packages/_rocm_sdk_libraries/lib:$LD_LIBRARY_PATH