TensorFlow Installation for ROCm#
9 min read time
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.
ROCm 5.6 and 5.7 deviates from the standard practice of supporting the last three TensorFlow versions. This is due to incompatibilities between earlier TensorFlow versions and changes introduced in the ROCm 5.6 compiler. Refer to the following version support matrix:
Last three versions at ROCm release.
The following sections contain options for installing TensorFlow.
Option 1: Install TensorFlow Using Docker Image#
To install ROCm on bare metal, follow the section Installation (Linux). The recommended option to get a TensorFlow environment is through Docker.
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. 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
Option 2: Install TensorFlow Using Wheels Package#
To install TensorFlow using the wheels package, follow these steps:
Check the Python version.
The Python version is less than 3.7
The Python version is more than 3.7
Skip this step and go to Step 3.
The supported Python versions are:
sudo apt-get install python3.7 # or python3.8 or python 3.9 or python 3.10
Set up multiple Python versions using update-alternatives.
update-alternatives --query python3 sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python[version] [priority]
Follow the instruction in Step 2 for incompatible Python versions.
sudo update-alternatives --config python3
Follow the screen prompts, and select the Python version installed in Step 2.
Install or upgrade PIP.
sudo apt install python3-pip
To install PIP, use the following:
/usr/bin/python[version] -m pip install --upgrade pip
Upgrade PIP for Python version installed in step 2:
sudo pip3 install --upgrade pip
Install TensorFlow for the Python version as indicated in Step 2.
/usr/bin/python[version] -m pip install --user tensorflow-rocm==[wheel-version] --upgrade
For a valid wheel version for a ROCm release, refer to the instruction below:
sudo apt install rocm-libs rccl
protobufto 3.19 or lower.
/usr/bin/python3.7 -m pip install protobuf=3.19.0 sudo pip3 install tensorflow
Set the environment variable
export PYTHONPATH="./.local/lib/python[version]/site-packages:$PYTHONPATH" #Use same python version as in step 2
sudo apt install rocm-libs rccl
python3 -c 'import tensorflow' 2> /dev/null && echo 'Success' || echo 'Failure'
For details on
tensorflow-rocmwheels and ROCm version compatibility, see: ROCmSoftwarePlatform/tensorflow-upstream
Test the TensorFlow Installation#
To test the installation of TensorFlow, 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’
Run a 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.
Follow these steps:
Clone the TensorFlow example repository.
cd ~ git clone https://github.com/tensorflow/models.git
Install the dependencies of the code, and run the code.
#pip3 install requirement.txt #python mnist_tf.py
C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens and Z. Wojna, “Rethinking the Inception Architecture for Computer Vision,” CoRR, p. abs/1512.00567, 2015
PyTorch, [Online]. Available: https://pytorch.org/vision/stable/index.html
PyTorch, [Online]. Available: https://pytorch.org/hub/pytorch_vision_inception_v3/
Stanford, [Online]. Available: http://cs231n.stanford.edu/
Wikipedia, [Online]. Available: https://en.wikipedia.org/wiki/Cross_entropy
AMD, “ROCm issues,” [Online]. Available: RadeonOpenCompute/ROCm#issues
PyTorch, [Online image]. https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf
TensorFlow, [Online image]. https://www.tensorflow.org/extras/tensorflow_brand_guidelines.pdf
MAGMA, [Online image]. https://bitbucket.org/icl/magma/src/master/docs/
Advanced Micro Devices, Inc., [Online]. Available: https://rocmsoftwareplatform.github.io/AMDMIGraphX/doc/html/
Advanced Micro Devices, Inc., [Online]. Available: ROCmSoftwarePlatform/AMDMIGraphX
Docker, [Online]. https://docs.docker.com/get-started/overview/
Torchvision, [Online]. Available https://pytorch.org/vision/master/index.html?highlight=torchvision#module-torchvision