Installing TensorFlow with ROCm Acceleration on Ubuntu 24.04

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, letting you use the power of AMD GPUs for machine learning tasks without manually setting up or managing dependencies. This guide downloads and runs a ROCm-supported TensorFlow container on a GPU-enabled server with an AMD GPU, then installs TensorFlow directly on the host using pip for the ROCm compute platform. By the end, you'll have TensorFlow running with GPU acceleration, both inside a container and installed natively on the host.

Prerequisites: an Ubuntu 24.04 server with an AMD GPU (ROCm-capable), and Docker with ROCm GPU support already installed and configured.

---

Use ROCm Supported TensorFlow Containers

Download and run a ROCm-supported TensorFlow container, and check GPU availability from inside it.

Before proceeding, make sure Docker with ROCm GPU support is installed and configured on your server.

1. Pull the ROCm supported container for TensorFlow:

$ docker pull rocm/tensorflow:latest

2. Run a temporary Docker container:

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

This runs a temporary container enabling access to GPU devices (/dev/kfd and /dev/dri) for ROCm-supported TensorFlow workloads.

3. Verify GPU availability from the temporary container:

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

The output should display all the devices along with their specifications.

4. Exit and destroy the temporary container:

$ exit

Install TensorFlow on Host using Pip

Install TensorFlow on the host machine using pip and check for GPU availability.

1. Fetch the ROCm version:

$ amd-smi version

2. Navigate to the 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, the compatible TensorFlow versions would be 2.14, 2.15, 2.16.

3. Install TensorFlow using a pre-built wheel:

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

Replace {TENSORFLOW_VERSION} with the version from the support matrix in the previous step, and {ROCM_VERSION} with only the first two numbers of your ROCm version.

For example, if your ROCm version is 6.2.x, your command should look like this:

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

4. Verify GPU availability:

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

The output should display all the devices along with their specifications.

Next Steps

• Install PyTorch with ROCm acceleration for additional deep learning workflows
• Set up a ROCm GPU operator if you're running workloads on a Kubernetes cluster
• Explore JAX with ROCm acceleration as an alternative ML framework
• Benchmark model training performance to validate GPU utilization under real workloads

For the full guide with additional tips, visit the original article on Vultr Docs.