Install Tensorflow 2 & PyTorch for AMD GPUs

AMD has released ROCm, a Deep Learning driver to run Tensorflow and PyTorch on AMD GPUs. Hence, I provided the installation instructions of Tensorflow and PyTorch for AMD GPUs below.

In this post, I guided you to install ROCm (AMD GPU driver) and Tensorflow ROCm (version compatible with AMD hardware) from Debian repository on Ubuntu. At the moment this post published, ROCm runs only on Linux-like platforms (e.g. Ubuntu, Centos). To run Deep Learning with AMD GPUs on MacOS, you can use PlaidML owned and maintained by PlaidML. So far, I have not seen packages to run AMD-based Deep Learning on Windows.

Updated on September. 7, 2021

Installation steps:

• Install GPU driver, ROCm
• Install AMD-compatible Tensorflow version, Tensorflow ROCm
• Install AMD-compatiblle PyTorch version

Notice: there is often a version mismatch between ROCm and TF-ROCm, and therefore, please check this link to find the compatible ROCm and TF-ROCm versions. PyTorch can be installed as Python package on AMD GPUs with ROCm 4.0 and above.

Prerequisites

• Operating Systems: Ubuntu 18.04.5 (Kernel 5.4) and other supported OSs.
• Hardware: Radeon VII (I owned it)…