Install TensorFlow for ROCm

Install TensorFlow for ROCm#

TensorFlow is an open-source library for solving machine-learning, deep-learning, and artificial-intelligence problems. It can be used to solve many problems across different sectors and industries but primarily focuses on training and inference in neural networks. It is one of the most popular and in-demand frameworks and is very active in open source contribution and development.

As of ROCm 6.1, tensorflow-rocm packages are found at https://repo.radeon.com/rocm/manylinux. Prior to ROCm 6.1, packages were found at https://pypi.org/project/tensorflow-rocm.
Refer to the following version support matrix:

ROCm version

TensorFlow version

6.1.x

2.13.1, 2.14.0, 2.15.0

6.0.x

2.12, 2.13.1, 2.14.0

Pre-requisites#

  • Radeon software for Linux (with ROCm) must be installed.

  • MIGraphX must be installed for TensorFlow to build the correct mig execution provider.

Select installation method#

AMD recommends the PIP install method to create a TensorFlow environment when working with ROCm for machine learning development.

Note
The latest version of Python module numpy v2.0 is incompatible with the TensorFlow wheels for this version. Downgrade to an older version is required.
Example: pip3 install numpy==1.26.4

To install TensorFlow,

Download and install the TensorFlow wheel.

pip3 uninstall tensorflow-rocm numpy
pip3 install https://repo.radeon.com/rocm/manylinux/rocm-rel-6.1.3/tensorflow_rocm-2.15.1-cp310-cp310-manylinux_2_28_x86_64.whl numpy==1.26.4

Next, verify your TensorFlow installation.

To install ROCm on bare metal, follow ROCm installation options.

Using Docker provides portability and access to a prebuilt Docker container that has been rigorously tested within AMD. This might also save compilation time and should perform as tested without facing potential installation issues.

To use Docker image with pre-installed TensorFlow

  1. Enter the following command to pull the latest public TensorFlow Docker image.

    docker pull rocm/tensorflow:latest

  2. Run the downloaded image.

    docker run -it --network=host --device=/dev/kfd --device=/dev/dri \
--ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined rocm/tensorflow:latest

    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 TensorFlow installation.

Note See Tensorflow Installation for ROCm for more information.

Verify TensorFlow installation#

To test the TensorFlow installation, run the container image as specified in the previous section Installing TensorFlow. Ensure you have access to the Python shell in the Docker container.

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

Next, run basic TensorFlow example.

Run basic TensorFlow example#

The TensorFlow examples repository provides basic examples that exercise the framework’s functionality.

The MNIST database is a collection of handwritten digits that may be used to train a Convolutional Neural Network for handwriting recognition.

This dataset is included with your TensorFlow installation.

  1. Run the following sample code to load the MNIST dataset, then train and evaluate it.

    import tensorflow as tf
print("TensorFlow version:", tf.__version__)
mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])
predictions = model(x_train[:1]).numpy()
tf.nn.softmax(predictions).numpy()
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
loss_fn(y_train[:1], predictions).numpy()
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test,  y_test, verbose=2)

  2. If successful, you should see the following output indicating the image classifier is now trained to around 98% accuracy on this dataset.

    TensorFlow test output

Environment set-up is complete, and the system is ready for use with TensorFlow to work with machine learning models, and algorithms.

Contents