TensorFlow on ROCm#
2024-12-17
6 min read time
TensorFlow is an open-source library for solving machine learning, deep learning, and AI problems. It can solve many problems across different sectors and industries, but primarily focuses on neural network training and inference. It is one of the most popular and in-demand frameworks and is very active in open-source contribution and development.
To install TensorFlow for ROCm, you have the following options:
Using a Docker image with TensorFlow pre-installed (recommended)
Note
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.
ROCm version |
TensorFlow version |
|---|---|
6.3.x |
2.15.1, 2.16.2, 2.17.0 |
6.2.x |
2.14.1, 2.15.1, 2.16.1 |
6.1.x |
2.13.1, 2.14.0, 2.15.0 |
6.0.x |
2.12.1, 2.13.1, 2.14.0 |
Using a Docker image with TensorFlow pre-installed#
To install ROCm on bare metal, follow ROCm installation overview. The recommended option to get a TensorFlow environment is through Docker.
Using Docker provides portability and access to a prebuilt Docker image that has been rigorously tested within AMD. This can also save compilation time and should perform as tested and mitigate potential installation issues. See Docker image support
Follow these steps:
Pull the latest public TensorFlow Docker image.
docker pull rocm/tensorflow:latest
Once you have pulled the image, run it by using the command below:
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
Docker image support#
AMD validates and publishes ready-made TensorFlow images with ROCm backends on Docker Hub. The following Docker image tags and associated inventories are validated for ROCm 6.3.
Using a wheels package#
To install TensorFlow using the wheels package, use the following command.
pip install --user tensorflow-rocm==[wheel-version] -f [repo] --upgrade
[wheel-version]is the TensorFlow version.[repo]ishttps://repo.radeon.com/rocm/manylinux/rocm-rel-X.Y/for versions 6.1 and later, whereX.Yindicates the ROCm version.
Note
Prior to ROCm 6.1, [wheel-version] followed the <TensorFlowVersion>.<ROCmVersion> format.
Testing the TensorFlow installation#
To test the installation of TensorFlow, run the container as specified in Installing TensorFlow. Ensure you have access to the Python shell in the Docker container.
python -c 'import tensorflow' 2> /dev/null && echo ‘Success’ || echo ‘Failure’
Running a basic TensorFlow example#
To quickly validate your TensorFlow environment, run a basic TensorFlow example.
The MNIST dataset is a collection of handwritten digits that may be used to train a Convolutional Neural Network (CNN) for handwriting recognition. This dataset is included with your TensorFlow installation.
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)
If successful, you should see the following output indicating the image classifier is now trained to around 98% accuracy on this dataset.
