Есть ли жизнь без Nvidia? Запускаем TensorFlow на Radeon RX 580

Предыстория

На новом рабочем месте меня посадили за ПК, оборудованный процессором Ryzen 2600 и видеокартой Radeon RX 580. Попробовав обучать нейронные сети на процессоре, я понял, что это не дело: уж слишком медленным был процесс. После недолгих поисков я узнал, что существует как минимум 2 способа запуска современных библиотек машинного обучения на видеокартах Radeon: PlaidML и ROCm. Я попробовал оба и хочу поделиться результатами.

PlaidML: кроссплатформенность во главе угла

TensorFlow служит бэкэндом для Keras, интерпретируя его синтаксис и преобразуя его в инструкции, которые могут выполняться на процессоре или GPU. К сожалению, он поддерживает только видеокарты с технологией Nvidia CUDA.

PlaidML - альтернативный бэкэнд для Keras с поддержкой OpenCL. Его можно использовать для обучения моделей Keras на встроенной графике процессора, дискретном или даже внешнем графическом процессоре AMD. Он работает на Windows, Linux и Mac.

UPD: В своём комментарии пользователь MikeLP указал, что:

Keras задумывался как высокоуровневый интерфейс. Соответственно остальные бекенды, как PlaidML могли обеспечить работу с другими видеоускорителями без закрытой технологии CUDA. И так было до тех пор, пока авторы не заявили, что Keras будет развиваться только в рамках Tensorflow. Все — «накрылась медным тазом ваша качалка». И получается что PlaidML может работать только с Keras, a не с tf.Keras — API разбежалось. Cоответственно смысл юзать PlaidML, когда он не совместим с последней версией Keras API в Tensorflow. Пацаны из PlaidML (они часть Intel) были вынуждены поменять roadmap и идти в другом направлении.

По моему мнению, PlaidML ещё актуален, т.к. включение Keras в Tensorflow произошло недавно и API не успел существено поменяться. Однако дальшейшие перспективы Keras+PlaidML выглядят туманно.

Установка PlaidML очень проста. Нужно поставить Python-пакет и выбрать устройство в пошаговом конфигураторе:

 pip install plaidml-keras && plaidml-setup

В коде нужно прописать использование PlaidML как бэкэнда:

os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"

После этого все вычисления будут выполняться на видеокарте. Запустим бенчмарк:

plaidbench keras mobilenet

Сравним результаты нашего видеоадаптера с некоторыми другими. RX 580 и Ryzen 2600 я тестировал локально, результаты остальных устройств взяты от других пользователей:

Видеоадаптеры AMD показывают себя ощутимо медленнее видеокарт Nvidia. Впрочем, разрыв между видеокартами на порядок меньше пропасти между видеокартами и процессорами. Исключением выступает Intel HD Graphics: PlaidML совершенно не оптимизирован для видеокарт Intel и их использование не имеет особого смысла.

ROCm: Старший брат

Со своей платформой Radeon Open Compute Platform (ROCm), AMD обеспечивает интерфейс для выполнения кода на GPU и процессорах. ROCm конкурирует с платформой NVIDIA CUDA. Он работает только на UNIX-системах.

На использовании ROCm я в итоге и остановился. По моим ощущениям, он работает в ~2 раза быстрее, чем PlaidML, и позволяет загружать в память более сложные модели, когда PlaidML выбрасывает Out Of Memory.

Я установил свежую Ubuntu 20.04 LTS. Для установки и настройки ROCm я предпринял следующие шаги: (они немного отличаются от официальной документации, т.к. она успела несколько устареть)

1. Обновил пакеты:

   sudo apt update && sudo apt dist-upgrade && sudo apt install libnuma-dev && sudo reboot

2. Откатил ядро с версии 5.8 до 5.6. Для этого я зашел на http://kernel.ubuntu.com/~kernel-ppa/mainline/ и скачал файлы:

   linux-headers-VERSION-NUMBER_all.deb
   linux-headers-VERSION-NUMBER_amd64.deb
   linux-image-VERSION-NUMBER_amd64.deb
   linux-modules-VERSION-NUMBER_amd64.deb

3. Далее скопировал их в отдельную папку, в которой выполнил:

   sudo dpkg -i *.deb

4. Перезагрузил ПК с ядром 5.6, выбрав его в GRUB. После этого удалил ядро 5.8:

   sudo apt-get purge *5.8.0*

5. Перезагрузил ПК.

6. Добавил репозиторий:

   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

7. Установил ROCm:

   sudo apt update && sudo apt install rocm-dkms && sudo reboot

8. Настроил права:

   sudo usermod -a -G video $LOGNAME
   sudo usermod -a -G render $LOGNAME

9. Для проверки успешности установки запустил эти команды. Если установка прошла успешно, обе из них выведут название GPU:

   /opt/rocm/bin/rocminfo
   /opt/rocm/opencl/bin/clinfo

10. Добавил ROCm в PATH:

    echo 'export PATH=$PATH:/opt/rocm/bin:/opt/rocm/rocprofiler/bin:/opt/rocm/opencl/bin' | sudo tee -a /etc/profile.d/rocm.sh

11. После этого установил пакеты для машинного обучения:

    sudo apt install rocm-libs miopen-hip rccl

12. Установил Pip и TensorFlow:

    sudo apt install python3-pip
    pip3 install --user tensorflow-rocm

Готово! После этого TensorFlow автоматически стал использовать мой GPU:

Давайте посмотрим на бенчмарки:

python3 ./tf_cnn_benchmarks.py --num_gpus=1 --batch_size=64 --model=resnet50

RX 580 я тестировал локально, результаты остальных устройств взяты от других пользователей:

Выводы

В 2021 году в сфере машинного обучения становится всё менее и менее важно, какая у вас видеокарта. Видеокарты AMD всё ещё ощутимо отстают в производительности от собратьев из Nvidia, однако при этом и стоят дешевле. Тем не менее, поддержка видеокарт Radeon далека от идеальной, и необходимая настройка для них сложнее и дольше.

Если у вас уже есть видеокарта от AMD, то возможно для целей разработки вам не понадобится менять её на видеокарту Nvidia. Оптимальным выбором для машинного обучения на видеокартах Radeon я считаю фреймворк ROCm. Он обеспечивает приличную производительность, и позволяет запустить TensorFlow 2.2.0 и Keras.

UPD: В комментариях выяснили, что последняя версия связки ROCm+TensorFlow не всегда корректно работает на чипе Polaris 10 (Radeon RX 580). Написал туториал, как поставить ROCm 3.5.1 + TF 2.2: https://github.com/boriswinner/RX580-rocM-tensorflow-ubuntu20.4-guide