How to Install TensorFlow with ROCm Acceleration on Ubuntu 24.04

Introduction

TensorFlow is an open-source framework widely used for building and training machine learning models, including deep learning. ROCm is AMD's platform for GPU computing, designed to deliver high-performance computing and machine learning capabilities on AMD GPUs. A ROCm-enabled TensorFlow container is a pre-configured, portable environment that includes TensorFlow optimized for running efficiently on AMD GPUs. Using this container allows you to utilize the power of AMD GPUs for machine learning tasks without the complexity of manually setting up or managing dependencies. It streamlines the process of deploying TensorFlow applications on GPU hardware.

In this article, you are to download and run a ROCm supported TensorFlow container, and install TensorFlow using Pip for ROCm compute platform.

Use ROCm Supported TensorFlow Containers

In this section, you are to download and run a ROCm supported TensorFlow container. You are to also check the GPU availability from the container.

Before proceeding with the steps in this section, ensure you have downloaded Docker with ROCm GPU Support on Ubuntu 24.04.

1. Pull the ROCm supported container for TensorFlow.

   CONSOLE

   Copy

   $ docker pull rocm/tensorflow:latest

   Explain Code

   Powered by Vultr AgentBeta

2. Run a temporary Docker container.

   CONSOLE

   Copy

   $ docker run --rm -it --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --shm-size 8G rocm/tensorflow:latest

   Explain Code

   Powered by Vultr AgentBeta

   The above command runs a temporary container enabling access to GPU devices (/dev/kfd and /dev/dri) for ROCm-supported TensorFlow workloads. To learn more about ROCm devices.

3. Verify GPU availability from the temporary container.

   CONSOLE

   Copy

   $ rocm-smi
   $ python3 -c 'import tensorflow; print(tensorflow.config.list_physical_devices())'

   Explain Code

   Powered by Vultr AgentBeta

   The above command's output should display all the devices along with their specifications.

4. Exit and destory the temporary container.

   CONSOLE

   Copy

   $ exit

   Explain Code

   Powered by Vultr AgentBeta

Install TensorFlow on Host using Pip

In this section, you are to install TensorFlow on the host machine using Pip and check for GPU availability.

1. Fetch the ROCm version.

   CONSOLE

   Copy

   $ amd-smi version

   Explain Code

   Powered by Vultr AgentBeta

2. Navigate to official installation page and review the support matrix to find the compatible TensorFlow version to accompany ROCm. For example, if your ROCm version is 6.2.x then the compatible TensorFlow versions would be 2.14, 2.15, 2.16.

3. Install TensorFlow using a pre-built wheel.

   CONSOLE

   Copy

   $ python3 -m pip install tensorflow-rocm=={TENSORFLOW_VERSION} -f https://repo.radeon.com/rocm/manylinux/rocm-rel-{ROCM_VERSION}

   Explain Code

   Powered by Vultr AgentBeta

   Replace {TENSORFLOW_VERSION} with suitable version according to support matrix mentioned in the previous instruction. Additionally, replace {ROCM_VERSION} with only the first two numbers.

   For example, if your ROCm version is 6.2.x then your command should look like the below example command.

   python3 -m pip install tensorflow-rocm==2.14 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.2

4. Verify GPU availability.

   CONSOLE

   Copy

   $ python3 -c 'import tensorflow; print(tensorflow.config.list_physical_devices())'

   Explain Code

   Powered by Vultr AgentBeta

   The above command's output should display all the devices along with their specifications.

Conclusion

In this article, you downloaded and ran a ROCm supported TensorFlow container with access to ROCm supported devices. Furthermore, you also installed TensorFlow using Pip for ROCm compute platform.

More Resources

• Install Docker with ROCm GPU Support on Ubuntu 24.04
• Install PyTorch with ROCm Acceleration on Ubuntu 24.04.
• Install K3s with ROCm GPU Operator on Ubuntu 24.04.
• Install JAX with ROCm Acceleration on Ubuntu 24.04.