Running Tensorflow 2.4 on AMD RX580 graphic card on Ubuntu Budgie 20.04

Deep learning on GPUs is generally linked to NVidia graphics card. Anyway it is possible to use also AMD GPUs thanks to ROCm and the Tensorflow ROCm port.

I’ve tested the solution on a RX580 8Gb (gfx903) on Ubuntu Budgie 20.04 and encountered some troubles.

Following the official guide I installed the latest rocm-dkms driver:

sudo apt update
sudo apt dist-upgrade
sudo apt install libnuma-dev
sudo reboot
wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/debian/ xenial main' | sudo tee /etc/apt/sources.list.d/rocm.list
sudo apt update
sudo apt install rocm-dkms && sudo reboot
sudo usermod -a -G video $LOGNAME
sudo usermod -a -G render $LOGNAME

After installation is finished and confirmed to be working by running the rocminfo and clinfo commands, I’ve followed the instructions on the Tensorflow ROCm port github page, but using the latest docker image available, at the time of writing at least, won’t work. So I tested various images and found that the one working is the rocm/tensorflow:rocm4.0-tf2.4-dev.

I modified the command from the instruction as follow:

alias drun='sudo 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 \
      -v $HOME/dockerx:/dockerx'

drun rocm/tensorflow:rocm4.0-tf2.4-dev

To test it out run:

$ python3
>>> import tensorflow as tf
>>> tf.add(1, 2).numpy()
3
>>> hello = tf.constant('Hello, TensorFlow!')
>>> hello.numpy()
b'Hello, TensorFlow!'

Don’t know if this will work for other AMD graphic cards, but those of the same family of the RX580 should be working.

I’ve tried to use other suggested workarounds like use ROCm 3.5.1 drivers as suggested from github user Grench6, but with no luck.

The AMD support seems to be still not very mature, or at least be somewhat unstable especially for some card families, but things are getting better.

It should be also possible to build and run Pytorch and Caffe2 with GPU support, but this is material for another story.

Hope this will be helpful for someone else.

Stay safe!