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.

Important!
ROCm on Radeon scripts uses Keras 2, but newest wheels uses Keras 3 by default.

A manual install of the tf-keras package is required to enable Keras 2 on TensorFlow.

Use the following command to install the ROCm compatible TensorFlow wheel.

`pip install tf-keras --no-deps`

Additional information 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.

PIP installation#

Use 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

Library Compatibility
When installing TensorFlow, it is essential to ensure that any additional TensorFlow-related libraries, or dependencies are compatible with the version of TensorFlow that you are using.

TensorFlow libraries (such as tensorflow-hub, etc.) often have specific version requirements that depend on the main TensorFlow version.

Be aware that installing or upgrading TensorFlow libraries may inadvertently replace the ROCm-supported TensorFlow package with a non-ROCm supported version.

To avoid this, ensure that you are explicitly specifying the ROCm-compatible version during installation.

To install TensorFlow,

Download and install the TensorFlow wheel.

pip3 uninstall tensorflow-rocm
pip3 install https://repo.radeon.com/rocm/manylinux/rocm-rel-6.2.3/tensorflow_rocm-2.16.2-cp310-cp310-manylinux_2_28_x86_64.whl

Next, verify your TensorFlow installation.

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