Install TensorFlow for ROCm

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

ROCm version

10.1.0

10.0.0

7.14.1

Device family

AMD Instinct™

AMD Radeon™

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)

Radeon GPU

AMD Radeon RX 9070 XT (gfx1201)AMD Radeon RX 9070 GRE (gfx1201)AMD Radeon RX 9070 (gfx1201)AMD Radeon RX 9060 XT LP (gfx1200)AMD Radeon RX 9060 XT (gfx1200)AMD Radeon RX 9060 (gfx1200)AMD Radeon RX 9050 (gfx1200)AMD Radeon RX 9050 (4GB) (gfx1200)AMD Radeon RX 7900 XTX (gfx1100)AMD Radeon RX 7900 XT (gfx1100)AMD Radeon RX 7900 GRE (gfx1100)AMD Radeon RX 7800 XT (gfx1101)AMD Radeon RX 7700 XT (gfx1101)AMD Radeon RX 7700 (gfx1101)AMD Radeon RX 7600 (gfx1102)AMD Radeon PRO W7900 Dual Slot (gfx1100)AMD Radeon PRO W7900 (gfx1100)AMD Radeon PRO W7800 48GB (gfx1100)AMD Radeon PRO W7800 (gfx1100)AMD Radeon PRO W7700 (gfx1101)AMD Radeon PRO V710 (gfx1101)AMD Radeon AI PRO R9700S (gfx1201)AMD Radeon AI PRO R9700 (gfx1201)AMD Radeon AI PRO R9600D (gfx1201)AMD Radeon AI PRO R9600 (gfx1201)

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)

TensorFlow version

2.21

2.20

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 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. For more guidance on running ROCm workloads in Docker containers, see Run ROCm Docker containers.

• 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.1.0) for instructions.

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

Get started

1. Pull the ROCm TensorFlow 2.21 Docker image.

   Ubuntu 24.04

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

   Ubuntu 22.04

   docker pull rocm/tensorflow:rocm10.1-ubuntu22.04-py3.12-tf2.21
   

1. Pull the ROCm TensorFlow 2.20 Docker image.

   Ubuntu 24.04

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

   Ubuntu 22.04

   docker pull rocm/tensorflow:rocm10.1-ubuntu22.04-py3.12-tf2.20
   

2. Start the Docker container.

   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.1-ubuntu24.04-py3.12-tf2.21 \
      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.1-ubuntu22.04-py3.12-tf2.21 \
      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.1-ubuntu24.04-py3.12-tf2.20 \
      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.1-ubuntu22.04-py3.12-tf2.20 \
      bash
   

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

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

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

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

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

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

   python -m pip install --index-url https://stable.repo.amd.com/rocm/whl-next/ \
       "rocm[libraries,device-gfx1102]==10.1.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.1.0"
   

   python -m pip install --extra-index-url https://rocm.frameworks.amd.com/whl-multi-arch/ \
       "tensorflow-rocm==2.20.0+rocm10.1.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
   

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

6. 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 ImportErrors.

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